From 14a74cbcbd1d1d7b9293969201933c6b11608301 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 28 Jul 2026 23:55:29 +0000 Subject: [PATCH 1/9] Add nvfp4_act_headroom activation calibration for NVFP4 NVFP4 scales a tensor with an FP8-E4M3 scale per 16-element block plus one per-tensor global scale. Plain max calibration sets that global scale from the largest block seen during calibration, leaving no room above it: any activation larger than the calibration max saturates. nvfp4_act_headroom instead anchors the global scale to a low percentile of the per-block amax distribution, amax = max(rho * anchor, floor), placing the calibrated blocks in the lower part of the FP8 scale range and leaving the rest as headroom. A high-percentile floor keeps the scale at or above the calibrated blocks so calibration data is never clipped, and a guardrail warns when the observed range is wider than FP8 scales can represent. Applies to NVFP4 dynamic-block input quantizers only; weight quantizers and all other quantizers are calibrated with plain max. - calib/nvfp4_act_headroom.py: NVFP4ActHeadroomCalibrator - config.py: NVFP4ActHeadroomCalibConfig (anchor_percentile, rho) - model_calib.py: nvfp4_act_headroom_calibrate - mode.py: NVFP4ActHeadroomCalibrateModeDescriptor - modelopt_recipes: a4_nvfp4 unit and nvfp4_act_only_headroom-kv_fp8_cast, an activation-only recipe whose module coverage matches nvfp4_default-kv_fp8_cast Pre-commit hooks were run manually on these files (all pass); the commit hook is skipped only because its autostash cannot write to a read-only cache in this environment. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 4 + modelopt/torch/quantization/calib/__init__.py | 1 + .../quantization/calib/nvfp4_act_headroom.py | 188 +++++++++++++++++ modelopt/torch/quantization/config.py | 46 +++++ modelopt/torch/quantization/mode.py | 19 ++ modelopt/torch/quantization/model_calib.py | 80 +++++++- .../ptq/nvfp4_act_headroom-kv_fp8_cast.yaml | 41 ++++ .../quantization/test_nvfp4_act_headroom.py | 191 ++++++++++++++++++ 8 files changed, 569 insertions(+), 1 deletion(-) create mode 100644 modelopt/torch/quantization/calib/nvfp4_act_headroom.py create mode 100644 modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml create mode 100644 tests/unit/torch/quantization/test_nvfp4_act_headroom.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 329e5d21f05..4b04e2423ea 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,10 @@ Changelog **New Features** +*Quantization* + +- Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Plain ``max`` sets a tensor's global scale from the largest per-block amax seen during calibration, leaving no room above it, so any activation larger than the calibration max saturates. ``nvfp4_act_headroom`` instead anchors the global scale to a low percentile of the per-block amax distribution, ``amax = max(rho * anchor, floor)``, placing the calibrated blocks in the lower part of the FP8 block-scale range and leaving the rest as headroom; a high-percentile ``floor`` keeps the scale at or above the calibrated blocks so calibration data is never clipped, and a guardrail warns when the observed range is wider than FP8 scales can represent. Tunable via ``anchor_percentile`` (default 1) and ``rho`` (default 16384). Applies only to NVFP4 dynamic-block input quantizers; weight quantizers and all other quantizers keep plain ``max``. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` (dynamic NVFP4 W4A4 plus FP8 KV-cache cast) with only the calibration algorithm swapped, so it exports a standard NVFP4 checkpoint. + **Backward Breaking Changes** **Deprecations** diff --git a/modelopt/torch/quantization/calib/__init__.py b/modelopt/torch/quantization/calib/__init__.py index 9accd0af67e..b48e3374886 100644 --- a/modelopt/torch/quantization/calib/__init__.py +++ b/modelopt/torch/quantization/calib/__init__.py @@ -24,3 +24,4 @@ from .histogram import * from .max import * from .mse import * +from .nvfp4_act_headroom import * diff --git a/modelopt/torch/quantization/calib/nvfp4_act_headroom.py b/modelopt/torch/quantization/calib/nvfp4_act_headroom.py new file mode 100644 index 00000000000..959dbb0499f --- /dev/null +++ b/modelopt/torch/quantization/calib/nvfp4_act_headroom.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Headroom calibrator for the NVFP4 activation global scale.""" + +import warnings + +import torch + +from ..utils import reduce_block_amax +from .calibrator import _Calibrator + +__all__ = ["NVFP4ActHeadroomCalibrator"] + +# FP8-E4M3 normal dynamic range (max_normal / min_normal = 448 / 2**-6). Per-block scales +# outside this ratio cannot all be represented as normal FP8 values by one global scale. +_FP8_NORMAL_DYNAMIC_RANGE = 28672.0 + +# High percentile of the per-block amax distribution used as the no-clipping floor. +_FLOOR_PERCENTILE = 99.99 + +# Per-block amaxes below ``floor / _ANCHOR_FLOOR_RATIO`` are ignored when locating the +# anchor, so a tail of near-zero blocks cannot drag the global scale down. +_ANCHOR_FLOOR_RATIO = 1e6 + + +class NVFP4ActHeadroomCalibrator(_Calibrator): + """Calibrates the NVFP4 activation global scale with headroom for unseen outliers. + + NVFP4 scales a tensor in two levels: an FP8-E4M3 scale per 16-element block, and one + per-tensor global scale (``amax``). Plain max calibration sets the global scale from the + largest block seen during calibration, which leaves no room above it: any activation + larger than the calibration max saturates. + + This calibrator instead anchors the global scale to a low percentile of the per-block + amax distribution:: + + amax = max(rho * anchor, floor) + + where ``anchor`` is the per-block amax at ``anchor_percentile`` and ``floor`` is a high + percentile of the same distribution. Multiplying the low anchor by ``rho`` places the + calibrated blocks in the lower part of the FP8 scale range and leaves the rest as upward + headroom, while ``floor`` keeps the scale at or above the calibrated blocks so + calibration data is never clipped. + + Args: + num_bits: quantizer ``num_bits`` (``(2, 1)`` for NVFP4); kept for interface parity. + axis: unused (the global scale is per-tensor); kept for interface parity. + unsigned: unused; kept for interface parity. + block_size: NVFP4 block width along the last dim. + anchor_percentile: percentile of the per-block amaxes used as the anchor. + rho: headroom factor applied to the anchor. Must be in ``(0, 28672)``. + num_bins: number of log2 histogram bins. + log2_min: low end of the log2 range covered by the histogram. + log2_max: high end of the log2 range covered by the histogram. + """ + + def __init__( + self, + num_bits=(2, 1), + axis=None, + unsigned=False, + *, + block_size=16, + anchor_percentile=1.0, + rho=16384.0, + num_bins=512, + log2_min=-40.0, + log2_max=40.0, + ): + """Initialize.""" + super().__init__(num_bits, axis, unsigned) + if not (0.0 < rho < _FP8_NORMAL_DYNAMIC_RANGE): + raise ValueError( + f"rho must be in (0, {_FP8_NORMAL_DYNAMIC_RANGE}); got {rho}. Larger rho gives " + "more headroom above the calibrated range but less room below it." + ) + if not (0.0 <= anchor_percentile <= 100.0): + raise ValueError(f"anchor_percentile must be in [0, 100]; got {anchor_percentile}.") + self._block_size = block_size + self._anchor_percentile = float(anchor_percentile) + self._rho = float(rho) + self._num_bins = int(num_bins) + self._log2_min = float(log2_min) + self._log2_max = float(log2_max) + self._hist = None # int64 histogram over log2(per-block amax), created on first collect + self._running_max = None # literal per-tensor max, used as the fallback amax + + def _bin_index(self, log2_vals: torch.Tensor) -> torch.Tensor: + frac = (log2_vals - self._log2_min) / (self._log2_max - self._log2_min) + idx = (frac * self._num_bins).floor().long() + return idx.clamp_(0, self._num_bins - 1) + + @torch.no_grad() + def collect(self, x: torch.Tensor) -> None: + """Accumulate the per-block amax histogram for one activation batch.""" + # reduce_block_amax already takes abs(). + block_amax = reduce_block_amax(x.detach(), block_sizes={-1: self._block_size}) + block_amax = block_amax.flatten().float() + + cur_max = block_amax.max() + self._running_max = ( + cur_max if self._running_max is None else torch.maximum(self._running_max, cur_max) + ) + + nonzero = block_amax[block_amax > 0] + if nonzero.numel() == 0: + return + + idx = self._bin_index(torch.log2(nonzero)) + counts = torch.bincount(idx, minlength=self._num_bins) + if self._hist is None: + self._hist = torch.zeros(self._num_bins, dtype=torch.long, device=x.device) + self._hist += counts.to(self._hist.device) + + def _percentile(self, percentile: float, floor_value: float | None = None) -> float | None: + """Value at ``percentile`` of the histogram, ignoring bins below ``floor_value``.""" + if self._hist is None: + return None + counts = self._hist.float().clone() + if floor_value is not None and floor_value > 0: + floor_bin = int(self._bin_index(torch.log2(torch.tensor(floor_value))).item()) + counts[:floor_bin] = 0 + total = counts.sum() + if total <= 0: + return None + target = percentile / 100.0 * total + cdf = torch.cumsum(counts, dim=0) + bin_idx = int(torch.searchsorted(cdf, target).clamp(0, self._num_bins - 1).item()) + # Bin center back to a value in linear space. + log2_val = self._log2_min + (bin_idx + 0.5) / self._num_bins * ( + self._log2_max - self._log2_min + ) + return float(2.0**log2_val) + + @torch.no_grad() + def compute_amax(self) -> torch.Tensor | None: + """Return the calibrated NVFP4 activation global scale.""" + if self._hist is None or self._running_max is None: + return None + + floor = self._percentile(_FLOOR_PERCENTILE) + anchor = ( + self._percentile( + self._anchor_percentile, floor_value=floor / _ANCHOR_FLOOR_RATIO if floor else None + ) + if floor + else None + ) + if not floor or floor <= 0 or not anchor or anchor <= 0: + # Percentiles collapsed despite a non-empty histogram: fall back to plain max. + return self._running_max.clone() + + if floor / anchor > _FP8_NORMAL_DYNAMIC_RANGE: + warnings.warn( + f"[nvfp4_act_headroom] per-block amax range {floor / anchor:.1f} exceeds the FP8 " + f"scale range ({_FP8_NORMAL_DYNAMIC_RANGE:.0f}); no global scale can hold both " + "ends. Using the no-clipping floor, which leaves no headroom. Reduce the range " + "with outlier mitigation (SmoothQuant / per-channel / higher precision)." + ) + amax = floor + else: + amax = max(self._rho * anchor, floor) + + return torch.tensor(float(amax), dtype=torch.float32, device=self._running_max.device) + + def reset(self) -> None: + """Reset the collected histogram.""" + self._hist = None + self._running_max = None + + def __repr__(self): + return ( + f"NVFP4ActHeadroomCalibrator(anchor_percentile={self._anchor_percentile}, " + f"rho={self._rho}, block_size={self._block_size})" + ) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 6257623d108..3e119e01ed2 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -986,6 +986,52 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): ) +class NVFP4ActHeadroomCalibConfig(QuantizeAlgorithmConfig): + """Config for the ``nvfp4_act_headroom`` calibration algorithm. + + Calibrates the per-tensor global scale of NVFP4 *activation* (input) quantizers so that + the calibrated range sits in the lower part of the FP8 block-scale range, leaving headroom + above it for activations larger than any seen during calibration. Weight quantizers and + all non-NVFP4 quantizers are calibrated with plain ``max``. + + See :class:`NVFP4ActHeadroomCalibrator + ` for the formula. + """ + + _mutates_weights: ClassVar[bool] = False + + method: Literal["nvfp4_act_headroom"] = ModeloptField("nvfp4_act_headroom") + + anchor_percentile: float = ModeloptField( + default=1.0, + ge=0.0, + le=100.0, + title="Percentile of the per-block activation amaxes used as the anchor.", + description=( + "The global scale is anchored to this percentile of the per-block amax " + "distribution. Lower values anchor further into the low tail, which yields a " + "smaller global scale and less headroom." + ), + ) + + rho: float = ModeloptField( + default=16384.0, + gt=0.0, + lt=28672.0, + title="Headroom factor applied to the anchor (amax = rho * anchor).", + description=( + "Larger rho leaves more headroom above the calibrated range and less room below " + "it. Must stay below 28672, the FP8-E4M3 normal dynamic range." + ), + ) + + distributed_sync: bool | None = ModeloptField( + default=True, + title="Whether to sync the amax across the distributed processes.", + description="If True, the amax will be synced across the distributed processes.", + ) + + class MseCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): """Configuration for per-tensor MSE calibration. diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py index 8f58d4c9dbb..c096aaeb00e 100644 --- a/modelopt/torch/quantization/mode.py +++ b/modelopt/torch/quantization/mode.py @@ -42,6 +42,7 @@ LSQConfig, MaxCalibConfig, MseCalibConfig, + NVFP4ActHeadroomCalibConfig, QuantizeAlgoCfgType, QuantizeAlgorithmConfig, QuantizeConfig, @@ -66,6 +67,7 @@ lsq, max_calibrate, mse_calibrate, + nvfp4_act_headroom_calibrate, smoothquant, svdquant, ) @@ -429,6 +431,23 @@ def config_class(self) -> type[QuantizeAlgorithmConfig]: _calib_func = max_calibrate +@CalibrateModeRegistry.register_mode +class NVFP4ActHeadroomCalibrateModeDescriptor(BaseCalibrateModeDescriptor): + """Mode for the ``nvfp4_act_headroom`` calibration algorithm. + + Headroom-aware global scales for NVFP4 activation quantizers; plain max for everything + else (see :class:`NVFP4ActHeadroomCalibConfig + `). + """ + + @property + def config_class(self) -> type[QuantizeAlgorithmConfig]: + """Specifies the config class for the mode.""" + return NVFP4ActHeadroomCalibConfig + + _calib_func = nvfp4_act_headroom_calibrate + + @CalibrateModeRegistry.register_mode class MseCalibrateModeDescriptor(BaseCalibrateModeDescriptor): """Mode for mse calibration algorithm.""" diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 3fe38610a74..0fb6fe18d65 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -41,7 +41,7 @@ from modelopt.torch.utils.distributed import size as dist_size from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method -from .calib import MseCalibrator, NVFP4MSECalibrator, _Calibrator +from .calib import MseCalibrator, NVFP4ActHeadroomCalibrator, NVFP4MSECalibrator, _Calibrator from .conversion import create_and_replace_svdquant_linear_on_the_fly, set_quantizer_by_cfg_context from .nn import QuantModule, SequentialQuantizer, StaticBlockScaleQuantizer, TensorQuantizer from .utils import ( @@ -66,6 +66,7 @@ "local_hessian_calibrate", "lsq", "max_calibrate", + "nvfp4_act_headroom_calibrate", "smoothquant", "svdquant", ] @@ -490,6 +491,83 @@ def sync_quantizer_amax_across_tp( _finalize_with_shared_state(model, weight_patterns) +def _is_nvfp4_dynamic_input_quantizer(name: str, module: nn.Module) -> bool: + """True for an enabled NVFP4 (E2M1 + E4M3 scale) dynamic-block *input* quantizer. + + These carry their per-tensor global scale in ``_amax`` and are calibrated like plain max + (the quantizer itself is not ``_dynamic``; only its per-block scales are). + """ + if not isinstance(module, TensorQuantizer) or not name.endswith("input_quantizer"): + return False + if getattr(module, "_disabled", False) or getattr(module, "_dynamic", False): + return False + block_sizes = module.block_sizes + return ( + block_sizes is not None + and block_sizes.get("type", None) == "dynamic" + and module.num_bits == (2, 1) + and block_sizes.get("scale_bits", None) == (4, 3) + ) + + +def _swap_in_nvfp4_act_headroom_calibrators( + model: nn.Module, *, anchor_percentile: float, rho: float +) -> int: + """Swap in an :class:`NVFP4ActHeadroomCalibrator` for each qualifying NVFP4 input quantizer. + + Returns the number of quantizers swapped. + """ + count = 0 + for name, module in model.named_modules(): + if not _is_nvfp4_dynamic_input_quantizer(name, module): + continue + module._calibrator = NVFP4ActHeadroomCalibrator( + module.num_bits, + None, + module._unsigned, + block_size=module.block_sizes.get(-1, 16), + anchor_percentile=anchor_percentile, + rho=rho, + ) + count += 1 + return count + + +@torch.no_grad() +def nvfp4_act_headroom_calibrate( + model: nn.Module, + forward_loop: ForwardLoop | None = None, + *, + anchor_percentile: float = 1.0, + rho: float = 16384.0, + distributed_sync: bool = True, +): + """Calibrate NVFP4 activation global scales with headroom; everything else uses max. + + For NVFP4 dynamic-block *input* quantizers, the per-tensor global scale is derived from + the distribution of per-block activation amaxes so that headroom is left above the + calibrated range (see :class:`NVFP4ActHeadroomCalibrator + `). Weight quantizers and + every other quantizer are calibrated exactly as in :func:`max_calibrate`. + + Args: + model: model to be calibrated. + forward_loop: callable that runs calibration data through the model. + anchor_percentile: percentile of the per-block amaxes used as the anchor. + rho: headroom factor; ``amax = rho * anchor``. Must be in ``(0, 28672)``. + distributed_sync: whether to sync amax across distributed processes + (see :func:`max_calibrate`). + """ + n = _swap_in_nvfp4_act_headroom_calibrators(model, anchor_percentile=anchor_percentile, rho=rho) + print_rank_0( + f"nvfp4_act_headroom: calibrating {n} NVFP4 activation quantizer(s) " + f"(anchor_percentile={anchor_percentile}, rho={rho})." + ) + # max_calibrate runs the forward loop once: the swapped-in calibrators accumulate their + # per-block histograms in the same pass that collects max stats for every other quantizer. + max_calibrate(model, forward_loop, distributed_sync=distributed_sync) + + def _mse_quant_func(x, amax, quantizer): """Quantization function for MSE calibration.""" original_amax = quantizer._amax.clone() if hasattr(quantizer, "_amax") else None diff --git a/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml b/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml new file mode 100644 index 00000000000..a436e57759a --- /dev/null +++ b/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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 dynamic NVFP4 W4A4 model quantization with FP8 KV-cache cast mode, +# using nvfp4_act_headroom calibration: NVFP4 weights on plain max, NVFP4 activation global +# scales anchored with headroom. Module coverage matches nvfp4_default-kv_fp8_cast. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + w4a4_nvfp4_nvfp4: configs/ptq/units/w4a4_nvfp4_nvfp4 + kv_fp8_cast: configs/ptq/units/kv_fp8_cast + +metadata: + recipe_type: ptq + description: >- + Composes dynamic NVFP4 W4A4 model quantization with FP8 KV-cache cast mode using constant + amax; uses nvfp4_act_headroom calibration. Weight quantizers use plain max, while the NVFP4 + activation global scales are anchored to a low percentile of the per-block amax distribution + to leave headroom for activations larger than any seen during calibration. Module coverage + matches nvfp4_default-kv_fp8_cast. +quantize: + algorithm: + method: nvfp4_act_headroom + anchor_percentile: 1 + quant_cfg: + - $import: base_disable_all + - $import: w4a4_nvfp4_nvfp4 + - $import: kv_fp8_cast + - $import: default_disabled_quantizers diff --git a/tests/unit/torch/quantization/test_nvfp4_act_headroom.py b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py new file mode 100644 index 00000000000..9fc90b92133 --- /dev/null +++ b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Tests for the ``nvfp4_act_headroom`` activation global-scale calibration.""" + +import pytest +import torch + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.calib import NVFP4ActHeadroomCalibrator +from modelopt.torch.quantization.model_calib import ( + _is_nvfp4_dynamic_input_quantizer, + _swap_in_nvfp4_act_headroom_calibrators, +) +from modelopt.torch.quantization.nn import TensorQuantizer + +NVFP4_CFG = { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, +} + +ACT_ONLY_CFG = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*input_quantizer", "cfg": NVFP4_CFG}, + ], + "algorithm": {"method": "nvfp4_act_headroom", "anchor_percentile": 1}, +} + + +class _Net(torch.nn.Module): + def __init__(self): + super().__init__() + self.fc1 = torch.nn.Linear(64, 64) + self.fc2 = torch.nn.Linear(64, 64) + + def forward(self, x): + return self.fc2(self.fc1(x)) + + +def _calibrate(calibrator, x): + calibrator.collect(x) + return float(calibrator.compute_amax()) + + +def test_amax_is_rho_times_anchor(): + """On a uniform per-block distribution the anchor term sets amax = rho * anchor.""" + torch.manual_seed(0) + # All blocks share one magnitude, so anchor == floor and rho * anchor dominates. + x = torch.full((256, 64), 0.5) + amax = _calibrate(NVFP4ActHeadroomCalibrator(rho=1024.0, anchor_percentile=1.0), x) + assert amax == pytest.approx(1024.0 * 0.5, rel=0.05) + + +def test_amax_never_below_calibrated_max(): + """A range wider than FP8 can hold warns and falls back to the no-clipping floor.""" + x = torch.zeros(64, 64) + x[0, :16] = 100.0 # one large block, the rest tiny + x[1:, :] = 1e-3 + with pytest.warns(UserWarning, match="exceeds the FP8 scale range"): + amax = _calibrate(NVFP4ActHeadroomCalibrator(rho=1.0, anchor_percentile=1.0), x) + assert amax >= 100.0 * 0.9 + + +def test_lower_anchor_percentile_gives_smaller_amax(): + """anchor_percentile is tunable and monotone: a lower percentile anchors lower.""" + torch.manual_seed(0) + x = torch.randn(512, 64).abs() + 1e-4 + amax_p1 = _calibrate(NVFP4ActHeadroomCalibrator(anchor_percentile=1.0), x) + amax_p50 = _calibrate(NVFP4ActHeadroomCalibrator(anchor_percentile=50.0), x) + assert amax_p1 < amax_p50 + + +def test_headroom_exceeds_plain_max(): + """The calibrated scale leaves headroom above what plain max would pick.""" + torch.manual_seed(0) + x = torch.randn(512, 64).abs() + 1e-4 + amax = _calibrate(NVFP4ActHeadroomCalibrator(anchor_percentile=1.0), x) + assert amax > float(x.abs().max()) + + +def test_all_zero_activation_yields_no_scale(): + """An all-zero activation carries no scale information, so no amax is inferred.""" + calibrator = NVFP4ActHeadroomCalibrator() + calibrator.collect(torch.zeros(32, 64)) + assert calibrator.compute_amax() is None + + +def test_rho_out_of_range_rejected(): + with pytest.raises(ValueError, match="rho must be in"): + NVFP4ActHeadroomCalibrator(rho=28672.0) + + +def test_reset_clears_state(): + calibrator = NVFP4ActHeadroomCalibrator() + calibrator.collect(torch.randn(32, 64).abs() + 1e-3) + calibrator.reset() + assert calibrator.compute_amax() is None + + +def test_only_nvfp4_input_quantizers_are_selected(): + """The swap targets NVFP4 dynamic-block input quantizers, not weights or other formats.""" + nvfp4_in = TensorQuantizer(mtq.config.QuantizerAttributeConfig(**NVFP4_CFG)) + assert _is_nvfp4_dynamic_input_quantizer("layer.input_quantizer", nvfp4_in) + # Same config but a weight quantizer -> not selected. + assert not _is_nvfp4_dynamic_input_quantizer("layer.weight_quantizer", nvfp4_in) + # FP8 input quantizer -> not selected. + fp8_in = TensorQuantizer(mtq.config.QuantizerAttributeConfig(num_bits=(4, 3))) + assert not _is_nvfp4_dynamic_input_quantizer("layer.input_quantizer", fp8_in) + + +def test_swap_installs_calibrator_with_config_values(): + model = mtq.quantize(_Net(), ACT_ONLY_CFG, lambda m: m(torch.randn(8, 64))) + n = _swap_in_nvfp4_act_headroom_calibrators(model, anchor_percentile=5.0, rho=1024.0) + assert n == 2 # fc1 and fc2 input quantizers + cal = model.fc1.input_quantizer._calibrator + assert isinstance(cal, NVFP4ActHeadroomCalibrator) + assert cal._anchor_percentile == 5.0 + assert cal._rho == 1024.0 + + +def test_activation_only_quantize_leaves_weights_untouched(): + """End-to-end: input quantizers get a scale, weight quantizers stay disabled.""" + torch.manual_seed(0) + model = _Net() + weights_before = {n: p.clone() for n, p in model.named_parameters()} + + model = mtq.quantize(model, ACT_ONLY_CFG, lambda m: m(torch.randn(16, 64))) + + for name in ("fc1", "fc2"): + layer = getattr(model, name) + assert layer.input_quantizer.is_enabled + assert layer.input_quantizer.amax is not None + assert float(layer.input_quantizer.amax) > 0 + assert not layer.weight_quantizer.is_enabled + + for n, p in model.named_parameters(): + if n in weights_before: + assert torch.equal(p, weights_before[n]), f"{n} was modified" + + +def test_w4a4_weights_use_max_activations_use_headroom(): + """W4A4: weights fall back to plain max, only activations get the headroom scale.""" + torch.manual_seed(0) + data = torch.randn(16, 64) + w4a4_cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": NVFP4_CFG}, + {"quantizer_name": "*input_quantizer", "cfg": NVFP4_CFG}, + ], + "algorithm": {"method": "nvfp4_act_headroom", "anchor_percentile": 1}, + } + model = _Net() + weight_max = float(model.fc1.weight.abs().max()) + model = mtq.quantize(model, w4a4_cfg, lambda m: m(data)) + + # Weight quantizer is calibrated with plain max: amax == the literal weight max. + assert model.fc1.weight_quantizer.is_enabled + assert float(model.fc1.weight_quantizer.amax) == pytest.approx(weight_max, rel=1e-5) + # Activation quantizer gets the headroom scale, well above the observed activation max. + assert float(model.fc1.input_quantizer.amax) > float(data.abs().max()) + + +def test_anchor_percentile_changes_model_scales(): + """The knob propagates through mtq.quantize to the calibrated activation scales.""" + torch.manual_seed(0) + data = torch.randn(16, 64) + + def _amax_for(percentile): + torch.manual_seed(0) + cfg = { + **ACT_ONLY_CFG, + "algorithm": {"method": "nvfp4_act_headroom", "anchor_percentile": percentile}, + } + model = mtq.quantize(_Net(), cfg, lambda m: m(data)) + return float(model.fc1.input_quantizer.amax) + + assert _amax_for(1) < _amax_for(50) From 2076b6d0a52d6c3fe2aba854ecf6337a20041a28 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Wed, 29 Jul 2026 17:59:36 +0000 Subject: [PATCH 2/9] Address review: no-clipping invariant, calibrator restore, recipe docs - calib/nvfp4_act_headroom.py: floor the result at the largest observed per-block amax instead of the 99.99th percentile. The percentile floor could land below the observed max (bin centers round down, and 0.01% of blocks sit above it by construction), so the guardrail path could return a scale that clips calibration data -- worse than plain max. Warn whenever the anchor term cannot clear the observed max (headroom collapses at rho, not at the FP8 window), zero pad a ragged last dim instead of tripping the divisibility assert in reduce_block_amax, and assert on NaN/Inf as MaxCalibrator does. - Reject anchor_percentile=0 in both the config and the calibrator: searchsorted returns bin 0 for a zero target regardless of the low-tail mask, pinning the anchor to 2**-40. - model_calib.py: restore the original calibrators after calibration so the algorithm cannot leak into a later max_calibrate on the same model, and warn when no NVFP4 activation quantizer matches. - Document the recipe in modelopt_recipes/ptq.md and add it to the builtin recipe list in tests/unit/recipe/test_loader.py (two recipe-doc tests were failing). - Tests: assert the no-clipping invariant strictly across six distributions including the guardrail regime, plus regressions for percentile 0, ragged last dim, NaN/Inf, and calibrator restore. Pre-commit hooks were run manually on these files (all pass); the commit hook is skipped only because its autostash cannot write to a read-only cache in this environment. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 2 +- .../quantization/calib/nvfp4_act_headroom.py | 63 +++++++++---- modelopt/torch/quantization/config.py | 2 +- modelopt/torch/quantization/model_calib.py | 32 +++++-- modelopt_recipes/ptq.md | 3 +- tests/unit/recipe/test_loader.py | 1 + .../quantization/test_nvfp4_act_headroom.py | 88 +++++++++++++++++-- 7 files changed, 156 insertions(+), 35 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4b04e2423ea..3ef5c89276e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,7 +8,7 @@ Changelog *Quantization* -- Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Plain ``max`` sets a tensor's global scale from the largest per-block amax seen during calibration, leaving no room above it, so any activation larger than the calibration max saturates. ``nvfp4_act_headroom`` instead anchors the global scale to a low percentile of the per-block amax distribution, ``amax = max(rho * anchor, floor)``, placing the calibrated blocks in the lower part of the FP8 block-scale range and leaving the rest as headroom; a high-percentile ``floor`` keeps the scale at or above the calibrated blocks so calibration data is never clipped, and a guardrail warns when the observed range is wider than FP8 scales can represent. Tunable via ``anchor_percentile`` (default 1) and ``rho`` (default 16384). Applies only to NVFP4 dynamic-block input quantizers; weight quantizers and all other quantizers keep plain ``max``. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` (dynamic NVFP4 W4A4 plus FP8 KV-cache cast) with only the calibration algorithm swapped, so it exports a standard NVFP4 checkpoint. +- Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Plain ``max`` sets a tensor's global scale from the largest per-block amax seen during calibration, leaving no room above it, so any activation larger than the calibration max saturates. ``nvfp4_act_headroom`` instead anchors the global scale to a low percentile of the per-block amax distribution, ``amax = max(rho * anchor, floor)``, placing the calibrated blocks in the lower part of the FP8 block-scale range and leaving the rest as headroom; the result is floored at the largest per-block amax observed during calibration so the scale never sits below the calibrated data, and the calibrator warns and degenerates to plain ``max`` when the per-block range is too wide for ``rho`` to clear. Tunable via ``anchor_percentile`` (default 1) and ``rho`` (default 16384). Applies only to NVFP4 dynamic-block input quantizers; weight quantizers and all other quantizers keep plain ``max``. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` (dynamic NVFP4 W4A4 plus FP8 KV-cache cast) with only the calibration algorithm swapped, so it exports a standard NVFP4 checkpoint. **Backward Breaking Changes** diff --git a/modelopt/torch/quantization/calib/nvfp4_act_headroom.py b/modelopt/torch/quantization/calib/nvfp4_act_headroom.py index 959dbb0499f..2d46bc90ac7 100644 --- a/modelopt/torch/quantization/calib/nvfp4_act_headroom.py +++ b/modelopt/torch/quantization/calib/nvfp4_act_headroom.py @@ -47,20 +47,21 @@ class NVFP4ActHeadroomCalibrator(_Calibrator): This calibrator instead anchors the global scale to a low percentile of the per-block amax distribution:: - amax = max(rho * anchor, floor) + amax = max(rho * anchor, observed_max) - where ``anchor`` is the per-block amax at ``anchor_percentile`` and ``floor`` is a high - percentile of the same distribution. Multiplying the low anchor by ``rho`` places the - calibrated blocks in the lower part of the FP8 scale range and leaves the rest as upward - headroom, while ``floor`` keeps the scale at or above the calibrated blocks so - calibration data is never clipped. + where ``anchor`` is the per-block amax at ``anchor_percentile``. Multiplying the low anchor + by ``rho`` places the calibrated blocks in the lower part of the FP8 scale range and leaves + the rest as upward headroom. The result is floored at the largest per-block amax observed + during calibration, so the scale never sits below the calibrated data; when that floor wins + (a per-block range too wide for ``rho`` to clear) the calibrator warns and the result + degenerates to plain ``max``. Args: num_bits: quantizer ``num_bits`` (``(2, 1)`` for NVFP4); kept for interface parity. axis: unused (the global scale is per-tensor); kept for interface parity. unsigned: unused; kept for interface parity. block_size: NVFP4 block width along the last dim. - anchor_percentile: percentile of the per-block amaxes used as the anchor. + anchor_percentile: percentile of the per-block amaxes used as the anchor; must be > 0. rho: headroom factor applied to the anchor. Must be in ``(0, 28672)``. num_bins: number of log2 histogram bins. log2_min: low end of the log2 range covered by the histogram. @@ -87,8 +88,8 @@ def __init__( f"rho must be in (0, {_FP8_NORMAL_DYNAMIC_RANGE}); got {rho}. Larger rho gives " "more headroom above the calibrated range but less room below it." ) - if not (0.0 <= anchor_percentile <= 100.0): - raise ValueError(f"anchor_percentile must be in [0, 100]; got {anchor_percentile}.") + if not (0.0 < anchor_percentile <= 100.0): + raise ValueError(f"anchor_percentile must be in (0, 100]; got {anchor_percentile}.") self._block_size = block_size self._anchor_percentile = float(anchor_percentile) self._rho = float(rho) @@ -106,10 +107,24 @@ def _bin_index(self, log2_vals: torch.Tensor) -> torch.Tensor: @torch.no_grad() def collect(self, x: torch.Tensor) -> None: """Accumulate the per-block amax histogram for one activation batch.""" + x = x.detach() + # Zero-pad a ragged tail so activations whose last dim is not a multiple of the block + # size are handled instead of tripping the divisibility assert in reduce_block_amax. + # Padding with zeros cannot change any block's amax. + remainder = x.shape[-1] % self._block_size + if remainder: + x = torch.nn.functional.pad(x, (0, self._block_size - remainder)) # reduce_block_amax already takes abs(). - block_amax = reduce_block_amax(x.detach(), block_sizes={-1: self._block_size}) + block_amax = reduce_block_amax(x, block_sizes={-1: self._block_size}) block_amax = block_amax.flatten().float() + assert not torch.any(torch.isnan(block_amax)), ( + f"detected nan values in amax. nan in original tensor: {torch.any(torch.isnan(x))}" + ) + assert not torch.any(torch.isinf(block_amax)), ( + f"detected inf values in amax. inf in original tensor: {torch.any(torch.isinf(x))}" + ) + cur_max = block_amax.max() self._running_max = ( cur_max if self._running_max is None else torch.maximum(self._running_max, cur_max) @@ -129,7 +144,7 @@ def _percentile(self, percentile: float, floor_value: float | None = None) -> fl """Value at ``percentile`` of the histogram, ignoring bins below ``floor_value``.""" if self._hist is None: return None - counts = self._hist.float().clone() + counts = self._hist.float() if floor_value is not None and floor_value > 0: floor_bin = int(self._bin_index(torch.log2(torch.tensor(floor_value))).item()) counts[:floor_bin] = 0 @@ -151,6 +166,7 @@ def compute_amax(self) -> torch.Tensor | None: if self._hist is None or self._running_max is None: return None + running_max = float(self._running_max) floor = self._percentile(_FLOOR_PERCENTILE) anchor = ( self._percentile( @@ -163,16 +179,27 @@ def compute_amax(self) -> torch.Tensor | None: # Percentiles collapsed despite a non-empty histogram: fall back to plain max. return self._running_max.clone() - if floor / anchor > _FP8_NORMAL_DYNAMIC_RANGE: + # The observed max is the no-clipping bound: the returned scale never sits below it, so + # calibration data is never clipped. Headroom is whatever the anchor term adds on top. + headroom_amax = self._rho * anchor + if headroom_amax <= running_max: + exceeds_window = floor / anchor > _FP8_NORMAL_DYNAMIC_RANGE warnings.warn( - f"[nvfp4_act_headroom] per-block amax range {floor / anchor:.1f} exceeds the FP8 " - f"scale range ({_FP8_NORMAL_DYNAMIC_RANGE:.0f}); no global scale can hold both " - "ends. Using the no-clipping floor, which leaves no headroom. Reduce the range " - "with outlier mitigation (SmoothQuant / per-channel / higher precision)." + f"[nvfp4_act_headroom] per-block amax range {floor / anchor:.1f} leaves no " + f"headroom at rho={self._rho:g}; falling back to the observed max. " + + ( + f"The range also exceeds the FP8 scale range " + f"({_FP8_NORMAL_DYNAMIC_RANGE:.0f}), so no global scale can hold both ends. " + "Reduce the range with outlier mitigation (SmoothQuant / per-channel / " + "higher precision)." + if exceeds_window + else "Lower anchor_percentile or raise rho for more headroom." + ), + stacklevel=2, ) - amax = floor + amax = running_max else: - amax = max(self._rho * anchor, floor) + amax = headroom_amax return torch.tensor(float(amax), dtype=torch.float32, device=self._running_max.device) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 3e119e01ed2..9e9c2ad70a8 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -1004,7 +1004,7 @@ class NVFP4ActHeadroomCalibConfig(QuantizeAlgorithmConfig): anchor_percentile: float = ModeloptField( default=1.0, - ge=0.0, + gt=0.0, le=100.0, title="Percentile of the per-block activation amaxes used as the anchor.", description=( diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 0fb6fe18d65..9a257067c78 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -21,7 +21,7 @@ import warnings from collections.abc import Callable, Mapping, Sequence from functools import partial -from typing import TypeAlias +from typing import Any, TypeAlias import torch import torch.distributed as dist @@ -512,15 +512,16 @@ def _is_nvfp4_dynamic_input_quantizer(name: str, module: nn.Module) -> bool: def _swap_in_nvfp4_act_headroom_calibrators( model: nn.Module, *, anchor_percentile: float, rho: float -) -> int: +) -> list[tuple[nn.Module, Any]]: """Swap in an :class:`NVFP4ActHeadroomCalibrator` for each qualifying NVFP4 input quantizer. - Returns the number of quantizers swapped. + Returns ``(quantizer, original_calibrator)`` pairs so the caller can restore them. """ - count = 0 + swapped: list[tuple[nn.Module, Any]] = [] for name, module in model.named_modules(): if not _is_nvfp4_dynamic_input_quantizer(name, module): continue + swapped.append((module, module._calibrator)) module._calibrator = NVFP4ActHeadroomCalibrator( module.num_bits, None, @@ -529,8 +530,7 @@ def _swap_in_nvfp4_act_headroom_calibrators( anchor_percentile=anchor_percentile, rho=rho, ) - count += 1 - return count + return swapped @torch.no_grad() @@ -558,14 +558,28 @@ def nvfp4_act_headroom_calibrate( distributed_sync: whether to sync amax across distributed processes (see :func:`max_calibrate`). """ - n = _swap_in_nvfp4_act_headroom_calibrators(model, anchor_percentile=anchor_percentile, rho=rho) + swapped = _swap_in_nvfp4_act_headroom_calibrators( + model, anchor_percentile=anchor_percentile, rho=rho + ) + if not swapped: + warn_rank_0( + "nvfp4_act_headroom: no NVFP4 dynamic-block input quantizer matched, so this is " + "equivalent to plain max calibration. Check that the recipe enables NVFP4 " + "activation quantizers." + ) print_rank_0( - f"nvfp4_act_headroom: calibrating {n} NVFP4 activation quantizer(s) " + f"nvfp4_act_headroom: calibrating {len(swapped)} NVFP4 activation quantizer(s) " f"(anchor_percentile={anchor_percentile}, rho={rho})." ) # max_calibrate runs the forward loop once: the swapped-in calibrators accumulate their # per-block histograms in the same pass that collects max stats for every other quantizer. - max_calibrate(model, forward_loop, distributed_sync=distributed_sync) + # The calibrators are restored afterwards so this algorithm does not leak into a later + # calibration of the same model, and so a repeat run starts from a fresh histogram. + try: + max_calibrate(model, forward_loop, distributed_sync=distributed_sync) + finally: + for quantizer, original_calibrator in swapped: + quantizer._calibrator = original_calibrator def _mse_quant_func(x, amax, quantizer): diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 255544ffe1e..29da1ff5d45 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -29,7 +29,7 @@ supported combinations. ### The shipped recipes
-All 20 general/ptq/ recipes (click to expand) +All 21 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -37,6 +37,7 @@ supported combinations. | `fp8_default-kv_fp8_cast` | FP8 W8A8, all linears | FP8 (constant amax) | max | | `nvfp4_default-kv_fp8` | NVFP4 W4A4, all linears | FP8 (calibrated) | max | | `nvfp4_default-kv_fp8_cast` | NVFP4 W4A4, all linears | FP8 (constant amax) | max | +| `nvfp4_act_headroom-kv_fp8_cast` | NVFP4 W4A4, all linears | FP8 (constant amax) | nvfp4_act_headroom | | `nvfp4_default-kv_nvfp4_cast` | NVFP4 W4A4, all linears | NVFP4 (constant amax) | max | | `nvfp4_default-kv_none-gptq` | NVFP4 W4A4 (static W), all linears | none | GPTQ (layerwise) | | `nvfp4_mlp_only-kv_fp8` | NVFP4 W4A4, MLP + MoE experts | FP8 (calibrated) | max | diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index e15b897a224..d1f6579019d 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -159,6 +159,7 @@ def test_load_recipe_builtin_description(): "general/ptq/fp8_default-kv_fp8", "general/ptq/fp8_default-kv_fp8_cast", "general/ptq/int4_blockwise_weight_only", + "general/ptq/nvfp4_act_headroom-kv_fp8_cast", "general/ptq/nvfp4_default-kv_fp8", "general/ptq/nvfp4_default-kv_fp8_cast", "general/ptq/nvfp4_default-kv_nvfp4_cast", diff --git a/tests/unit/torch/quantization/test_nvfp4_act_headroom.py b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py index 9fc90b92133..293ba605a1c 100644 --- a/tests/unit/torch/quantization/test_nvfp4_act_headroom.py +++ b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py @@ -15,6 +15,8 @@ """Tests for the ``nvfp4_act_headroom`` activation global-scale calibration.""" +import warnings + import pytest import torch @@ -25,6 +27,7 @@ _swap_in_nvfp4_act_headroom_calibrators, ) from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.utils import reduce_block_amax NVFP4_CFG = { "num_bits": (2, 1), @@ -65,13 +68,88 @@ def test_amax_is_rho_times_anchor(): def test_amax_never_below_calibrated_max(): - """A range wider than FP8 can hold warns and falls back to the no-clipping floor.""" + """A range too wide for rho warns and degenerates to plain max -- never below it.""" x = torch.zeros(64, 64) x[0, :16] = 100.0 # one large block, the rest tiny x[1:, :] = 1e-3 - with pytest.warns(UserWarning, match="exceeds the FP8 scale range"): + with pytest.warns(UserWarning, match="leaves no headroom"): amax = _calibrate(NVFP4ActHeadroomCalibrator(rho=1.0, anchor_percentile=1.0), x) - assert amax >= 100.0 * 0.9 + assert amax == pytest.approx(100.0) + + +@pytest.mark.parametrize( + "case", + [ + "gaussian", + "heavy_tail", + "sparse", + "outlier", + "wide_dynamic_range", + "constant", + ], +) +def test_amax_is_never_below_observed_max(case): + """The no-clipping invariant holds on every distribution, guardrail path included.""" + g = torch.Generator().manual_seed(7) + if case == "gaussian": + x = torch.randn(512, 256, generator=g) + elif case == "heavy_tail": + x = torch.randn(512, 256, generator=g) * torch.exp(torch.randn(512, 256, generator=g) * 2) + elif case == "sparse": + x = torch.randn(512, 256, generator=g) * (torch.rand(512, 256, generator=g) > 0.9) + elif case == "outlier": + x = torch.randn(512, 256, generator=g) + x[0, 0] = 5000.0 + elif case == "wide_dynamic_range": + # Per-token magnitudes spanning ~6 decades: the regime that trips the guardrail. + x = torch.randn(2048, 1024, generator=g) * torch.logspace(-4, 2, 2048).unsqueeze(1) + else: + x = torch.full((256, 256), 0.5) + + cal = NVFP4ActHeadroomCalibrator(anchor_percentile=1.0) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + amax = _calibrate(cal, x) + observed_max = float(reduce_block_amax(x, block_sizes={-1: 16}).max()) + assert amax >= observed_max, f"{case}: amax {amax} clips observed max {observed_max}" + + +def test_anchor_percentile_zero_rejected(): + """percentile 0 would defeat the low-tail mask, so it is rejected up front.""" + with pytest.raises(ValueError, match="anchor_percentile must be in"): + NVFP4ActHeadroomCalibrator(anchor_percentile=0.0) + + +def test_ragged_last_dim_is_handled(): + """Activations whose last dim is not a multiple of the block size must not crash.""" + g = torch.Generator().manual_seed(3) + x = torch.randn(32, 70, generator=g) # 70 % 16 != 0 + amax = _calibrate(NVFP4ActHeadroomCalibrator(anchor_percentile=1.0), x) + assert amax > 0 + assert amax >= float(x.abs().max()) + + +@pytest.mark.parametrize("bad", [float("nan"), float("inf")]) +def test_nan_inf_activations_rejected(bad): + """NaN/Inf must be reported rather than silently mis-binned.""" + x = torch.randn(32, 64) + x[0, 0] = bad + with pytest.raises(AssertionError, match=r"nan|inf"): + NVFP4ActHeadroomCalibrator().collect(x) + + +def test_calibrator_is_restored_after_calibration(): + """The swapped calibrator must not leak into later calibrations of the same model.""" + from modelopt.torch.quantization.model_calib import max_calibrate + + data = torch.randn(16, 64) + model = mtq.quantize(_Net(), ACT_ONLY_CFG, lambda m: m(data)) + assert not isinstance(model.fc1.input_quantizer._calibrator, NVFP4ActHeadroomCalibrator) + + headroom_amax = float(model.fc1.input_quantizer.amax) + # A subsequent plain-max calibration must produce a plain-max scale, not a headroom one. + max_calibrate(model, lambda m: m(data)) + assert float(model.fc1.input_quantizer.amax) < headroom_amax def test_lower_anchor_percentile_gives_smaller_amax(): @@ -123,8 +201,8 @@ def test_only_nvfp4_input_quantizers_are_selected(): def test_swap_installs_calibrator_with_config_values(): model = mtq.quantize(_Net(), ACT_ONLY_CFG, lambda m: m(torch.randn(8, 64))) - n = _swap_in_nvfp4_act_headroom_calibrators(model, anchor_percentile=5.0, rho=1024.0) - assert n == 2 # fc1 and fc2 input quantizers + swapped = _swap_in_nvfp4_act_headroom_calibrators(model, anchor_percentile=5.0, rho=1024.0) + assert len(swapped) == 2 # fc1 and fc2 input quantizers cal = model.fc1.input_quantizer._calibrator assert isinstance(cal, NVFP4ActHeadroomCalibrator) assert cal._anchor_percentile == 5.0 From f09f019b3b7a95647c00e06d27591e37da779ce3 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Wed, 29 Jul 2026 18:22:15 +0000 Subject: [PATCH 3/9] Expose shared_states and sync_expert_weight_amax on nvfp4_act_headroom This algorithm delegates all weight calibration to max_calibrate, but did not accept the shared-state knobs its siblings expose, so switching a recipe from `max` to `nvfp4_act_headroom` silently lost the ability to override the weight global_amax grouping patterns or to share one weight amax across local experts in a SequentialMLP MoE layer. Defaults are unchanged: resolve_patterns(None) already returns the default patterns, so behavior only changes when a recipe sets these fields. Also document the data-parallel sync semantics: max_calibrate combines per-rank scales with a MAX all-reduce, so the result is the largest per-rank headroom scale rather than the scale implied by pooling every rank's distribution. Pre-commit hooks were run manually on these files (all pass); the commit hook is skipped only because its autostash cannot write to a read-only cache in this environment. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/config.py | 11 ++++++- modelopt/torch/quantization/model_calib.py | 18 ++++++++++- .../quantization/test_nvfp4_act_headroom.py | 30 +++++++++++++++++++ 3 files changed, 57 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 9e9c2ad70a8..738d05275d8 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -986,7 +986,7 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): ) -class NVFP4ActHeadroomCalibConfig(QuantizeAlgorithmConfig): +class NVFP4ActHeadroomCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): """Config for the ``nvfp4_act_headroom`` calibration algorithm. Calibrates the per-tensor global scale of NVFP4 *activation* (input) quantizers so that @@ -1031,6 +1031,15 @@ class NVFP4ActHeadroomCalibConfig(QuantizeAlgorithmConfig): description="If True, the amax will be synced across the distributed processes.", ) + sync_expert_weight_amax: bool = ModeloptField( + default=False, + title="Share one weight amax across local experts in a SequentialMLP MoE layer.", + description=( + "Forwarded to max calibration, which handles the weight quantizers for this " + "algorithm. See :class:`MaxCalibConfig` for details." + ), + ) + 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 9a257067c78..d0c9288b332 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -541,6 +541,8 @@ def nvfp4_act_headroom_calibrate( anchor_percentile: float = 1.0, rho: float = 16384.0, distributed_sync: bool = True, + shared_states: Mapping[str, Mapping[str, Sequence[str]]] | None = None, + sync_expert_weight_amax: bool = False, ): """Calibrate NVFP4 activation global scales with headroom; everything else uses max. @@ -557,6 +559,14 @@ def nvfp4_act_headroom_calibrate( rho: headroom factor; ``amax = rho * anchor``. Must be in ``(0, 28672)``. distributed_sync: whether to sync amax across distributed processes (see :func:`max_calibrate`). + shared_states: shared quantization-state grouping, forwarded to :func:`max_calibrate`. + sync_expert_weight_amax: share one weight amax across local experts in a SequentialMLP + MoE layer, forwarded to :func:`max_calibrate`. + + .. note:: + Under data parallelism the per-rank scales are combined by :func:`max_calibrate` with a + ``MAX`` all-reduce, so the result is the largest per-rank headroom scale rather than the + scale implied by pooling every rank's per-block distribution. """ swapped = _swap_in_nvfp4_act_headroom_calibrators( model, anchor_percentile=anchor_percentile, rho=rho @@ -576,7 +586,13 @@ def nvfp4_act_headroom_calibrate( # The calibrators are restored afterwards so this algorithm does not leak into a later # calibration of the same model, and so a repeat run starts from a fresh histogram. try: - max_calibrate(model, forward_loop, distributed_sync=distributed_sync) + max_calibrate( + model, + forward_loop, + distributed_sync=distributed_sync, + sync_expert_weight_amax=sync_expert_weight_amax, + shared_states=shared_states, + ) finally: for quantizer, original_calibrator in swapped: quantizer._calibrator = original_calibrator diff --git a/tests/unit/torch/quantization/test_nvfp4_act_headroom.py b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py index 293ba605a1c..bb530f0881d 100644 --- a/tests/unit/torch/quantization/test_nvfp4_act_headroom.py +++ b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py @@ -152,6 +152,36 @@ def test_calibrator_is_restored_after_calibration(): assert float(model.fc1.input_quantizer.amax) < headroom_amax +def test_shared_state_knobs_are_forwarded_to_max_calibrate(): + """shared_states / sync_expert_weight_amax reach max_calibrate, as they do for `max`.""" + import modelopt.torch.quantization.model_calib as mc + + seen = {} + real = mc.max_calibrate + + def spy(model, forward_loop=None, distributed_sync=True, **kwargs): + seen.update(kwargs) + return real(model, forward_loop, distributed_sync) + + cfg = { + **ACT_ONLY_CFG, + "algorithm": { + "method": "nvfp4_act_headroom", + "anchor_percentile": 1, + "shared_states": {"weight_global_amax": {"patterns": [r".*fc1"]}}, + "sync_expert_weight_amax": True, + }, + } + mc.max_calibrate = spy + try: + mtq.quantize(_Net(), cfg, lambda m: m(torch.randn(8, 64))) + finally: + mc.max_calibrate = real + + assert seen["sync_expert_weight_amax"] is True + assert seen["shared_states"] == {"weight_global_amax": {"patterns": [r".*fc1"]}} + + def test_lower_anchor_percentile_gives_smaller_amax(): """anchor_percentile is tunable and monotone: a lower percentile anchors lower.""" torch.manual_seed(0) From e935774d6dd7a5f939125663ef09a0066ed9a0c0 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Wed, 29 Jul 2026 18:45:10 +0000 Subject: [PATCH 4/9] Note that the percentile bin representative is deliberately fixed Changing how a percentile is reconstructed from the histogram (bin center vs. a bin edge vs. cdf interpolation within the bin) shifts every calibrated scale, so record that the current choice is intentional rather than an oversight. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/calib/nvfp4_act_headroom.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/calib/nvfp4_act_headroom.py b/modelopt/torch/quantization/calib/nvfp4_act_headroom.py index 2d46bc90ac7..267537420b9 100644 --- a/modelopt/torch/quantization/calib/nvfp4_act_headroom.py +++ b/modelopt/torch/quantization/calib/nvfp4_act_headroom.py @@ -154,7 +154,9 @@ def _percentile(self, percentile: float, floor_value: float | None = None) -> fl target = percentile / 100.0 * total cdf = torch.cumsum(counts, dim=0) bin_idx = int(torch.searchsorted(cdf, target).clamp(0, self._num_bins - 1).item()) - # Bin center back to a value in linear space. + # Reconstruct at the bin center. Which representative is used (center vs. a bin edge + # vs. interpolating within the bin from the cdf) shifts every calibrated scale, so it + # is kept fixed to keep scales comparable across runs and existing checkpoints. log2_val = self._log2_min + (bin_idx + 0.5) / self._num_bins * ( self._log2_max - self._log2_min ) From f4ccecce0ceddc41b878ef936f218ee37243183a Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 30 Jul 2026 06:28:08 +0000 Subject: [PATCH 5/9] Restore the configurable calibrated-range top instead of forcing the observed max An earlier commit made the observed per-block max the hard floor for the global scale, on the grounds that a high percentile does not guarantee every calibration block is representable. Measuring the resulting quantization error shows that was the wrong trade. Flooring at the literal max means a single freak block drags the global scale up until every other block's FP8 block scale falls below subnormal: on a tensor with one block seven orders of magnitude above the rest, it flushed 99.998% of elements to zero, versus 6.7% when the rare blocks are clipped instead. On a benign wide-range tensor the two differ by 0.4% relative MSE, so the strict no-clipping floor buys almost nothing where it is safe and is destructive where it is not. Restore the top of the calibrated range as a configurable percentile (``upper_percentile``, default 99.99, ``100`` selects the literal observed max), and state the trade-off in the docstrings rather than asserting a no-clipping guarantee the default does not provide. That overclaim, not the percentile, was the actual defect. This also makes the guardrail ratio self-consistent again: the reported range is measured against the bound that is actually operative. Also address review feedback: move test imports to module level, and use the canonical 2026 copyright year on the two new files. Pre-commit hooks were run manually on these files (all pass); the commit hook is skipped only because its autostash cannot write to a read-only cache in this environment. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 2 +- .../quantization/calib/nvfp4_act_headroom.py | 71 +++++++++++-------- modelopt/torch/quantization/config.py | 17 +++++ modelopt/torch/quantization/model_calib.py | 11 ++- .../quantization/test_nvfp4_act_headroom.py | 56 +++++++++++---- 5 files changed, 109 insertions(+), 48 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3ef5c89276e..02590724c65 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,7 +8,7 @@ Changelog *Quantization* -- Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Plain ``max`` sets a tensor's global scale from the largest per-block amax seen during calibration, leaving no room above it, so any activation larger than the calibration max saturates. ``nvfp4_act_headroom`` instead anchors the global scale to a low percentile of the per-block amax distribution, ``amax = max(rho * anchor, floor)``, placing the calibrated blocks in the lower part of the FP8 block-scale range and leaving the rest as headroom; the result is floored at the largest per-block amax observed during calibration so the scale never sits below the calibrated data, and the calibrator warns and degenerates to plain ``max`` when the per-block range is too wide for ``rho`` to clear. Tunable via ``anchor_percentile`` (default 1) and ``rho`` (default 16384). Applies only to NVFP4 dynamic-block input quantizers; weight quantizers and all other quantizers keep plain ``max``. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` (dynamic NVFP4 W4A4 plus FP8 KV-cache cast) with only the calibration algorithm swapped, so it exports a standard NVFP4 checkpoint. +- Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Plain ``max`` sets a tensor's global scale from the largest per-block amax seen during calibration, leaving no room above it, so any activation larger than the calibration max saturates. ``nvfp4_act_headroom`` instead anchors the global scale to a low percentile of the per-block amax distribution, ``amax = max(rho * anchor, upper)``, placing the calibrated blocks in the lower part of the FP8 block-scale range and leaving the rest as headroom. ``upper`` is the top of the range the scale commits to representing and defaults to the 99.99th percentile rather than the literal maximum: chasing a lone freak block would drag the global scale up until every other block's FP8 block scale falls below subnormal and flushes to zero, so the rarest blocks are clipped instead. Set ``upper_percentile=100`` to use the literal observed max, which guarantees no calibration data is clipped. The calibrator warns when the per-block range is too wide for ``rho`` to clear any headroom. Tunable via ``anchor_percentile`` (default 1), ``upper_percentile`` (default 99.99) and ``rho`` (default 16384). Applies only to NVFP4 dynamic-block input quantizers; weight quantizers and all other quantizers keep plain ``max``. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` (dynamic NVFP4 W4A4 plus FP8 KV-cache cast) with only the calibration algorithm swapped, so it exports a standard NVFP4 checkpoint. **Backward Breaking Changes** diff --git a/modelopt/torch/quantization/calib/nvfp4_act_headroom.py b/modelopt/torch/quantization/calib/nvfp4_act_headroom.py index 267537420b9..2c94f02a819 100644 --- a/modelopt/torch/quantization/calib/nvfp4_act_headroom.py +++ b/modelopt/torch/quantization/calib/nvfp4_act_headroom.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -28,10 +28,7 @@ # outside this ratio cannot all be represented as normal FP8 values by one global scale. _FP8_NORMAL_DYNAMIC_RANGE = 28672.0 -# High percentile of the per-block amax distribution used as the no-clipping floor. -_FLOOR_PERCENTILE = 99.99 - -# Per-block amaxes below ``floor / _ANCHOR_FLOOR_RATIO`` are ignored when locating the +# Per-block amaxes below ``upper / _ANCHOR_FLOOR_RATIO`` are ignored when locating the # anchor, so a tail of near-zero blocks cannot drag the global scale down. _ANCHOR_FLOOR_RATIO = 1e6 @@ -47,14 +44,19 @@ class NVFP4ActHeadroomCalibrator(_Calibrator): This calibrator instead anchors the global scale to a low percentile of the per-block amax distribution:: - amax = max(rho * anchor, observed_max) + amax = max(rho * anchor, upper) + + where ``anchor`` is the per-block amax at ``anchor_percentile`` and ``upper`` is the + per-block amax at ``upper_percentile``. Multiplying the low anchor by ``rho`` places the + calibrated blocks in the lower part of the FP8 scale range and leaves the rest as upward + headroom; ``upper`` is the top of the range the scale commits to representing. - where ``anchor`` is the per-block amax at ``anchor_percentile``. Multiplying the low anchor - by ``rho`` places the calibrated blocks in the lower part of the FP8 scale range and leaves - the rest as upward headroom. The result is floored at the largest per-block amax observed - during calibration, so the scale never sits below the calibrated data; when that floor wins - (a per-block range too wide for ``rho`` to clear) the calibrator warns and the result - degenerates to plain ``max``. + ``upper_percentile`` defaults to 99.99 rather than the literal maximum on purpose. A single + freak block far above the rest would otherwise drag the global scale up so far that every + other block's FP8 block scale falls below subnormal and flushes to zero -- losing the whole + tensor to protect one value. The default clips those rare blocks instead. Set + ``upper_percentile=100`` to use the literal observed max, which guarantees no calibration + data is clipped at the cost of that exposure. Args: num_bits: quantizer ``num_bits`` (``(2, 1)`` for NVFP4); kept for interface parity. @@ -62,6 +64,8 @@ class NVFP4ActHeadroomCalibrator(_Calibrator): unsigned: unused; kept for interface parity. block_size: NVFP4 block width along the last dim. anchor_percentile: percentile of the per-block amaxes used as the anchor; must be > 0. + upper_percentile: percentile of the per-block amaxes used as the top of the calibrated + range. ``100`` uses the literal observed max (no calibration data is clipped). rho: headroom factor applied to the anchor. Must be in ``(0, 28672)``. num_bins: number of log2 histogram bins. log2_min: low end of the log2 range covered by the histogram. @@ -76,6 +80,7 @@ def __init__( *, block_size=16, anchor_percentile=1.0, + upper_percentile=99.99, rho=16384.0, num_bins=512, log2_min=-40.0, @@ -91,7 +96,10 @@ def __init__( if not (0.0 < anchor_percentile <= 100.0): raise ValueError(f"anchor_percentile must be in (0, 100]; got {anchor_percentile}.") self._block_size = block_size + if not (0.0 < upper_percentile <= 100.0): + raise ValueError(f"upper_percentile must be in (0, 100]; got {upper_percentile}.") self._anchor_percentile = float(anchor_percentile) + self._upper_percentile = float(upper_percentile) self._rho = float(rho) self._num_bins = int(num_bins) self._log2_min = float(log2_min) @@ -168,38 +176,42 @@ def compute_amax(self) -> torch.Tensor | None: if self._hist is None or self._running_max is None: return None - running_max = float(self._running_max) - floor = self._percentile(_FLOOR_PERCENTILE) + upper = ( + float(self._running_max) + if self._upper_percentile >= 100.0 + else self._percentile(self._upper_percentile) + ) anchor = ( self._percentile( - self._anchor_percentile, floor_value=floor / _ANCHOR_FLOOR_RATIO if floor else None + self._anchor_percentile, floor_value=upper / _ANCHOR_FLOOR_RATIO if upper else None ) - if floor + if upper else None ) - if not floor or floor <= 0 or not anchor or anchor <= 0: + if not upper or upper <= 0 or not anchor or anchor <= 0: # Percentiles collapsed despite a non-empty histogram: fall back to plain max. return self._running_max.clone() - # The observed max is the no-clipping bound: the returned scale never sits below it, so - # calibration data is never clipped. Headroom is whatever the anchor term adds on top. + # ``upper`` is the top of the range the scale commits to representing; the anchor term + # only ever raises the scale above it. Blocks above ``upper`` are clipped -- see the + # class docstring for why that is preferred to chasing a lone outlier. headroom_amax = self._rho * anchor - if headroom_amax <= running_max: - exceeds_window = floor / anchor > _FP8_NORMAL_DYNAMIC_RANGE + if headroom_amax <= upper: + span = upper / anchor warnings.warn( - f"[nvfp4_act_headroom] per-block amax range {floor / anchor:.1f} leaves no " - f"headroom at rho={self._rho:g}; falling back to the observed max. " + f"[nvfp4_act_headroom] per-block amax range {span:.1f} leaves no headroom at " + f"rho={self._rho:g}; the scale falls back to the top of the calibrated range. " + ( f"The range also exceeds the FP8 scale range " - f"({_FP8_NORMAL_DYNAMIC_RANGE:.0f}), so no global scale can hold both ends. " - "Reduce the range with outlier mitigation (SmoothQuant / per-channel / " - "higher precision)." - if exceeds_window + f"({_FP8_NORMAL_DYNAMIC_RANGE:.0f}), so no permitted rho can clear it and no " + "single global scale holds both ends. Reduce the range with outlier " + "mitigation (SmoothQuant / per-channel / higher precision)." + if span > _FP8_NORMAL_DYNAMIC_RANGE else "Lower anchor_percentile or raise rho for more headroom." ), stacklevel=2, ) - amax = running_max + amax = upper else: amax = headroom_amax @@ -213,5 +225,6 @@ def reset(self) -> None: def __repr__(self): return ( f"NVFP4ActHeadroomCalibrator(anchor_percentile={self._anchor_percentile}, " - f"rho={self._rho}, block_size={self._block_size})" + f"upper_percentile={self._upper_percentile}, rho={self._rho}, " + f"block_size={self._block_size})" ) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 738d05275d8..5178eec095e 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -994,6 +994,10 @@ class NVFP4ActHeadroomCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): above it for activations larger than any seen during calibration. Weight quantizers and all non-NVFP4 quantizers are calibrated with plain ``max``. + The top of the calibrated range is ``upper_percentile`` (default 99.99) rather than the + literal maximum, so rare blocks far above the rest are clipped instead of dragging the + global scale up until every other block flushes to zero. + See :class:`NVFP4ActHeadroomCalibrator ` for the formula. """ @@ -1014,6 +1018,19 @@ class NVFP4ActHeadroomCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): ), ) + upper_percentile: float = ModeloptField( + default=99.99, + gt=0.0, + le=100.0, + title="Percentile of the per-block activation amaxes used as the top of the range.", + description=( + "The global scale is floored at this percentile, so per-block amaxes above it are " + "clipped. The default excludes the rarest blocks on purpose: chasing a lone outlier " + "pushes every other block's FP8 block scale below subnormal. Set to 100 to use the " + "literal observed max, which guarantees no calibration data is clipped." + ), + ) + rho: float = ModeloptField( default=16384.0, gt=0.0, diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index d0c9288b332..fcecd645605 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -511,7 +511,7 @@ def _is_nvfp4_dynamic_input_quantizer(name: str, module: nn.Module) -> bool: def _swap_in_nvfp4_act_headroom_calibrators( - model: nn.Module, *, anchor_percentile: float, rho: float + model: nn.Module, *, anchor_percentile: float, upper_percentile: float, rho: float ) -> list[tuple[nn.Module, Any]]: """Swap in an :class:`NVFP4ActHeadroomCalibrator` for each qualifying NVFP4 input quantizer. @@ -528,6 +528,7 @@ def _swap_in_nvfp4_act_headroom_calibrators( module._unsigned, block_size=module.block_sizes.get(-1, 16), anchor_percentile=anchor_percentile, + upper_percentile=upper_percentile, rho=rho, ) return swapped @@ -539,6 +540,7 @@ def nvfp4_act_headroom_calibrate( forward_loop: ForwardLoop | None = None, *, anchor_percentile: float = 1.0, + upper_percentile: float = 99.99, rho: float = 16384.0, distributed_sync: bool = True, shared_states: Mapping[str, Mapping[str, Sequence[str]]] | None = None, @@ -556,6 +558,8 @@ def nvfp4_act_headroom_calibrate( model: model to be calibrated. forward_loop: callable that runs calibration data through the model. anchor_percentile: percentile of the per-block amaxes used as the anchor. + upper_percentile: percentile of the per-block amaxes used as the top of the calibrated + range; ``100`` uses the literal observed max. rho: headroom factor; ``amax = rho * anchor``. Must be in ``(0, 28672)``. distributed_sync: whether to sync amax across distributed processes (see :func:`max_calibrate`). @@ -569,7 +573,7 @@ def nvfp4_act_headroom_calibrate( scale implied by pooling every rank's per-block distribution. """ swapped = _swap_in_nvfp4_act_headroom_calibrators( - model, anchor_percentile=anchor_percentile, rho=rho + model, anchor_percentile=anchor_percentile, upper_percentile=upper_percentile, rho=rho ) if not swapped: warn_rank_0( @@ -579,7 +583,8 @@ def nvfp4_act_headroom_calibrate( ) print_rank_0( f"nvfp4_act_headroom: calibrating {len(swapped)} NVFP4 activation quantizer(s) " - f"(anchor_percentile={anchor_percentile}, rho={rho})." + f"(anchor_percentile={anchor_percentile}, upper_percentile={upper_percentile}, " + f"rho={rho})." ) # max_calibrate runs the forward loop once: the swapped-in calibrators accumulate their # per-block histograms in the same pass that collects max stats for every other quantizer. diff --git a/tests/unit/torch/quantization/test_nvfp4_act_headroom.py b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py index bb530f0881d..e655477a08c 100644 --- a/tests/unit/torch/quantization/test_nvfp4_act_headroom.py +++ b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -21,10 +21,12 @@ import torch import modelopt.torch.quantization as mtq +from modelopt.torch.quantization import model_calib from modelopt.torch.quantization.calib import NVFP4ActHeadroomCalibrator from modelopt.torch.quantization.model_calib import ( _is_nvfp4_dynamic_input_quantizer, _swap_in_nvfp4_act_headroom_calibrators, + max_calibrate, ) from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.quantization.utils import reduce_block_amax @@ -67,14 +69,14 @@ def test_amax_is_rho_times_anchor(): assert amax == pytest.approx(1024.0 * 0.5, rel=0.05) -def test_amax_never_below_calibrated_max(): - """A range too wide for rho warns and degenerates to plain max -- never below it.""" +def test_range_too_wide_for_rho_falls_back_to_the_calibrated_top(): + """A range too wide for rho warns and falls back to the top of the calibrated range.""" x = torch.zeros(64, 64) x[0, :16] = 100.0 # one large block, the rest tiny x[1:, :] = 1e-3 with pytest.warns(UserWarning, match="leaves no headroom"): amax = _calibrate(NVFP4ActHeadroomCalibrator(rho=1.0, anchor_percentile=1.0), x) - assert amax == pytest.approx(100.0) + assert amax == pytest.approx(100.0, rel=0.06) # within one histogram bin of the top block @pytest.mark.parametrize( @@ -88,8 +90,8 @@ def test_amax_never_below_calibrated_max(): "constant", ], ) -def test_amax_is_never_below_observed_max(case): - """The no-clipping invariant holds on every distribution, guardrail path included.""" +def test_upper_percentile_100_never_clips(case): + """upper_percentile=100 is the documented no-clipping setting; verify it holds.""" g = torch.Generator().manual_seed(7) if case == "gaussian": x = torch.randn(512, 256, generator=g) @@ -106,7 +108,7 @@ def test_amax_is_never_below_observed_max(case): else: x = torch.full((256, 256), 0.5) - cal = NVFP4ActHeadroomCalibrator(anchor_percentile=1.0) + cal = NVFP4ActHeadroomCalibrator(anchor_percentile=1.0, upper_percentile=100.0) with warnings.catch_warnings(): warnings.simplefilter("ignore") amax = _calibrate(cal, x) @@ -114,6 +116,31 @@ def test_amax_is_never_below_observed_max(case): assert amax >= observed_max, f"{case}: amax {amax} clips observed max {observed_max}" +def test_default_clips_rare_outliers_to_protect_the_bulk(): + """A lone freak block is clipped by default; chasing it would flush the tensor to zero.""" + g = torch.Generator().manual_seed(11) + x = torch.randn(4096, 256, generator=g).abs() + 1e-3 + x[0, :16] = 3e7 # one block ~7 orders of magnitude above the rest + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + default = _calibrate(NVFP4ActHeadroomCalibrator(anchor_percentile=1.0), x) + literal = _calibrate( + NVFP4ActHeadroomCalibrator(anchor_percentile=1.0, upper_percentile=100.0), x + ) + + # The default ignores the outlier, keeping the scale near the bulk of the distribution. + assert default < 3e7 + # Honouring it instead drags the scale orders of magnitude higher, which is exactly the + # exposure upper_percentile=100 opts into. + assert literal > 100 * default + + +def test_upper_percentile_out_of_range_rejected(): + with pytest.raises(ValueError, match="upper_percentile must be in"): + NVFP4ActHeadroomCalibrator(upper_percentile=0.0) + + def test_anchor_percentile_zero_rejected(): """percentile 0 would defeat the low-tail mask, so it is rejected up front.""" with pytest.raises(ValueError, match="anchor_percentile must be in"): @@ -140,8 +167,6 @@ def test_nan_inf_activations_rejected(bad): def test_calibrator_is_restored_after_calibration(): """The swapped calibrator must not leak into later calibrations of the same model.""" - from modelopt.torch.quantization.model_calib import max_calibrate - data = torch.randn(16, 64) model = mtq.quantize(_Net(), ACT_ONLY_CFG, lambda m: m(data)) assert not isinstance(model.fc1.input_quantizer._calibrator, NVFP4ActHeadroomCalibrator) @@ -154,10 +179,8 @@ def test_calibrator_is_restored_after_calibration(): def test_shared_state_knobs_are_forwarded_to_max_calibrate(): """shared_states / sync_expert_weight_amax reach max_calibrate, as they do for `max`.""" - import modelopt.torch.quantization.model_calib as mc - seen = {} - real = mc.max_calibrate + real = model_calib.max_calibrate def spy(model, forward_loop=None, distributed_sync=True, **kwargs): seen.update(kwargs) @@ -172,11 +195,11 @@ def spy(model, forward_loop=None, distributed_sync=True, **kwargs): "sync_expert_weight_amax": True, }, } - mc.max_calibrate = spy + model_calib.max_calibrate = spy try: mtq.quantize(_Net(), cfg, lambda m: m(torch.randn(8, 64))) finally: - mc.max_calibrate = real + model_calib.max_calibrate = real assert seen["sync_expert_weight_amax"] is True assert seen["shared_states"] == {"weight_global_amax": {"patterns": [r".*fc1"]}} @@ -231,11 +254,14 @@ def test_only_nvfp4_input_quantizers_are_selected(): def test_swap_installs_calibrator_with_config_values(): model = mtq.quantize(_Net(), ACT_ONLY_CFG, lambda m: m(torch.randn(8, 64))) - swapped = _swap_in_nvfp4_act_headroom_calibrators(model, anchor_percentile=5.0, rho=1024.0) + swapped = _swap_in_nvfp4_act_headroom_calibrators( + model, anchor_percentile=5.0, upper_percentile=99.99, rho=1024.0 + ) assert len(swapped) == 2 # fc1 and fc2 input quantizers cal = model.fc1.input_quantizer._calibrator assert isinstance(cal, NVFP4ActHeadroomCalibrator) assert cal._anchor_percentile == 5.0 + assert cal._upper_percentile == 99.99 assert cal._rho == 1024.0 From 6a302f313557e0d2d2c9922f0c11c493e477c126 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 30 Jul 2026 17:24:18 +0000 Subject: [PATCH 6/9] Show all three nvfp4_act_headroom knobs in the shipped recipe The recipe only set anchor_percentile, leaving upper_percentile and rho implicit. Spell all three out at their defaults so the recipe documents the tunables and the formula they feed, and note that upper_percentile=100 selects the literal observed max. Values match the defaults, so calibration behavior is unchanged. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Chenjie Luo --- .../general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml b/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml index a436e57759a..bf81b47bbc9 100644 --- a/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml @@ -33,7 +33,13 @@ metadata: quantize: algorithm: method: nvfp4_act_headroom + # amax = max(rho * anchor, upper), where anchor and upper are the per-block activation + # amaxes at anchor_percentile and upper_percentile. All three are shown at their defaults. + # Set upper_percentile to 100 to use the literal observed max, which never clips + # calibration data but lets a single outlier block dominate the global scale. anchor_percentile: 1 + upper_percentile: 99.99 + rho: 16384 quant_cfg: - $import: base_disable_all - $import: w4a4_nvfp4_nvfp4 From 136e0361be88a0f396d40b73bd033304d3a25196 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 30 Jul 2026 17:48:03 +0000 Subject: [PATCH 7/9] Calibrate SequentialQuantizer activation leaves; fix the recipe license year A quant_cfg whose input_quantizer cfg is a list builds a SequentialQuantizer, whose leaves are submodules named ``...input_quantizer.``. The selector matched on the leaf name ending in ``input_quantizer``, so neither the container (not a TensorQuantizer) nor its leaves (name has a numeric suffix) qualified, and those NVFP4 activations silently fell back to plain max. Match on the container name and walk the leaves via the existing ``_iter_leaf_quantizers`` helper, with a regression test asserting all four leaves of a two-layer model are calibrated and receive the headroom scale rather than a max scale. The new recipe also still carried a 2024 copyright header, copied from the unit file it was derived from; it now uses the canonical 2026 header like the other two new files. Pre-commit hooks were run manually on these files (all pass); the commit hook is skipped only because its autostash cannot write to a read-only cache in this environment. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/model_calib.py | 44 +++++++++++++------ .../ptq/nvfp4_act_headroom-kv_fp8_cast.yaml | 2 +- .../quantization/test_nvfp4_act_headroom.py | 32 +++++++++++++- 3 files changed, 62 insertions(+), 16 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index fcecd645605..9cda40ed61b 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -491,13 +491,13 @@ def sync_quantizer_amax_across_tp( _finalize_with_shared_state(model, weight_patterns) -def _is_nvfp4_dynamic_input_quantizer(name: str, module: nn.Module) -> bool: - """True for an enabled NVFP4 (E2M1 + E4M3 scale) dynamic-block *input* quantizer. +def _is_nvfp4_dynamic_activation_quantizer(module: nn.Module) -> bool: + """True for an enabled NVFP4 (E2M1 + E4M3 scale) dynamic-block quantizer leaf. These carry their per-tensor global scale in ``_amax`` and are calibrated like plain max (the quantizer itself is not ``_dynamic``; only its per-block scales are). """ - if not isinstance(module, TensorQuantizer) or not name.endswith("input_quantizer"): + if not isinstance(module, TensorQuantizer): return False if getattr(module, "_disabled", False) or getattr(module, "_dynamic", False): return False @@ -510,6 +510,18 @@ def _is_nvfp4_dynamic_input_quantizer(name: str, module: nn.Module) -> bool: ) +def _is_nvfp4_dynamic_input_quantizer(name: str, module: nn.Module) -> bool: + """True for an *input* quantizer holding at least one qualifying NVFP4 leaf. + + ``input_quantizer`` may be a :class:`SequentialQuantizer`, whose leaves are submodules + named ``...input_quantizer.``. Matching on the container name and then walking the + leaves keeps those wrapped quantizers from silently falling back to plain max. + """ + if not name.endswith("input_quantizer"): + return False + return any(_is_nvfp4_dynamic_activation_quantizer(q) for q in _iter_leaf_quantizers(module)) + + def _swap_in_nvfp4_act_headroom_calibrators( model: nn.Module, *, anchor_percentile: float, upper_percentile: float, rho: float ) -> list[tuple[nn.Module, Any]]: @@ -519,18 +531,22 @@ def _swap_in_nvfp4_act_headroom_calibrators( """ swapped: list[tuple[nn.Module, Any]] = [] for name, module in model.named_modules(): - if not _is_nvfp4_dynamic_input_quantizer(name, module): + if not name.endswith("input_quantizer"): continue - swapped.append((module, module._calibrator)) - module._calibrator = NVFP4ActHeadroomCalibrator( - module.num_bits, - None, - module._unsigned, - block_size=module.block_sizes.get(-1, 16), - anchor_percentile=anchor_percentile, - upper_percentile=upper_percentile, - rho=rho, - ) + # Walk SequentialQuantizer leaves so wrapped NVFP4 inputs are calibrated too. + for leaf in _iter_leaf_quantizers(module): + if not _is_nvfp4_dynamic_activation_quantizer(leaf): + continue + swapped.append((leaf, leaf._calibrator)) + leaf._calibrator = NVFP4ActHeadroomCalibrator( + leaf.num_bits, + None, + leaf._unsigned, + block_size=leaf.block_sizes.get(-1, 16), + anchor_percentile=anchor_percentile, + upper_percentile=upper_percentile, + rho=rho, + ) return swapped diff --git a/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml b/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml index bf81b47bbc9..21a71f4a10b 100644 --- a/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); diff --git a/tests/unit/torch/quantization/test_nvfp4_act_headroom.py b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py index e655477a08c..7fdb24510a5 100644 --- a/tests/unit/torch/quantization/test_nvfp4_act_headroom.py +++ b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py @@ -28,7 +28,7 @@ _swap_in_nvfp4_act_headroom_calibrators, max_calibrate, ) -from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.utils import reduce_block_amax NVFP4_CFG = { @@ -252,6 +252,36 @@ def test_only_nvfp4_input_quantizers_are_selected(): assert not _is_nvfp4_dynamic_input_quantizer("layer.input_quantizer", fp8_in) +def test_sequential_quantizer_leaves_are_calibrated(): + """A SequentialQuantizer input must not silently fall back to plain max. + + Its leaves are submodules named ``...input_quantizer.``, so a plain + ``endswith("input_quantizer")`` check on the leaf name would skip them. + """ + cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*input_quantizer", "cfg": [NVFP4_CFG, NVFP4_CFG]}, + ], + "algorithm": {"method": "nvfp4_act_headroom", "anchor_percentile": 1}, + } + torch.manual_seed(0) + data = torch.randn(16, 64) + model = mtq.quantize(_Net(), cfg, lambda m: m(data)) + assert isinstance(model.fc1.input_quantizer, SequentialQuantizer) + + swapped = _swap_in_nvfp4_act_headroom_calibrators( + model, anchor_percentile=1.0, upper_percentile=99.99, rho=16384.0 + ) + # Both leaves of both layers, not zero and not just the containers. + assert len(swapped) == 4 + assert all(isinstance(q, TensorQuantizer) for q, _ in swapped) + + # And the calibrated scale really is the headroom one, not plain max. + for leaf in model.fc1.input_quantizer: + assert float(leaf.amax) > float(data.abs().max()) + + def test_swap_installs_calibrator_with_config_values(): model = mtq.quantize(_Net(), ACT_ONLY_CFG, lambda m: m(torch.randn(8, 64))) swapped = _swap_in_nvfp4_act_headroom_calibrators( From 4620291db35f608468350be469ab8c3063072889 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 30 Jul 2026 17:51:52 +0000 Subject: [PATCH 8/9] Reject SequentialQuantizer activations instead of calibrating their leaves A SequentialQuantizer wraps several quantizers over one activation, while this algorithm derives a single per-tensor global scale, so the composition has no defined behavior here. Calibrating each leaf independently, as the previous commit did, papers over that rather than settling it. Reject it explicitly instead: if an input quantizer is a SequentialQuantizer wrapping an NVFP4 dynamic-block quantizer, raise NotImplementedError naming the module. Validation runs before any calibrator is swapped, so a rejected model is left untouched. SequentialQuantizers with no NVFP4 leaf are irrelevant to this algorithm and are ignored rather than rejected. Pre-commit hooks were run manually on these files (all pass); the commit hook is skipped only because its autostash cannot write to a read-only cache in this environment. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 2 +- modelopt/torch/quantization/model_calib.py | 55 +++++++++------- .../quantization/test_nvfp4_act_headroom.py | 66 ++++++++++++++----- 3 files changed, 85 insertions(+), 38 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 02590724c65..1dcd3b636fc 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,7 +8,7 @@ Changelog *Quantization* -- Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Plain ``max`` sets a tensor's global scale from the largest per-block amax seen during calibration, leaving no room above it, so any activation larger than the calibration max saturates. ``nvfp4_act_headroom`` instead anchors the global scale to a low percentile of the per-block amax distribution, ``amax = max(rho * anchor, upper)``, placing the calibrated blocks in the lower part of the FP8 block-scale range and leaving the rest as headroom. ``upper`` is the top of the range the scale commits to representing and defaults to the 99.99th percentile rather than the literal maximum: chasing a lone freak block would drag the global scale up until every other block's FP8 block scale falls below subnormal and flushes to zero, so the rarest blocks are clipped instead. Set ``upper_percentile=100`` to use the literal observed max, which guarantees no calibration data is clipped. The calibrator warns when the per-block range is too wide for ``rho`` to clear any headroom. Tunable via ``anchor_percentile`` (default 1), ``upper_percentile`` (default 99.99) and ``rho`` (default 16384). Applies only to NVFP4 dynamic-block input quantizers; weight quantizers and all other quantizers keep plain ``max``. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` (dynamic NVFP4 W4A4 plus FP8 KV-cache cast) with only the calibration algorithm swapped, so it exports a standard NVFP4 checkpoint. +- Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Plain ``max`` sets a tensor's global scale from the largest per-block amax seen during calibration, leaving no room above it, so any activation larger than the calibration max saturates. ``nvfp4_act_headroom`` instead anchors the global scale to a low percentile of the per-block amax distribution, ``amax = max(rho * anchor, upper)``, placing the calibrated blocks in the lower part of the FP8 block-scale range and leaving the rest as headroom. ``upper`` is the top of the range the scale commits to representing and defaults to the 99.99th percentile rather than the literal maximum: chasing a lone freak block would drag the global scale up until every other block's FP8 block scale falls below subnormal and flushes to zero, so the rarest blocks are clipped instead. Set ``upper_percentile=100`` to use the literal observed max, which guarantees no calibration data is clipped. The calibrator warns when the per-block range is too wide for ``rho`` to clear any headroom. Tunable via ``anchor_percentile`` (default 1), ``upper_percentile`` (default 99.99) and ``rho`` (default 16384). Applies only to NVFP4 dynamic-block input quantizers; weight quantizers and all other quantizers keep plain ``max``. ``SequentialQuantizer`` activation quantizers are not supported and raise. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` (dynamic NVFP4 W4A4 plus FP8 KV-cache cast) with only the calibration algorithm swapped, so it exports a standard NVFP4 checkpoint. **Backward Breaking Changes** diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 9cda40ed61b..5c5963e25e2 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -511,15 +511,12 @@ def _is_nvfp4_dynamic_activation_quantizer(module: nn.Module) -> bool: def _is_nvfp4_dynamic_input_quantizer(name: str, module: nn.Module) -> bool: - """True for an *input* quantizer holding at least one qualifying NVFP4 leaf. + """True for a plain NVFP4 dynamic-block *input* quantizer. - ``input_quantizer`` may be a :class:`SequentialQuantizer`, whose leaves are submodules - named ``...input_quantizer.``. Matching on the container name and then walking the - leaves keeps those wrapped quantizers from silently falling back to plain max. + :class:`SequentialQuantizer` activation quantizers are not supported by this algorithm and + are rejected in :func:`_swap_in_nvfp4_act_headroom_calibrators` rather than matched here. """ - if not name.endswith("input_quantizer"): - return False - return any(_is_nvfp4_dynamic_activation_quantizer(q) for q in _iter_leaf_quantizers(module)) + return name.endswith("input_quantizer") and _is_nvfp4_dynamic_activation_quantizer(module) def _swap_in_nvfp4_act_headroom_calibrators( @@ -529,24 +526,34 @@ def _swap_in_nvfp4_act_headroom_calibrators( Returns ``(quantizer, original_calibrator)`` pairs so the caller can restore them. """ - swapped: list[tuple[nn.Module, Any]] = [] + # Validate before mutating anything, so a rejected model is left untouched. A + # SequentialQuantizer wraps several quantizers over one activation; this algorithm derives a + # single global scale per tensor and has no defined behavior for that composition, so reject + # it explicitly instead of silently leaving those activations on plain max. for name, module in model.named_modules(): - if not name.endswith("input_quantizer"): + if not name.endswith("input_quantizer") or not isinstance(module, SequentialQuantizer): continue - # Walk SequentialQuantizer leaves so wrapped NVFP4 inputs are calibrated too. - for leaf in _iter_leaf_quantizers(module): - if not _is_nvfp4_dynamic_activation_quantizer(leaf): - continue - swapped.append((leaf, leaf._calibrator)) - leaf._calibrator = NVFP4ActHeadroomCalibrator( - leaf.num_bits, - None, - leaf._unsigned, - block_size=leaf.block_sizes.get(-1, 16), - anchor_percentile=anchor_percentile, - upper_percentile=upper_percentile, - rho=rho, + if any(_is_nvfp4_dynamic_activation_quantizer(q) for q in _iter_leaf_quantizers(module)): + raise NotImplementedError( + f"nvfp4_act_headroom does not support SequentialQuantizer activation quantizers, " + f"but {name!r} is one wrapping an NVFP4 dynamic-block quantizer. Use a single " + f"activation quantizer config for these modules, or calibrate with 'max'." ) + + swapped: list[tuple[nn.Module, Any]] = [] + for name, module in model.named_modules(): + if not _is_nvfp4_dynamic_input_quantizer(name, module): + continue + swapped.append((module, module._calibrator)) + module._calibrator = NVFP4ActHeadroomCalibrator( + module.num_bits, + None, + module._unsigned, + block_size=module.block_sizes.get(-1, 16), + anchor_percentile=anchor_percentile, + upper_percentile=upper_percentile, + rho=rho, + ) return swapped @@ -579,6 +586,10 @@ def nvfp4_act_headroom_calibrate( rho: headroom factor; ``amax = rho * anchor``. Must be in ``(0, 28672)``. distributed_sync: whether to sync amax across distributed processes (see :func:`max_calibrate`). + + Raises: + NotImplementedError: if an NVFP4 activation quantizer is a ``SequentialQuantizer``, + which this algorithm does not support. shared_states: shared quantization-state grouping, forwarded to :func:`max_calibrate`. sync_expert_weight_amax: share one weight amax across local experts in a SequentialMLP MoE layer, forwarded to :func:`max_calibrate`. diff --git a/tests/unit/torch/quantization/test_nvfp4_act_headroom.py b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py index 7fdb24510a5..c7ec9598e42 100644 --- a/tests/unit/torch/quantization/test_nvfp4_act_headroom.py +++ b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py @@ -252,11 +252,11 @@ def test_only_nvfp4_input_quantizers_are_selected(): assert not _is_nvfp4_dynamic_input_quantizer("layer.input_quantizer", fp8_in) -def test_sequential_quantizer_leaves_are_calibrated(): - """A SequentialQuantizer input must not silently fall back to plain max. +def test_sequential_quantizer_activation_is_rejected(): + """SequentialQuantizer activations are unsupported and must fail loudly, not silently. - Its leaves are submodules named ``...input_quantizer.``, so a plain - ``endswith("input_quantizer")`` check on the leaf name would skip them. + Their leaves are submodules named ``...input_quantizer.``, so a name-based match would + skip them and leave those activations on plain max without any signal. """ cfg = { "quant_cfg": [ @@ -266,20 +266,56 @@ def test_sequential_quantizer_leaves_are_calibrated(): "algorithm": {"method": "nvfp4_act_headroom", "anchor_percentile": 1}, } torch.manual_seed(0) - data = torch.randn(16, 64) - model = mtq.quantize(_Net(), cfg, lambda m: m(data)) - assert isinstance(model.fc1.input_quantizer, SequentialQuantizer) + with pytest.raises(NotImplementedError, match="does not support SequentialQuantizer"): + mtq.quantize(_Net(), cfg, lambda m: m(torch.randn(16, 64))) - swapped = _swap_in_nvfp4_act_headroom_calibrators( - model, anchor_percentile=1.0, upper_percentile=99.99, rho=16384.0 + +def test_rejection_leaves_the_model_untouched(): + """Validation happens before any calibrator is swapped, so a rejected model is unchanged.""" + torch.manual_seed(0) + model = mtq.quantize( + _Net(), + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*fc1*input_quantizer", "cfg": NVFP4_CFG}, + {"quantizer_name": "*fc2*input_quantizer", "cfg": [NVFP4_CFG, NVFP4_CFG]}, + ], + "algorithm": "max", + }, + lambda m: m(torch.randn(16, 64)), ) - # Both leaves of both layers, not zero and not just the containers. - assert len(swapped) == 4 - assert all(isinstance(q, TensorQuantizer) for q, _ in swapped) + before = model.fc1.input_quantizer._calibrator + with pytest.raises(NotImplementedError, match="does not support SequentialQuantizer"): + _swap_in_nvfp4_act_headroom_calibrators( + model, anchor_percentile=1.0, upper_percentile=99.99, rho=16384.0 + ) + assert model.fc1.input_quantizer._calibrator is before - # And the calibrated scale really is the headroom one, not plain max. - for leaf in model.fc1.input_quantizer: - assert float(leaf.amax) > float(data.abs().max()) + +def test_non_nvfp4_sequential_activation_is_ignored(): + """Only SequentialQuantizers wrapping NVFP4 leaves are rejected; others are irrelevant.""" + torch.manual_seed(0) + fp8 = {"num_bits": (4, 3)} + model = mtq.quantize( + _Net(), + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*input_quantizer", "cfg": [fp8, fp8]}, + ], + "algorithm": "max", + }, + lambda m: m(torch.randn(16, 64)), + ) + assert isinstance(model.fc1.input_quantizer, SequentialQuantizer) + # No NVFP4 leaf anywhere, so nothing to reject and nothing to swap. + assert ( + _swap_in_nvfp4_act_headroom_calibrators( + model, anchor_percentile=1.0, upper_percentile=99.99, rho=16384.0 + ) + == [] + ) def test_swap_installs_calibrator_with_config_values(): From a596c0e4f1266698dd28251e93cc05186ad89642 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 30 Jul 2026 19:36:55 +0000 Subject: [PATCH 9/9] Select the weight-scale algorithm independently of the activation policy Weight-scale calibration is orthogonal to the activation global scale this algorithm sets, but it was hard-wired to max_calibrate, so a recipe could not combine, say, MSE weights with headroom activations. Chaining is not a substitute: [mse, nvfp4_act_headroom] has the second step's max_calibrate reset the MSE weights, and [nvfp4_act_headroom, mse] has MSE's internal max_calibrate reset the activation scales. Route weight calibration through the existing dispatch helper, renamed _run_weight_scale_calibration to say what it actually governs, and expose it as a nested, typed weight_scale_algorithm on the config (max / mse / local_hessian plus that algorithm's own options). Default explicitly to {"method": "max"} rather than inheriting the helper's None -> mse default, so selecting this activation policy never silently changes how weights are calibrated. distributed_sync, shared_states and sync_expert_weight_amax move off the top level into that nested config, where the union types them per method -- MaxCalibConfig defines sync_expert_weight_amax and the others do not. A field_serializer keeps the sparse public dict shape, which is load-bearing: wrapped_calib_func does a full model_dump, so without it every nested default would be forwarded as a kwarg and max_calibrate would raise on moe_calib_experts_ratio and layerwise. Tests cover the max default with sparse dispatch, mse selection with its options, and nested knob ownership. Asserting that MSE actually moves the weights needs triton (static NVFP4) or CUDA (dynamic), so these assert the dispatch and that the headroom activation scale survives. Pre-commit hooks were run manually on these files (all pass); the commit hook is skipped only because its autostash cannot write to a read-only cache in this environment. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 2 +- modelopt/torch/quantization/config.py | 147 +++++++++--------- modelopt/torch/quantization/model_calib.py | 48 +++--- .../ptq/nvfp4_act_headroom-kv_fp8_cast.yaml | 4 + tests/unit/torch/quantization/test_lsq.py | 4 +- .../quantization/test_nvfp4_act_headroom.py | 86 +++++++--- 6 files changed, 172 insertions(+), 119 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1dcd3b636fc..885dbf293ae 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,7 +8,7 @@ Changelog *Quantization* -- Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Plain ``max`` sets a tensor's global scale from the largest per-block amax seen during calibration, leaving no room above it, so any activation larger than the calibration max saturates. ``nvfp4_act_headroom`` instead anchors the global scale to a low percentile of the per-block amax distribution, ``amax = max(rho * anchor, upper)``, placing the calibrated blocks in the lower part of the FP8 block-scale range and leaving the rest as headroom. ``upper`` is the top of the range the scale commits to representing and defaults to the 99.99th percentile rather than the literal maximum: chasing a lone freak block would drag the global scale up until every other block's FP8 block scale falls below subnormal and flushes to zero, so the rarest blocks are clipped instead. Set ``upper_percentile=100`` to use the literal observed max, which guarantees no calibration data is clipped. The calibrator warns when the per-block range is too wide for ``rho`` to clear any headroom. Tunable via ``anchor_percentile`` (default 1), ``upper_percentile`` (default 99.99) and ``rho`` (default 16384). Applies only to NVFP4 dynamic-block input quantizers; weight quantizers and all other quantizers keep plain ``max``. ``SequentialQuantizer`` activation quantizers are not supported and raise. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` (dynamic NVFP4 W4A4 plus FP8 KV-cache cast) with only the calibration algorithm swapped, so it exports a standard NVFP4 checkpoint. +- Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Plain ``max`` sets a tensor's global scale from the largest per-block amax seen during calibration, leaving no room above it, so any activation larger than the calibration max saturates. ``nvfp4_act_headroom`` instead anchors the global scale to a low percentile of the per-block amax distribution, ``amax = max(rho * anchor, upper)``, placing the calibrated blocks in the lower part of the FP8 block-scale range and leaving the rest as headroom. ``upper`` is the top of the range the scale commits to representing and defaults to the 99.99th percentile rather than the literal maximum: chasing a lone freak block would drag the global scale up until every other block's FP8 block scale falls below subnormal and flushes to zero, so the rarest blocks are clipped instead. Set ``upper_percentile=100`` to use the literal observed max, which guarantees no calibration data is clipped. The calibrator warns when the per-block range is too wide for ``rho`` to clear any headroom. Tunable via ``anchor_percentile`` (default 1), ``upper_percentile`` (default 99.99) and ``rho`` (default 16384). Applies only to NVFP4 dynamic-block input quantizers. Weight scales are an orthogonal axis, selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian`` with that algorithm's own options), so one recipe can combine a weight calibration with this activation policy in a single pass. ``SequentialQuantizer`` activation quantizers are not supported and raise. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` (dynamic NVFP4 W4A4 plus FP8 KV-cache cast) with only the calibration algorithm swapped, so it exports a standard NVFP4 checkpoint. **Backward Breaking Changes** diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 5178eec095e..7ca5fd1f7a6 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -986,78 +986,6 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): ) -class NVFP4ActHeadroomCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): - """Config for the ``nvfp4_act_headroom`` calibration algorithm. - - Calibrates the per-tensor global scale of NVFP4 *activation* (input) quantizers so that - the calibrated range sits in the lower part of the FP8 block-scale range, leaving headroom - above it for activations larger than any seen during calibration. Weight quantizers and - all non-NVFP4 quantizers are calibrated with plain ``max``. - - The top of the calibrated range is ``upper_percentile`` (default 99.99) rather than the - literal maximum, so rare blocks far above the rest are clipped instead of dragging the - global scale up until every other block flushes to zero. - - See :class:`NVFP4ActHeadroomCalibrator - ` for the formula. - """ - - _mutates_weights: ClassVar[bool] = False - - method: Literal["nvfp4_act_headroom"] = ModeloptField("nvfp4_act_headroom") - - anchor_percentile: float = ModeloptField( - default=1.0, - gt=0.0, - le=100.0, - title="Percentile of the per-block activation amaxes used as the anchor.", - description=( - "The global scale is anchored to this percentile of the per-block amax " - "distribution. Lower values anchor further into the low tail, which yields a " - "smaller global scale and less headroom." - ), - ) - - upper_percentile: float = ModeloptField( - default=99.99, - gt=0.0, - le=100.0, - title="Percentile of the per-block activation amaxes used as the top of the range.", - description=( - "The global scale is floored at this percentile, so per-block amaxes above it are " - "clipped. The default excludes the rarest blocks on purpose: chasing a lone outlier " - "pushes every other block's FP8 block scale below subnormal. Set to 100 to use the " - "literal observed max, which guarantees no calibration data is clipped." - ), - ) - - rho: float = ModeloptField( - default=16384.0, - gt=0.0, - lt=28672.0, - title="Headroom factor applied to the anchor (amax = rho * anchor).", - description=( - "Larger rho leaves more headroom above the calibrated range and less room below " - "it. Must stay below 28672, the FP8-E4M3 normal dynamic range." - ), - ) - - distributed_sync: bool | None = ModeloptField( - default=True, - title="Whether to sync the amax across the distributed processes.", - description="If True, the amax will be synced across the distributed processes.", - ) - - sync_expert_weight_amax: bool = ModeloptField( - default=False, - title="Share one weight amax across local experts in a SequentialMLP MoE layer.", - description=( - "Forwarded to max calibration, which handles the weight quantizers for this " - "algorithm. See :class:`MaxCalibConfig` for details." - ), - ) - - class MseCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): """Configuration for per-tensor MSE calibration. @@ -1373,6 +1301,81 @@ def _gptq_qdq_default(self): _ScaleCalibConfig: TypeAlias = MaxCalibConfig | MseCalibConfig | LocalHessianCalibConfig +class NVFP4ActHeadroomCalibConfig(QuantizeAlgorithmConfig): + """Config for the ``nvfp4_act_headroom`` calibration algorithm. + + Calibrates the per-tensor global scale of NVFP4 *activation* (input) quantizers so that + the calibrated range sits in the lower part of the FP8 block-scale range, leaving headroom + above it for activations larger than any seen during calibration. Weight quantizers and + all non-NVFP4 quantizers are calibrated with plain ``max``. + + The top of the calibrated range is ``upper_percentile`` (default 99.99) rather than the + literal maximum, so rare blocks far above the rest are clipped instead of dragging the + global scale up until every other block flushes to zero. + + See :class:`NVFP4ActHeadroomCalibrator + ` for the formula. + """ + + _mutates_weights: ClassVar[bool] = False + + method: Literal["nvfp4_act_headroom"] = ModeloptField("nvfp4_act_headroom") + + anchor_percentile: float = ModeloptField( + default=1.0, + gt=0.0, + le=100.0, + title="Percentile of the per-block activation amaxes used as the anchor.", + description=( + "The global scale is anchored to this percentile of the per-block amax " + "distribution. Lower values anchor further into the low tail, which yields a " + "smaller global scale and less headroom." + ), + ) + + upper_percentile: float = ModeloptField( + default=99.99, + gt=0.0, + le=100.0, + title="Percentile of the per-block activation amaxes used as the top of the range.", + description=( + "The global scale is floored at this percentile, so per-block amaxes above it are " + "clipped. The default excludes the rarest blocks on purpose: chasing a lone outlier " + "pushes every other block's FP8 block scale below subnormal. Set to 100 to use the " + "literal observed max, which guarantees no calibration data is clipped." + ), + ) + + rho: float = ModeloptField( + default=16384.0, + gt=0.0, + lt=28672.0, + title="Headroom factor applied to the anchor (amax = rho * anchor).", + description=( + "Larger rho leaves more headroom above the calibrated range and less room below " + "it. Must stay below 28672, the FP8-E4M3 normal dynamic range." + ), + ) + + weight_scale_algorithm: _ScaleCalibConfig = ModeloptField( + default={"method": "max"}, + title="Algorithm used to calibrate the weight scales.", + description=( + "Weight scales are set by an independent algorithm -- ``max`` (default), ``mse`` or " + "``local_hessian`` -- because this algorithm only decides the NVFP4 *activation* " + "global scale. Give the chosen algorithm's own options alongside ``method`` (for " + "example ``{'method': 'mse', 'fp8_scale_sweep': true}``); ``distributed_sync`` and " + "``shared_states`` belong to that weight calibration pass and are set there." + ), + validate_default=True, + ) + + @field_serializer("weight_scale_algorithm") + def _serialize_weight_scale_algorithm(self, value: _ScaleCalibConfig): + """Preserve the sparse public dict shape accepted by this field.""" + return {"method": value.method, **value.model_dump(exclude={"method"}, exclude_unset=True)} + + class LSQConfig(QuantizeAlgorithmConfig): """Config for LSQ (Learnt Scale Quantization) and Dual-LSQ algorithms. diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 5c5963e25e2..aa8f766d7be 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -565,17 +565,19 @@ def nvfp4_act_headroom_calibrate( anchor_percentile: float = 1.0, upper_percentile: float = 99.99, rho: float = 16384.0, - distributed_sync: bool = True, - shared_states: Mapping[str, Mapping[str, Sequence[str]]] | None = None, - sync_expert_weight_amax: bool = False, + weight_scale_algorithm: Any = None, ): - """Calibrate NVFP4 activation global scales with headroom; everything else uses max. + """Calibrate NVFP4 activation global scales with headroom. For NVFP4 dynamic-block *input* quantizers, the per-tensor global scale is derived from the distribution of per-block activation amaxes so that headroom is left above the calibrated range (see :class:`NVFP4ActHeadroomCalibrator - `). Weight quantizers and - every other quantizer are calibrated exactly as in :func:`max_calibrate`. + `). + + Weight scales are an orthogonal concern and are delegated to ``weight_scale_algorithm``, + so a recipe can combine this activation policy with ``max``, ``mse`` or ``local_hessian`` + weights. Each of those runs max calibration first, so the activation collectors installed + here are populated in that pass and any later refinement touches only weights. Args: model: model to be calibrated. @@ -584,20 +586,19 @@ def nvfp4_act_headroom_calibrate( upper_percentile: percentile of the per-block amaxes used as the top of the calibrated range; ``100`` uses the literal observed max. rho: headroom factor; ``amax = rho * anchor``. Must be in ``(0, 28672)``. - distributed_sync: whether to sync amax across distributed processes - (see :func:`max_calibrate`). + weight_scale_algorithm: config for the algorithm that calibrates the *weight* scales -- + ``{"method": "max"}`` (the default), ``"mse"`` or ``"local_hessian"``, plus that + algorithm's own options such as ``distributed_sync`` and ``shared_states``. + ``None`` is treated as ``{"method": "max"}``. Raises: NotImplementedError: if an NVFP4 activation quantizer is a ``SequentialQuantizer``, which this algorithm does not support. - shared_states: shared quantization-state grouping, forwarded to :func:`max_calibrate`. - sync_expert_weight_amax: share one weight amax across local experts in a SequentialMLP - MoE layer, forwarded to :func:`max_calibrate`. .. note:: - Under data parallelism the per-rank scales are combined by :func:`max_calibrate` with a - ``MAX`` all-reduce, so the result is the largest per-rank headroom scale rather than the - scale implied by pooling every rank's per-block distribution. + Under data parallelism the per-rank scales are combined with a ``MAX`` all-reduce, so + the result is the largest per-rank headroom scale rather than the scale implied by + pooling every rank's per-block distribution. """ swapped = _swap_in_nvfp4_act_headroom_calibrators( model, anchor_percentile=anchor_percentile, upper_percentile=upper_percentile, rho=rho @@ -617,13 +618,11 @@ def nvfp4_act_headroom_calibrate( # per-block histograms in the same pass that collects max stats for every other quantizer. # The calibrators are restored afterwards so this algorithm does not leak into a later # calibration of the same model, and so a repeat run starts from a fresh histogram. + # Default to max rather than the helper's own default, so selecting this activation policy + # never silently changes how weights are calibrated. try: - max_calibrate( - model, - forward_loop, - distributed_sync=distributed_sync, - sync_expert_weight_amax=sync_expert_weight_amax, - shared_states=shared_states, + _run_weight_scale_calibration( + model, forward_loop, weight_scale_algorithm or {"method": "max"} ) finally: for quantizer, original_calibrator in swapped: @@ -2265,8 +2264,11 @@ def _make_gptq_handle(name, m): print_rank_0(f"GPTQ time: {time.time() - total_start:.2f}s") -def _run_scale_calibration(model, forward_loop, scale_algorithm): - """Run scale calibration.""" +def _run_weight_scale_calibration(model, forward_loop, scale_algorithm): + """Run the weight-scale calibration algorithm. + + These algorithms set *weight* scales; activation scales are the caller's concern. + """ if scale_algorithm is None: scale_algorithm = {"method": "mse"} @@ -2309,7 +2311,7 @@ def lsq( tied_amax: If True, pre and post share a single tensor. quantize_pre_scale: If False, skip FP8 quantization for the LSQ pre scale. """ - _run_scale_calibration(model, forward_loop, scale_algorithm) + _run_weight_scale_calibration(model, forward_loop, scale_algorithm) name_to_module = dict(model.named_modules()) seen_modules: set[int] = set() diff --git a/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml b/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml index 21a71f4a10b..a0e79316eb4 100644 --- a/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml @@ -40,6 +40,10 @@ quantize: anchor_percentile: 1 upper_percentile: 99.99 rho: 16384 + # Weight scales are an orthogonal axis: swap in mse or local_hessian here to combine a + # different weight calibration with the same activation policy. + weight_scale_algorithm: + method: max quant_cfg: - $import: base_disable_all - $import: w4a4_nvfp4_nvfp4 diff --git a/tests/unit/torch/quantization/test_lsq.py b/tests/unit/torch/quantization/test_lsq.py index 50fd1ffc870..1a1136d6ece 100644 --- a/tests/unit/torch/quantization/test_lsq.py +++ b/tests/unit/torch/quantization/test_lsq.py @@ -58,7 +58,7 @@ def _make_int4_static_quantizer(): def _skip_scale_calibration(monkeypatch): monkeypatch.setattr( - "modelopt.torch.quantization.model_calib._run_scale_calibration", + "modelopt.torch.quantization.model_calib._run_weight_scale_calibration", lambda *args, **kwargs: None, ) @@ -130,7 +130,7 @@ def test_scale_algorithm_preserves_sparse_dict(self, monkeypatch): calibrate = create_autospec(model_calib_module.mse_calibrate) monkeypatch.setattr(model_calib_module, "mse_calibrate", calibrate) model = Mock() - model_calib_module._run_scale_calibration(model, None, cfg.scale_algorithm) + model_calib_module._run_weight_scale_calibration(model, None, cfg.scale_algorithm) calibrate.assert_called_once_with(model, forward_loop=None, fp8_scale_sweep=True) @pytest.mark.parametrize( diff --git a/tests/unit/torch/quantization/test_nvfp4_act_headroom.py b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py index c7ec9598e42..cc73532ed07 100644 --- a/tests/unit/torch/quantization/test_nvfp4_act_headroom.py +++ b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py @@ -177,32 +177,76 @@ def test_calibrator_is_restored_after_calibration(): assert float(model.fc1.input_quantizer.amax) < headroom_amax -def test_shared_state_knobs_are_forwarded_to_max_calibrate(): - """shared_states / sync_expert_weight_amax reach max_calibrate, as they do for `max`.""" - seen = {} - real = model_calib.max_calibrate +def _spy_on(name, calls, stand_in=None): + """Replace a model_calib entry point with a recording spy, returning the original.""" + real = getattr(model_calib, name) + run = stand_in if stand_in is not None else real - def spy(model, forward_loop=None, distributed_sync=True, **kwargs): - seen.update(kwargs) - return real(model, forward_loop, distributed_sync) + def spy(model, forward_loop=None, **kwargs): + calls.append((name, kwargs)) + return run(model, forward_loop) - cfg = { - **ACT_ONLY_CFG, - "algorithm": { - "method": "nvfp4_act_headroom", - "anchor_percentile": 1, - "shared_states": {"weight_global_amax": {"patterns": [r".*fc1"]}}, - "sync_expert_weight_amax": True, - }, - } - model_calib.max_calibrate = spy + setattr(model_calib, name, spy) + return real + + +def _quantize_with(weight_scale_algorithm, calls, spy_name, stand_in=None): + cfg = {**ACT_ONLY_CFG, "algorithm": {**ACT_ONLY_CFG["algorithm"]}} + if weight_scale_algorithm is not None: + cfg["algorithm"]["weight_scale_algorithm"] = weight_scale_algorithm + real = _spy_on(spy_name, calls, stand_in) try: - mtq.quantize(_Net(), cfg, lambda m: m(torch.randn(8, 64))) + torch.manual_seed(0) + return mtq.quantize(_Net(), cfg, lambda m: m(torch.randn(16, 64))) finally: - model_calib.max_calibrate = real + setattr(model_calib, spy_name, real) + + +def test_weight_scale_algorithm_defaults_to_max(): + """Choosing this activation policy must not change how weights are calibrated.""" + calls = [] + model = _quantize_with(None, calls, "max_calibrate") + assert [c[0] for c in calls] == ["max_calibrate"] + # Sparse dispatch: none of the nested config's unset defaults are forwarded as kwargs. + assert calls[0][1] == {} + assert float(model.fc1.input_quantizer.amax) > 0 + - assert seen["sync_expert_weight_amax"] is True - assert seen["shared_states"] == {"weight_global_amax": {"patterns": [r".*fc1"]}} +def test_weight_scale_algorithm_dispatches_to_mse_with_its_options(): + """A different weight algorithm is selectable, and the activation scale is unaffected. + + The real MSE weight pass needs triton (static NVFP4) or CUDA (dynamic), so this asserts the + dispatch and option plumbing and that the headroom activation scale still lands. + """ + calls = [] + model = _quantize_with( + {"method": "mse", "fp8_scale_sweep": True}, + calls, + "mse_calibrate", + stand_in=model_calib.max_calibrate, + ) + assert [c[0] for c in calls] == ["mse_calibrate"] + assert calls[0][1] == {"fp8_scale_sweep": True} + data_max = float(torch.randn(16, 64).abs().max()) + assert float(model.fc1.input_quantizer.amax) > data_max # still the headroom scale + + +def test_weight_scale_algorithm_owns_the_shared_state_knobs(): + """distributed_sync / shared_states / sync_expert_weight_amax belong to the weight pass.""" + calls = [] + _quantize_with( + { + "method": "max", + "sync_expert_weight_amax": True, + "shared_states": {"weight_global_amax": {"patterns": [r".*fc1"]}}, + }, + calls, + "max_calibrate", + ) + assert calls[0][1] == { + "sync_expert_weight_amax": True, + "shared_states": {"weight_global_amax": {"patterns": [r".*fc1"]}}, + } def test_lower_anchor_percentile_gives_smaller_amax():