diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index c86c90eaa59..72e33d7c254 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1330,13 +1330,16 @@ class _QuantFP8Linear(QuantModule): def _setup(self): self.input_quantizer = TensorQuantizer() self.weight_quantizer = TensorQuantizer() - assert self.weight_scale_inv.ndim == 2, "Weight scale inverse must be 2D" assert self.weight.ndim == 2, "Weight must be 2D" - self.block_size = max( - self.weight.shape[0] // self.weight_scale_inv.shape[0], - self.weight.shape[1] // self.weight_scale_inv.shape[1], - ) - assert self.block_size == 128, "Block size must be 128" + if self.weight_scale_inv.ndim == 0: + self.block_size = None + else: + assert self.weight_scale_inv.ndim == 2, "Weight scale inverse must be 0D or 2D" + self.block_size = max( + self.weight.shape[0] // self.weight_scale_inv.shape[0], + self.weight.shape[1] // self.weight_scale_inv.shape[1], + ) + assert self.block_size == 128, "Block size must be 128" def _get_weight_and_scale_inv(self): if isinstance(self.weight, torch.distributed.tensor.DTensor): @@ -1347,12 +1350,17 @@ def _get_weight_and_scale_inv(self): scale_inv = self.weight_scale_inv.contiguous() return weight, scale_inv - def forward(self, input: Tensor) -> Tensor: + def _dequantize_weight(self, dtype: torch.dtype) -> Tensor: + weight, scale_inv = self._get_weight_and_scale_inv() + if self.block_size is None: + return weight.to(dtype) * scale_inv.to(dtype) assert weight_dequant is not None, "Triton is not available" + return weight_dequant(weight, scale_inv, self.block_size, dtype=dtype) + + def forward(self, input: Tensor) -> Tensor: if self.weight.element_size() == 1: with torch.cuda.device(self.weight.device): - weight, scale_inv = self._get_weight_and_scale_inv() - weight = weight_dequant(weight, scale_inv, self.block_size, dtype=input.dtype) + weight = self._dequantize_weight(input.dtype) else: weight = self.weight return linear( @@ -1362,11 +1370,9 @@ def forward(self, input: Tensor) -> Tensor: ) def unpack_weight(self): - assert weight_dequant is not None, "Triton is not available" with torch.cuda.device(self.weight.device): - weight, scale_inv = self._get_weight_and_scale_inv() self.weight = nn.Parameter( - weight_dequant(weight, scale_inv, self.block_size, dtype=torch.get_default_dtype()), + self._dequantize_weight(torch.get_default_dtype()), requires_grad=False, ) if hasattr(self, "weight_scale_inv"): diff --git a/modelopt_recipes/huggingface/models/nvidia/Mistral-Medium-3.5-128B-NVFP4/ptq/nvfp4-max-calib.yaml b/modelopt_recipes/huggingface/models/nvidia/Mistral-Medium-3.5-128B-NVFP4/ptq/nvfp4-max-calib.yaml new file mode 100644 index 00000000000..43798137f0f --- /dev/null +++ b/modelopt_recipes/huggingface/models/nvidia/Mistral-Medium-3.5-128B-NVFP4/ptq/nvfp4-max-calib.yaml @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + fp8: configs/numerics/fp8 + kv_fp8: configs/ptq/units/kv_fp8 + +metadata: + recipe_type: ptq + description: >- + NVFP4 W4A4 on interior MLP layers, FP8 W8A8 on edge MLP and attention + layers, and an FP8 KV cache; uses max calibration. +quantize: + algorithm: + method: max + layerwise: false + quant_cfg: + - $import: base_disable_all + # NVFP4 on MLP and expert weight matrices. + - quantizer_name: '*mlp.experts*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mlp.experts*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*block_sparse_moe*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*block_sparse_moe*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mlp*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mlp*input_quantizer' + cfg: + $import: nvfp4 + # FP8 on language-model attention projections. + - quantizer_name: '*self_attn.q_proj*weight_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*self_attn.q_proj*input_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*self_attn.k_proj*weight_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*self_attn.k_proj*input_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*self_attn.v_proj*weight_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*self_attn.v_proj*input_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*self_attn.o_proj*weight_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*self_attn.o_proj*input_quantizer' + cfg: + $import: fp8 + # Keep the first four and final decoder MLP layers in FP8. + - quantizer_name: '*layers.0.mlp*weight_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*layers.0.mlp*input_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*layers.1.mlp*weight_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*layers.1.mlp*input_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*layers.2.mlp*weight_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*layers.2.mlp*input_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*layers.3.mlp*weight_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*layers.3.mlp*input_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*layers.87.mlp*weight_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*layers.87.mlp*input_quantizer' + cfg: + $import: fp8 + - $import: kv_fp8 + - $import: default_disabled_quantizers diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index 3aaacaa3e0e..f2bf8dce034 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -166,6 +166,7 @@ def test_load_recipe_builtin_description(): "general/ptq/nvfp4_experts_only-kv_fp8", "general/ptq/nvfp4_experts_only-kv_fp8_cast", "general/ptq/nvfp4_experts_only-kv_fp8_layerwise", + "huggingface/models/nvidia/Mistral-Medium-3.5-128B-NVFP4/ptq/nvfp4-max-calib", "general/ptq/nvfp4_mlp_only-kv_fp8", "general/ptq/nvfp4_mlp_only-novit-kv_fp8", "general/ptq/nvfp4_mlp_only-kv_fp8_cast", diff --git a/tests/unit/torch/quantization/plugins/test_huggingface.py b/tests/unit/torch/quantization/plugins/test_huggingface.py index 89d723d1248..b16ddde705c 100644 --- a/tests/unit/torch/quantization/plugins/test_huggingface.py +++ b/tests/unit/torch/quantization/plugins/test_huggingface.py @@ -45,6 +45,7 @@ import transformers from transformers import AutoModelForCausalLM, LlamaForCausalLM +from transformers.integrations.finegrained_fp8 import FP8Linear from transformers.models.dbrx.configuration_dbrx import DbrxConfig, DbrxFFNConfig from transformers.models.dbrx.modeling_dbrx import DbrxExpertGLU, DbrxExperts, DbrxFFN @@ -109,6 +110,21 @@ def test_convert_conv1d(): assert torch.allclose(out_1, out_2) +def test_fp8_linear_per_tensor_dequant(monkeypatch): + module = FP8Linear(2, 2, block_size=(128, 128)) + module.weight_scale_inv = nn.Parameter(torch.tensor(2.0)) + with torch.no_grad(): + module.weight.copy_(torch.tensor([[-2.0, 1.0], [0.5, 4.0]], dtype=torch.float8_e4m3fn)) + + mtq.replace_quant_module(module) + monkeypatch.setattr("modelopt.torch.quantization.plugins.huggingface.weight_dequant", None) + + assert module.block_size is None + torch.testing.assert_close( + module._dequantize_weight(torch.float32), module.weight.float() * 2.0 + ) + + @pytest.mark.skipif( Version(transformers.__version__) < Version("5.0"), reason="test_dbrx is not supported for transformers<5.0",