diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 329e5d21f05..885dbf293ae 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, 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** **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..2c94f02a819 --- /dev/null +++ b/modelopt/torch/quantization/calib/nvfp4_act_headroom.py @@ -0,0 +1,230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""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 + +# 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 + + +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, 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. + + ``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. + 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; 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. + 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, + upper_percentile=99.99, + 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 + 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) + 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.""" + 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, 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) + ) + + 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() + 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()) + # 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 + ) + 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 + + 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=upper / _ANCHOR_FLOOR_RATIO if upper else None + ) + if upper + else None + ) + 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() + + # ``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 <= upper: + span = upper / anchor + warnings.warn( + 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 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 = upper + else: + amax = headroom_amax + + 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"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 6257623d108..7ca5fd1f7a6 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -1301,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/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..aa8f766d7be 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 @@ -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,144 @@ def sync_quantizer_amax_across_tp( _finalize_with_shared_state(model, weight_patterns) +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): + 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 _is_nvfp4_dynamic_input_quantizer(name: str, module: nn.Module) -> bool: + """True for a plain NVFP4 dynamic-block *input* quantizer. + + :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. + """ + return name.endswith("input_quantizer") and _is_nvfp4_dynamic_activation_quantizer(module) + + +def _swap_in_nvfp4_act_headroom_calibrators( + 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. + + Returns ``(quantizer, original_calibrator)`` pairs so the caller can restore them. + """ + # 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") or not isinstance(module, SequentialQuantizer): + continue + 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 + + +@torch.no_grad() +def nvfp4_act_headroom_calibrate( + model: nn.Module, + forward_loop: ForwardLoop | None = None, + *, + anchor_percentile: float = 1.0, + upper_percentile: float = 99.99, + rho: float = 16384.0, + weight_scale_algorithm: Any = None, +): + """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 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. + 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)``. + 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. + + .. note:: + 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 + ) + 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 {len(swapped)} NVFP4 activation quantizer(s) " + 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. + # 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: + _run_weight_scale_calibration( + model, forward_loop, weight_scale_algorithm or {"method": "max"} + ) + finally: + for quantizer, original_calibrator in swapped: + quantizer._calibrator = original_calibrator + + def _mse_quant_func(x, amax, quantizer): """Quantization function for MSE calibration.""" original_amax = quantizer._amax.clone() if hasattr(quantizer, "_amax") else None @@ -2125,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"} @@ -2169,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 new file mode 100644 index 00000000000..a0e79316eb4 --- /dev/null +++ b/modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Composed PTQ recipe for 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 + # 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 + # 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 + - $import: kv_fp8_cast + - $import: default_disabled_quantizers 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_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 new file mode 100644 index 00000000000..cc73532ed07 --- /dev/null +++ b/tests/unit/torch/quantization/test_nvfp4_act_headroom.py @@ -0,0 +1,435 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the ``nvfp4_act_headroom`` activation global-scale calibration.""" + +import warnings + +import pytest +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 SequentialQuantizer, TensorQuantizer +from modelopt.torch.quantization.utils import reduce_block_amax + +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_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, rel=0.06) # within one histogram bin of the top block + + +@pytest.mark.parametrize( + "case", + [ + "gaussian", + "heavy_tail", + "sparse", + "outlier", + "wide_dynamic_range", + "constant", + ], +) +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) + 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, upper_percentile=100.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_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"): + 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.""" + 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 _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, **kwargs): + calls.append((name, kwargs)) + return run(model, forward_loop) + + 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: + torch.manual_seed(0) + return mtq.quantize(_Net(), cfg, lambda m: m(torch.randn(16, 64))) + finally: + 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 + + +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(): + """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_sequential_quantizer_activation_is_rejected(): + """SequentialQuantizer activations are unsupported and must fail loudly, not silently. + + 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": [ + {"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) + with pytest.raises(NotImplementedError, match="does not support SequentialQuantizer"): + mtq.quantize(_Net(), cfg, lambda m: m(torch.randn(16, 64))) + + +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)), + ) + 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 + + +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(): + 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, 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 + + +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)