From bd7ce1f19f669ed498a6323a659100211638e359 Mon Sep 17 00:00:00 2001 From: Zhiyu Cheng Date: Fri, 10 Jul 2026 15:28:05 -0700 Subject: [PATCH 1/3] feat(quant): skip max-calib forward when no activation needs data When all enabled activation quantizers use a constant amax (constant_amax / use_constant_amax) or dynamic quantization, max calibration's forward pass collects no data-driven statistics and is pure overhead. Add _needs_activation_forward_for_max_calib() and gate the forward_loop in max_calibrate() behind a new opt-in flag skip_forward_without_activation_calib. Weight quantizers are still calibrated on the weight tensors directly via weight_only_quantize(), so results are unchanged; only the wasted forward is avoided (e.g. the nvfp4_experts_only_input_scale1 recipe). The flag defaults to False in max_calibrate() so the advanced algorithms that call it directly and always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) are unaffected; it is enabled via MaxCalibConfig.skip_forward_without_activation_calib (default True) so only the top-level max path opts in. Also simplify a default-arg lambda in core_utils.sync_moe_expert_amax so mypy can infer its type (surfaced by the max_calibrate signature change); the forward_loop is invoked synchronously so closure capture is safe. Stacked on #1947 (depends on the constant_amax attribute). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Zhiyu Cheng --- CHANGELOG.rst | 1 + modelopt/torch/quantization/config.py | 13 ++++ modelopt/torch/quantization/model_calib.py | 48 +++++++++++- .../torch/quantization/utils/core_utils.py | 5 +- tests/unit/torch/quantization/test_calib.py | 78 ++++++++++++++++++- 5 files changed, 141 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index df3cdefdb5e..d77d8ba275a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -47,6 +47,7 @@ Changelog - 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. +- Max calibration now skips the calibration ``forward_loop`` when no enabled quantizer needs data-driven activation statistics — e.g. an experts-only recipe whose activation quantizers all use ``constant_amax`` / ``use_constant_amax``, or dynamic quantization. Weight calibration still runs on the weight tensors directly, so results are unchanged; only the wasted forward is avoided. Controlled by ``MaxCalibConfig.skip_forward_without_activation_calib`` (default ``True``); the advanced algorithms that always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) are unaffected. **Bug Fixes** diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index a8ed323a7b2..532a9cdf33b 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -939,6 +939,19 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): ), ) + skip_forward_without_activation_calib: bool = ModeloptField( + default=True, + title="Skip the calibration forward when no activation quantizer needs data.", + description=( + "If True (default), max calibration skips the ``forward_loop`` entirely when no " + "enabled quantizer collects data-driven activation statistics — e.g. an " + "experts-only recipe whose activation quantizers all use ``constant_amax`` or " + "``use_constant_amax``, or dynamic quantization. Weight calibration still runs on " + "the weight tensors directly, so results are unchanged; only the wasted forward is " + "avoided. Set to False to always run the provided ``forward_loop``." + ), + ) + class MseCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): """Configuration for per-tensor MSE calibration. diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index e03dff2e98b..d0eedfb6495 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -254,6 +254,38 @@ def _should_sync_amax_across_ep( return True +def _needs_activation_forward_for_max_calib(model: nn.Module) -> bool: + """Return True if any enabled quantizer still needs a calibration forward pass. + + Weight quantizers are calibrated directly on the weight tensors by + :func:`weight_only_quantize`, so they never need the data forward. Activation-side + quantizers (input/output/BMM) need it only when they collect data-driven statistics: + i.e. they are enabled, have a calibrator, and are neither dynamic nor pinned to a constant + amax (``use_constant_amax`` / ``constant_amax``). Constant-amax quantizers also skip bias + calibration in :func:`finish_stats_collection`, so they never require the forward either. + + When this returns False (e.g. an experts-only recipe whose activation quantizers all use + ``constant_amax``), the calibration forward can be skipped entirely and only weight + calibration is performed. + """ + for name, module in model.named_modules(): + if not isinstance(module, TensorQuantizer) or module._disabled: + continue + # Weight quantizers (incl. SequentialQuantizer stages named ``weight_quantizer.``) + # are calibrated on the weight tensor directly, not via the data forward. + if any(part.endswith("weight_quantizer") for part in name.split(".")): + continue + if ( + module._dynamic + or module._use_constant_amax + or getattr(module, "_constant_amax", None) is not None + ): + continue + if getattr(module, "_calibrator", None) is not None: + return True + return False + + @torch.no_grad() def max_calibrate( model: nn.Module, @@ -261,6 +293,7 @@ def max_calibrate( distributed_sync=True, sync_expert_weight_amax=False, shared_states: Mapping[str, Mapping[str, Sequence[str]]] | None = None, + skip_forward_without_activation_calib: bool = False, ): """Calibrate the model using max. @@ -274,6 +307,11 @@ def max_calibrate( shared_states: Optional dict keyed by shared-state name. ``"weight_global_amax"`` is implemented today and accepts ``{"patterns": [...]}``; omitted patterns use ``SHARED_PATTERNS``, while an empty list disables the state. + skip_forward_without_activation_calib: If True, skip the (potentially expensive) + ``forward_loop`` when no enabled quantizer needs data-driven activation statistics + (see :func:`_needs_activation_forward_for_max_calib`). Weight calibration still runs. + Only opt-in for the top-level ``max`` path; algorithms that always need activations + (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) call this with the default False. See :class:`MaxCalibConfig ` for details on the remaining arguments. @@ -292,7 +330,15 @@ def max_calibrate( enable_stats_collection(model) weight_only_quantize(model) if forward_loop is not None: - forward_loop(model) + if skip_forward_without_activation_calib and not _needs_activation_forward_for_max_calib( + model + ): + print_rank_0( + "max_calibrate: all enabled activation quantizers use constant/dynamic amax; " + "skipping the calibration forward pass (weight-only calibration)." + ) + else: + forward_loop(model) finish_stats_collection(model) # Sync quantizer amax across local experts within each rank (for SequentialMLP) diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 5c839f050ad..80352ced9eb 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -647,7 +647,10 @@ def sync_moe_expert_amax(experts, sync_weight_amax=False): if name.endswith("weight_quantizer") and module.is_enabled and module.amax is None: weight = expert.state_dict().get(name.replace("weight_quantizer", "weight")) if weight is not None: - max_calibrate(module, lambda m, w=weight: m(w), distributed_sync=False) + # max_calibrate invokes the forward_loop synchronously, so capturing + # ``weight`` by closure (rather than a default arg) is safe and lets mypy + # infer the lambda type. + max_calibrate(module, lambda m: m(weight), distributed_sync=False) @contextmanager diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index 64d89141bcd..a44d4636b19 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -23,13 +23,15 @@ from _test_utils.torch.quantization.quantize_common import get_awq_config import modelopt.torch.quantization as mtq -from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.config import MaxCalibConfig, QuantizerAttributeConfig from modelopt.torch.quantization.model_calib import ( + _needs_activation_forward_for_max_calib, apply_pre_quant_scale_and_smooth, disable_pre_quant_scale_and_resmooth, layerwise_calibrate, + max_calibrate, ) -from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.nn import QuantLinear, TensorQuantizer from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -583,3 +585,75 @@ def _pre_hook(_module, args): assert torch.allclose(observed_layer_inputs[1][0], torch.tensor([[2.0, 4.0]])) # Layer 2 gets output of layer 1 (prev * mask1 * scale1 = [2,4] * 0.5 * 0.5 = [0.5,1.0]) assert torch.allclose(observed_layer_inputs[2][0], torch.tensor([[0.5, 1.0]])) + + +def _make_quant_linear(input_cfg, fi=16, fo=8): + """QuantLinear with an int8 weight quantizer and a caller-specified input quantizer.""" + lin = QuantLinear(fi, fo, bias=False) + lin.weight_quantizer.set_from_attribute_config(QuantizerAttributeConfig(num_bits=8)) + lin.input_quantizer.set_from_attribute_config(input_cfg) + return lin + + +def test_max_calib_skips_forward_when_all_activations_constant(): + """constant_amax activation quantizer => no data forward; weights still calibrated.""" + lin = _make_quant_linear(QuantizerAttributeConfig(num_bits=8, constant_amax=127.0)) + assert not _needs_activation_forward_for_max_calib(lin) + + calls = [] + + def forward_loop(model): + calls.append(1) + model(torch.randn(4, 16)) + + max_calibrate( + lin, forward_loop, distributed_sync=False, skip_forward_without_activation_calib=True + ) + assert calls == [] # forward skipped entirely + assert hasattr( + lin.weight_quantizer, "_amax" + ) # weight still calibrated via weight_only_quantize + assert float(lin.input_quantizer._amax) == 127.0 # constant amax preserved + + +def test_max_calib_runs_forward_when_activation_needs_data(): + """A normal (non-constant) input quantizer still triggers the calibration forward.""" + lin = _make_quant_linear(QuantizerAttributeConfig(num_bits=8)) + assert _needs_activation_forward_for_max_calib(lin) + + calls = [] + + def forward_loop(model): + calls.append(1) + model(torch.randn(4, 16)) + + max_calibrate( + lin, forward_loop, distributed_sync=False, skip_forward_without_activation_calib=True + ) + assert calls == [1] + assert lin.input_quantizer.amax is not None + + +def test_max_calib_default_does_not_skip_forward(): + """Default flag is opt-out=False for direct callers: forward always runs (backward compat).""" + lin = _make_quant_linear(QuantizerAttributeConfig(num_bits=8, constant_amax=127.0)) + + calls = [] + + def forward_loop(model): + calls.append(1) + model(torch.randn(4, 16)) + + max_calibrate(lin, forward_loop, distributed_sync=False) # skip flag defaults to False + assert calls == [1] + + +def test_needs_activation_forward_ignores_disabled_and_weight_quantizers(): + """Disabled activation quantizers (weight-only recipe) need no forward.""" + lin = _make_quant_linear(QuantizerAttributeConfig(num_bits=8, enable=False)) + assert not _needs_activation_forward_for_max_calib(lin) + + +def test_max_calib_config_enables_skip_by_default(): + """The top-level max path opts in via MaxCalibConfig.""" + assert MaxCalibConfig().skip_forward_without_activation_calib is True From 3f9d4e28d69733109fc2527cc6fa8bee1b63df22 Mon Sep 17 00:00:00 2001 From: Zhiyu Cheng Date: Mon, 13 Jul 2026 14:55:44 -0700 Subject: [PATCH 2/3] fix(quant): make skip-forward opt-in + cover MX formats Address PR #1963 review feedback: - Make MaxCalibConfig.skip_forward_without_activation_calib opt-in (default False). Defaulting it True silently changed calibration for every max-PTQ user and broke weight-only quantization under DeepSpeed ZeRO-3 (test_nested_deepspeed_backward[3-quant_cfg0]): the calibration forward materializes sharded params, so skipping it left the ZeRO-3 param shards out of sync with the quantized weights. Enable it explicitly in the nvfp4_experts_only_input_scale1-kv_fp8_cast recipe. - Also skip the forward for MX formats (MXFP4/MXFP8) in _needs_activation_forward_for_max_calib (realAsma). MX configs set type: dynamic inside block_sizes (E8M0 per-block scales) but leave the top-level _dynamic False, so the _dynamic check did not catch them; they carry no per-tensor amax (is_mx_format -> amax is None), so add an explicit is_mx_format check. Update tests (opt-in default, MX/dynamic coverage) and changelog. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Zhiyu Cheng --- CHANGELOG.rst | 2 +- modelopt/torch/quantization/config.py | 18 ++++++++----- modelopt/torch/quantization/model_calib.py | 12 ++++++--- ...experts_only_input_scale1-kv_fp8_cast.yaml | 4 +++ tests/unit/torch/quantization/test_calib.py | 27 ++++++++++++++++--- 5 files changed, 49 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d77d8ba275a..4e0d092bdbe 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -47,7 +47,7 @@ Changelog - 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. -- Max calibration now skips the calibration ``forward_loop`` when no enabled quantizer needs data-driven activation statistics — e.g. an experts-only recipe whose activation quantizers all use ``constant_amax`` / ``use_constant_amax``, or dynamic quantization. Weight calibration still runs on the weight tensors directly, so results are unchanged; only the wasted forward is avoided. Controlled by ``MaxCalibConfig.skip_forward_without_activation_calib`` (default ``True``); the advanced algorithms that always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) are unaffected. +- Add ``MaxCalibConfig.skip_forward_without_activation_calib`` (opt-in, default ``False``): when enabled, max calibration skips the ``forward_loop`` if no enabled quantizer needs data-driven activation statistics — e.g. an experts-only recipe whose activation quantizers all use ``constant_amax`` / ``use_constant_amax``, or dynamic / MX (MXFP4/MXFP8) quantization. Weight calibration still runs on the weight tensors directly, so the quantized weights are unchanged; only the wasted forward is avoided. It is opt-in because the ``forward_loop`` can carry caller-side effects (notably materializing sharded parameters under DeepSpeed ZeRO-3). The advanced algorithms that always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) call ``max_calibrate`` directly and are unaffected. The ``nvfp4_experts_only_input_scale1-kv_fp8_cast`` recipe enables it. **Bug Fixes** diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 532a9cdf33b..6c07afab90c 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -940,15 +940,19 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): ) skip_forward_without_activation_calib: bool = ModeloptField( - default=True, + default=False, title="Skip the calibration forward when no activation quantizer needs data.", description=( - "If True (default), max calibration skips the ``forward_loop`` entirely when no " - "enabled quantizer collects data-driven activation statistics — e.g. an " - "experts-only recipe whose activation quantizers all use ``constant_amax`` or " - "``use_constant_amax``, or dynamic quantization. Weight calibration still runs on " - "the weight tensors directly, so results are unchanged; only the wasted forward is " - "avoided. Set to False to always run the provided ``forward_loop``." + "If True, max calibration skips the ``forward_loop`` entirely when no enabled " + "quantizer collects data-driven activation statistics — e.g. an experts-only recipe " + "whose activation quantizers all use ``constant_amax`` / ``use_constant_amax``, " + "dynamic, or MX (MXFP4/MXFP8) quantization. Weight calibration still runs on the " + "weight tensors directly, so the quantized weights are unchanged; only the wasted " + "forward is avoided. " + "Opt-in (default False) because the provided ``forward_loop`` can carry side " + "effects the caller relies on — most notably materializing sharded parameters under " + "DeepSpeed ZeRO-3 — so enable it per-recipe when the calibration data is known to be " + "unnecessary." ), ) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index d0eedfb6495..f06603dff16 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -260,9 +260,14 @@ def _needs_activation_forward_for_max_calib(model: nn.Module) -> bool: Weight quantizers are calibrated directly on the weight tensors by :func:`weight_only_quantize`, so they never need the data forward. Activation-side quantizers (input/output/BMM) need it only when they collect data-driven statistics: - i.e. they are enabled, have a calibrator, and are neither dynamic nor pinned to a constant - amax (``use_constant_amax`` / ``constant_amax``). Constant-amax quantizers also skip bias - calibration in :func:`finish_stats_collection`, so they never require the forward either. + i.e. they are enabled, have a calibrator, and are none of + + - dynamic (top-level ``type: dynamic``), which computes amax on the fly; + - MX formats (MXFP4/MXFP8), whose per-block E8M0 scales are dynamic and which carry no + per-tensor amax (``is_mx_format`` makes ``amax`` always ``None``); + - pinned to a constant amax (``use_constant_amax`` / ``constant_amax``). Constant-amax + quantizers also skip bias calibration in :func:`finish_stats_collection`, so they never + require the forward either. When this returns False (e.g. an experts-only recipe whose activation quantizers all use ``constant_amax``), the calibration forward can be skipped entirely and only weight @@ -277,6 +282,7 @@ def _needs_activation_forward_for_max_calib(model: nn.Module) -> bool: continue if ( module._dynamic + or module.is_mx_format or module._use_constant_amax or getattr(module, "_constant_amax", None) is not None ): 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 index 712a2d5d220..9d237eef31c 100644 --- 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 @@ -41,6 +41,10 @@ quantize: # 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 + # Every quantized activation quantizer here is either constant_amax (experts/block-sparse + # inputs) or use_constant_amax (KV cast), so no data-driven activation stats are collected. + # Skip the calibration forward and do weight-only calibration instead. + skip_forward_without_activation_calib: true quant_cfg: - $import: base_disable_all - quantizer_name: '*.experts.*weight_quantizer' diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index a44d4636b19..27a5b3905a8 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -654,6 +654,27 @@ def test_needs_activation_forward_ignores_disabled_and_weight_quantizers(): assert not _needs_activation_forward_for_max_calib(lin) -def test_max_calib_config_enables_skip_by_default(): - """The top-level max path opts in via MaxCalibConfig.""" - assert MaxCalibConfig().skip_forward_without_activation_calib is True +def test_needs_activation_forward_ignores_mx_and_dynamic_quantizers(): + """MX (block-dynamic E8M0, no per-tensor amax) and dynamic activations need no forward.""" + # MXFP8: block-dynamic with E8M0 scales; top-level ``_dynamic`` is False, so it must be + # excluded via ``is_mx_format`` rather than the ``_dynamic`` check. + mx = _make_quant_linear( + QuantizerAttributeConfig( + num_bits=(4, 3), block_sizes={-1: 32, "type": "dynamic", "scale_bits": (8, 0)} + ) + ) + assert mx.input_quantizer.is_mx_format + assert not mx.input_quantizer._dynamic + assert not _needs_activation_forward_for_max_calib(mx) + + # Top-level dynamic activation quantization also needs no calibration forward. + dyn = _make_quant_linear(QuantizerAttributeConfig(num_bits=8, type="dynamic")) + assert not _needs_activation_forward_for_max_calib(dyn) + + +def test_max_calib_config_skip_is_opt_in(): + """The flag is opt-in (default False) so it does not change behavior for direct callers.""" + assert MaxCalibConfig().skip_forward_without_activation_calib is False + assert MaxCalibConfig( + skip_forward_without_activation_calib=True + ).skip_forward_without_activation_calib From dec014d11b5a8135c8a66d29a1ecf9166d910abc Mon Sep 17 00:00:00 2001 From: Zhiyu Cheng Date: Tue, 14 Jul 2026 12:03:51 -0700 Subject: [PATCH 3/3] fix(quant): account for static bias_calibrator in skip-forward check Address review feedback (juhi10071998): a quantizer with a static bias_calibrator collects data during the calibration forward (finish_stats_collection -> load_calib_bias), so it needs the forward even when its amax is dynamic/MX (which carry no per-tensor amax). _needs_activation_forward_for_max_calib now mirrors finish_stats_collection exactly: a non-constant quantizer with a static bias_calibrator forces the forward, while constant-amax quantizers (which are exempted from bias calibration) and dynamic-bias quantizers still do not. Add tests for MX + static/dynamic bias and constant + static bias. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Zhiyu Cheng --- CHANGELOG.rst | 2 +- modelopt/torch/quantization/model_calib.py | 39 ++++++++++++--------- tests/unit/torch/quantization/test_calib.py | 23 ++++++++++++ 3 files changed, 47 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4e0d092bdbe..ed32e4bd8e7 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -47,7 +47,7 @@ Changelog - 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. -- Add ``MaxCalibConfig.skip_forward_without_activation_calib`` (opt-in, default ``False``): when enabled, max calibration skips the ``forward_loop`` if no enabled quantizer needs data-driven activation statistics — e.g. an experts-only recipe whose activation quantizers all use ``constant_amax`` / ``use_constant_amax``, or dynamic / MX (MXFP4/MXFP8) quantization. Weight calibration still runs on the weight tensors directly, so the quantized weights are unchanged; only the wasted forward is avoided. It is opt-in because the ``forward_loop`` can carry caller-side effects (notably materializing sharded parameters under DeepSpeed ZeRO-3). The advanced algorithms that always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) call ``max_calibrate`` directly and are unaffected. The ``nvfp4_experts_only_input_scale1-kv_fp8_cast`` recipe enables it. +- Add ``MaxCalibConfig.skip_forward_without_activation_calib`` (opt-in, default ``False``): when enabled, max calibration skips the ``forward_loop`` if no enabled quantizer needs data-driven activation statistics — e.g. an experts-only recipe whose activation quantizers all use ``constant_amax`` / ``use_constant_amax``, or dynamic / MX (MXFP4/MXFP8) quantization (and none has a static ``bias_calibrator``). Weight calibration still runs on the weight tensors directly, so the quantized weights are unchanged; only the wasted forward is avoided. It is opt-in because the ``forward_loop`` can carry caller-side effects (notably materializing sharded parameters under DeepSpeed ZeRO-3). The advanced algorithms that always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) call ``max_calibrate`` directly and are unaffected. The ``nvfp4_experts_only_input_scale1-kv_fp8_cast`` recipe enables it. **Bug Fixes** diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index f06603dff16..031e766c687 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -258,16 +258,17 @@ def _needs_activation_forward_for_max_calib(model: nn.Module) -> bool: """Return True if any enabled quantizer still needs a calibration forward pass. Weight quantizers are calibrated directly on the weight tensors by - :func:`weight_only_quantize`, so they never need the data forward. Activation-side - quantizers (input/output/BMM) need it only when they collect data-driven statistics: - i.e. they are enabled, have a calibrator, and are none of - - - dynamic (top-level ``type: dynamic``), which computes amax on the fly; - - MX formats (MXFP4/MXFP8), whose per-block E8M0 scales are dynamic and which carry no - per-tensor amax (``is_mx_format`` makes ``amax`` always ``None``); - - pinned to a constant amax (``use_constant_amax`` / ``constant_amax``). Constant-amax - quantizers also skip bias calibration in :func:`finish_stats_collection`, so they never - require the forward either. + :func:`weight_only_quantize`, so they never need the data forward. An activation-side + quantizer (input/output/BMM) needs it when it collects data-driven statistics during the + forward, which :func:`finish_stats_collection` does in two ways: + + - **amax**: loaded for a quantizer that has a calibrator and is not dynamic (top-level + ``type: dynamic``), MX (MXFP4/MXFP8, whose E8M0 per-block scales are dynamic and which + carry no per-tensor amax), or pinned to a constant amax + (``use_constant_amax`` / ``constant_amax``); + - **static bias**: loaded for a quantizer with a static ``bias_calibrator``. Constant-amax + quantizers are exempted from calibration entirely (they ``continue`` before the bias + block), so only non-constant quantizers with a static bias need the forward for it. When this returns False (e.g. an experts-only recipe whose activation quantizers all use ``constant_amax``), the calibration forward can be skipped entirely and only weight @@ -280,12 +281,18 @@ def _needs_activation_forward_for_max_calib(model: nn.Module) -> bool: # are calibrated on the weight tensor directly, not via the data forward. if any(part.endswith("weight_quantizer") for part in name.split(".")): continue - if ( - module._dynamic - or module.is_mx_format - or module._use_constant_amax - or getattr(module, "_constant_amax", None) is not None - ): + + is_constant = ( + module._use_constant_amax or getattr(module, "_constant_amax", None) is not None + ) + + # A static bias calibrator collects data during the forward. Constant-amax quantizers + # skip bias calibration, so only non-constant quantizers with a static bias need it. + if not is_constant and module.bias_calibrator is not None and module.bias_type == "static": + return True + + # amax is data-driven only for a calibrated, non-dynamic, non-MX, non-constant quantizer. + if is_constant or module._dynamic or module.is_mx_format: continue if getattr(module, "_calibrator", None) is not None: return True diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index 27a5b3905a8..b609761f12a 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -672,6 +672,29 @@ def test_needs_activation_forward_ignores_mx_and_dynamic_quantizers(): assert not _needs_activation_forward_for_max_calib(dyn) +def test_needs_activation_forward_for_static_bias_calibrator(): + """A static bias calibrator collects data during the forward, even on MX/dynamic amax.""" + static_bias = {-1: None, "type": "static"} + mx_cfg = {"num_bits": (4, 3), "block_sizes": {-1: 32, "type": "dynamic", "scale_bits": (8, 0)}} + + # MX amax (no per-tensor amax) but a *static* bias -> the forward is still required. + mx_static_bias = _make_quant_linear(QuantizerAttributeConfig(**mx_cfg, bias=static_bias)) + assert mx_static_bias.input_quantizer.bias_type == "static" + assert _needs_activation_forward_for_max_calib(mx_static_bias) + + # A *dynamic* bias is computed on the fly, so no forward is needed. + mx_dynamic_bias = _make_quant_linear( + QuantizerAttributeConfig(**mx_cfg, bias={-1: None, "type": "dynamic"}) + ) + assert not _needs_activation_forward_for_max_calib(mx_dynamic_bias) + + # constant_amax exempts the quantizer from calibration entirely (bias included). + const_static_bias = _make_quant_linear( + QuantizerAttributeConfig(num_bits=8, constant_amax=127.0, bias=static_bias) + ) + assert not _needs_activation_forward_for_max_calib(const_static_bias) + + def test_max_calib_config_skip_is_opt_in(): """The flag is opt-in (default False) so it does not change behavior for direct callers.""" assert MaxCalibConfig().skip_forward_without_activation_calib is False