diff --git a/CHANGELOG.rst b/CHANGELOG.rst index df3cdefdb5e..ed32e4bd8e7 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. +- 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/config.py b/modelopt/torch/quantization/config.py index a8ed323a7b2..6c07afab90c 100644 --- a/modelopt/torch/quantization/config.py +++ b/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. diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index e03dff2e98b..031e766c687 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -254,6 +254,51 @@ 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.``) + # 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, @@ -261,6 +306,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 +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 ` 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) 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/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 64d89141bcd..b609761f12a 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,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