feat(quant): skip max-calib forward when no activation needs data

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/<model>/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 (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**
 

modelopt/torch/quantization/config.py

@@ -939,6 +939,23 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig):
         ),
     )
 
+    skip_forward_without_activation_calib: bool = ModeloptField(
+        default=False,
+        title="Skip the calibration forward when no activation quantizer needs data.",
+        description=(
+            "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."
+        ),
+    )
+
 
 class MseCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig):
     """Configuration for per-tensor MSE calibration.

modelopt/torch/quantization/model_calib.py

@@ -254,13 +254,59 @@ 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. 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
+    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.<i>``)
+        # are calibrated on the weight tensor directly, not via the data forward.
+        if any(part.endswith("weight_quantizer") for part in name.split(".")):
+            continue
+
+        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
+    return False
+
+
 @torch.no_grad()
 def max_calibrate(
     model: nn.Module,
     forward_loop: ForwardLoop | None = None,
     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 +320,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 <modelopt.torch.quantization.config.MaxCalibConfig>` for
     details on the remaining arguments.
@@ -292,7 +343,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)

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

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'

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,119 @@ 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_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_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
+    assert MaxCalibConfig(
+        skip_forward_without_activation_calib=True
+    ).skip_forward_without_activation_calib