diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 03bdb027296..bb114c6e22b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -44,6 +44,7 @@ Changelog - Add Torch-TensorRT FP8 deployment example for HuggingFace ViT (``examples/torch_trt/``): ``torch_tensorrt_ptq.py`` covers ``mtq.quantize`` → ``torch_tensorrt.compile(ir="dynamo")``, and ``torch_tensorrt_accuracy.py`` reports the compiled model's ImageNet-1k top-1/top-5 accuracy via the ``onnx_ptq`` ``evaluate`` harness (the unquantized baseline is Torch-TensorRT-compiled too, for an apples-to-apples comparison). Ships a ViT-tuned FP8 PTQ recipe under ``modelopt_recipes/huggingface/vit/ptq/`` (``fp8.yaml``) composed from the shared ``modelopt_recipes/configs/`` units: it quantizes the encoder Linears, patch-embed ``nn.Conv2d``, ``classifier``, and per-block LayerNorm inputs plus the attention Q/K/V BMMs and softmax. Verified on ``google/vit-base-patch16-224`` (ImageNet-1k 50k validation): FP8 stays within 0.13 pp Top-1 of the FP16 baseline. - Add **AutoQuantize recipe** support: ``mtq.auto_quantize`` can be driven declaratively from a YAML recipe (``RecipeType.AUTO_QUANTIZE`` / ``AutoQuantizeConfig``) specifying candidate formats, the ``effective_bits`` target, cost model (incl. ``active_moe`` and ``excluded_module_name_patterns``), scoring method, and disabled layers. Adds an ``effective_bits`` cost-model override on ``QuantizeConfig`` / ``QuantizerAttributeConfig`` (block-scale-accurate NVFP4 = 4.5 via ``configs/numerics/nvfp4``). Shipped recipes live under ``modelopt_recipes/general/auto_quantize/`` and model-specific ones under ``modelopt_recipes/huggingface//auto_quantize/``. - Add ``rotate.mode`` to torch quantizer configs. The default ``"rotate"`` keeps the existing rotate-before-quantize behavior; ``"rotate_back"`` enables fake-quant rotate → quantize → rotate-back for TensorQuantizer. +- Add a ``constant_amax`` ``QuantizerAttributeConfig`` field that pins a quantizer's ``amax`` to a fixed value and skips activation calibration (no forward statistics collected). Unlike ``use_constant_amax`` (which hardcodes 448.0 for KV-cache cast math and registers no buffer), ``constant_amax`` stores the constant on the ``_amax`` buffer so it is used by both the fake-quant forward and the exported scaling factor. For NVFP4 activations the exported ``input_scale`` equals ``amax / (E2M1_MAX * E4M3_MAX)``, so ``constant_amax: 2688.0`` yields ``input_scale == 1.0``. Ships a new recipe ``modelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml`` that applies this to the MoE expert activation quantizers. **Bug Fixes** diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 0ca30d18448..fd40df74d25 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -671,6 +671,43 @@ def validate_calibrator(cls, v, info: ValidationInfo): """, ) + constant_amax: float | None = ModeloptField( + default=None, + title="Pin the quantizer amax to a constant value and skip calibration.", + description="""If set, the quantizer ``amax`` is fixed to this constant value and no + activation calibration is performed (no forward statistics are collected). The constant + is stored on the ``_amax`` buffer, so it is used by both the fake-quant forward pass and + the exported scaling factor (e.g. ``input_scale``). + + This differs from ``use_constant_amax``, which pins amax to the FP8 E4M3 range (448.0) + for KV-cache cast math and does not register an ``_amax`` buffer. For NVFP4 activations + the exported ``input_scale`` equals ``amax / (E2M1_MAX * E4M3_MAX) = amax / (6 * 448)``; + setting ``constant_amax`` to ``2688.0`` therefore yields an exported ``input_scale`` of + ``1.0``. + """, + ) + + @field_validator("constant_amax") + @classmethod + def validate_constant_amax(cls, v): + """Validate that constant_amax, when set, is a positive value.""" + assert v is None or v > 0, "constant_amax must be a positive value." + return v + + @model_validator(mode="after") + def validate_constant_amax_modes(self): + """Forbid combining ``use_constant_amax`` and ``constant_amax``. + + Both pin the amax but disagree on the value: ``use_constant_amax`` uses the FP8 E4M3 + range (448.0) for the fake-quant forward and registers no ``_amax`` buffer, while + ``constant_amax`` pins ``_amax`` to its configured value for both forward and export. + Setting both would make the simulated and exported scales silently diverge. + """ + assert not (self.use_constant_amax and self.constant_amax is not None), ( + "use_constant_amax and constant_amax are mutually exclusive; set only one." + ) + return self + class LayerwiseConfig(ModeloptBaseConfig): """Nested config for layer-by-layer calibration behavior.""" diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 7e5bb85c09b..ed7c44c4d09 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -920,8 +920,8 @@ def enable_stats_collection(model: nn.Module): """Enable stats collection for all quantizers in the model.""" for name, module in model.named_modules(): if isinstance(module, TensorQuantizer) and not module._disabled: - if module._use_constant_amax: - # use_constant_amax quantizers use a fixed amax and don't need calibration. + if module._use_constant_amax or module._constant_amax is not None: + # Quantizers with a constant amax use a fixed amax and don't need calibration. # Disable quantization during calibration so it doesn't affect other quantizers. module.disable_quant() continue @@ -938,8 +938,8 @@ def finish_stats_collection(model: nn.Module, method: str | None = None, **kwarg if not isinstance(module, TensorQuantizer) or module._disabled: continue - if module._use_constant_amax: - # Re-enable quantization for use_constant_amax quantizers disabled in enable_stats_collection. + if module._use_constant_amax or module._constant_amax is not None: + # Re-enable quantization for constant-amax quantizers disabled in enable_stats_collection. module.enable_quant() continue diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 98d9e0dcb1e..80954834f0a 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -202,6 +202,7 @@ def __init__( self.amax = amax self._use_constant_amax = False + self._constant_amax = None self.set_from_attribute_config(quant_attribute_cfg) self._if_quant = if_quant @@ -244,6 +245,13 @@ def _block_sizes_setter(val): self._calibrator._axis = None return val + def _constant_amax_setter(val): + if val is not None: + # Pin amax to the constant on the _amax buffer so it is used by both the + # fake-quant forward pass and export; calibration is skipped in model_calib. + self.amax = float(val) + return val + # Some attributes need custom handling. # By default, attributes from config are mapped to a name ``f"_{attribute}"`` _custom_setters: dict[str, tuple[str, Callable]] = { @@ -255,6 +263,7 @@ def _block_sizes_setter(val): "backend": ("backend", lambda val: val), "backend_extra_args": ("backend_extra_args", lambda val: val or {}), "use_constant_amax": ("_use_constant_amax", lambda val: val), + "constant_amax": ("_constant_amax", _constant_amax_setter), } for attribute, val in attribute_cfg.items(): @@ -722,6 +731,10 @@ def _get_amax(self, inputs): return torch.tensor(torch.finfo(torch.float8_e4m3fn).max, device=inputs.device) if hasattr(self, "_amax"): amax = self._amax + # A constant_amax buffer is registered at config time (on CPU) and may not have + # followed a later `model.to(device)`; align it with the input device on the fly. + if amax.device != inputs.device: + amax = amax.to(inputs.device) else: reduce_axis = quant_utils.convert_quantization_axis_to_reduce_axis(inputs, self._axis) amax = quant_utils.reduce_amax(inputs, axis=reduce_axis, keepdims=True).detach() diff --git a/modelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml b/modelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml new file mode 100644 index 00000000000..712a2d5d220 --- /dev/null +++ b/modelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml @@ -0,0 +1,63 @@ +# 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. + +# Composed PTQ recipe for expert-only NVFP4 quantization with a fixed activation input_scale of +# 1.0 (no activation calibration), plus FP8 KV-cache cast mode. +# +# The expert weight quantizers use standard NVFP4 (per-tensor weight scale + dynamic block +# scales, calibrated via max). The expert input (activation) quantizers pin the per-tensor amax +# to a constant 2688.0 = E2M1_MAX * E4M3_MAX (6 * 448) so the exported ``input_scale`` is exactly +# ``amax / (E2M1_MAX * E4M3_MAX) = 1.0``; the per-block E4M3 activation scales remain dynamic. +# No activation statistics are collected for these quantizers. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + kv_fp8_cast: configs/ptq/units/kv_fp8_cast + +metadata: + recipe_type: ptq + description: >- + Applies NVFP4 only to expert-layer weight and input quantizers, pinning the expert activation + input_scale to 1.0 (constant amax = 2688, no activation calibration), plus FP8 KV-cache cast + mode using constant amax; uses max calibration for the remaining (weight) quantizers. +quantize: + algorithm: + method: max + # Max calibration is fast and does not typically need checkpointing. + # layerwise=false required for VLMs where the decoder layers are nested under + # `model.language_model.layers` (layerwise_calibrate can't find them otherwise). + layerwise: false + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.experts.*input_quantizer' + cfg: + $import: nvfp4 + # 2688 = E2M1_MAX * E4M3_MAX (6 * 448) -> exported input_scale == 1.0, no calibration. + constant_amax: 2688.0 + - quantizer_name: '*block_sparse_moe*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*block_sparse_moe*input_quantizer' + cfg: + $import: nvfp4 + # 2688 = E2M1_MAX * E4M3_MAX (6 * 448) -> exported input_scale == 1.0, no calibration. + constant_amax: 2688.0 + - $import: kv_fp8_cast + - $import: default_disabled_quantizers diff --git a/tests/_test_utils/torch/quantization/tensor_quantizer_common.py b/tests/_test_utils/torch/quantization/tensor_quantizer_common.py index c9de64e6c60..dd8d790ee3f 100644 --- a/tests/_test_utils/torch/quantization/tensor_quantizer_common.py +++ b/tests/_test_utils/torch/quantization/tensor_quantizer_common.py @@ -28,6 +28,7 @@ max_calibrate, ) from modelopt.torch.quantization.nn import QuantLinear, SequentialQuantizer, TensorQuantizer +from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor class TensorQuantizerTester: @@ -267,6 +268,81 @@ def test_use_constant_amax_skips_calibration(self): assert not model["tq_const"]._disabled assert model["tq_const"]._if_quant + def test_constant_amax_registers_amax(self): + """constant_amax pins the _amax buffer to the configured value (used by fwd and export).""" + tq = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + constant_amax=2688.0, + ) + ).to(self.device) + + assert tq._constant_amax == 2688.0 + assert hasattr(tq, "_amax") + assert float(tq._amax) == 2688.0 + + # Default (unset) constant_amax must not register an _amax buffer. + tq_default = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + ) + ).to(self.device) + assert tq_default._constant_amax is None + assert not hasattr(tq_default, "_amax") + + def test_constant_amax_nvfp4_input_scale(self): + """For NVFP4, constant_amax=2688 (= E2M1_MAX * E4M3_MAX) exports input_scale == 1.0.""" + tq = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + constant_amax=2688.0, + ) + ).to(self.device) + + input_scale = NVFP4QTensor.get_activation_scaling_factor(tq) + assert torch.allclose(input_scale.float(), torch.ones_like(input_scale.float())) + + def test_constant_amax_skips_calibration(self): + """constant_amax quantizers are excluded from calibration and keep their fixed amax.""" + model = nn.ModuleDict( + { + "tq_const": TensorQuantizer( + QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + constant_amax=2688.0, + ) + ), + "tq_calib": TensorQuantizer(QuantizerAttributeConfig(num_bits=8)), + } + ).to(self.device) + + enable_stats_collection(model) + # constant_amax quantizer: quant disabled during calibration, not collecting stats. + assert not model["tq_const"]._if_calib + assert not model["tq_const"]._if_quant + assert float(model["tq_const"]._amax) == 2688.0 + + finish_stats_collection(model) + # Re-enabled and amax is NOT overwritten by calibration. + assert model["tq_const"]._if_quant + assert float(model["tq_const"]._amax) == 2688.0 + + def test_constant_amax_must_be_positive(self): + """constant_amax must be a positive value.""" + with pytest.raises(Exception, match="constant_amax must be a positive value"): + QuantizerAttributeConfig(num_bits=(2, 1), constant_amax=-1.0) + with pytest.raises(Exception, match="constant_amax must be a positive value"): + QuantizerAttributeConfig(num_bits=(2, 1), constant_amax=0.0) + + def test_constant_amax_mutually_exclusive_with_use_constant_amax(self): + """use_constant_amax and constant_amax cannot both be set.""" + with pytest.raises(Exception, match="mutually exclusive"): + QuantizerAttributeConfig(num_bits=8, use_constant_amax=True, constant_amax=2688.0) + def test_modelopt_state(self): # Test loading of amax from ref to test tensor_quantizer_ref = TensorQuantizer(QuantizerAttributeConfig(num_bits=4), amax=10.0)