diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 964fd8483fc..fe45b45daa5 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,7 @@ Changelog *Quantization* +- Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference: fake-quant weights affected by quantization, pre-quant scaling, or rotation are snapshotted on a configurable device, folded through the module's native ``fold_weight`` implementation for the duration of a context, then restored with their quantizer runtime state. Retained weight pre-quant scales are inactive while folded to avoid applying them twice. Sharing a weight or weight quantizer across quantized modules and ``SequentialQuantizer`` weights are not supported. - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. *Megatron Framework (M-LM / M-Bridge)* diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 966a3643fe3..315b163da88 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -20,6 +20,7 @@ import os import warnings from collections.abc import Callable, Iterable +from contextlib import contextmanager from typing import Any, cast import torch @@ -40,7 +41,7 @@ from .algorithms import get_auto_quantize_config as _get_auto_quantize_config from .config import QuantizeAlgoCfgType from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg -from .nn import QuantModule, TensorQuantizer +from .nn import QuantModule, SequentialQuantizer, TensorQuantizer from .utils import is_quantized __all__ = [ @@ -54,6 +55,7 @@ "postprocess_amax", "print_quant_summary", "quantize", + "temporarily_fold_weights", ] @@ -733,6 +735,92 @@ def fold_weight(model: nn.Module, keep_attrs: bool = False): module.fold_weight(keep_attrs) +@contextmanager +def temporarily_fold_weights( + model: nn.Module, + snapshot_device: torch.device | str | None = None, +): + """Temporarily fold fake-quant weights for a frozen inference region. + + Each :class:`QuantModule` performs its normal module-specific ``fold_weight`` operation. + Fake-quant weights affected by quantization, pre-quant scaling, or rotation and all fake-quant + runtime states are restored on exit, including after an exception. Weights are restored in + place so optimizer and distributed references remain valid. + + This context is intended for repeated no-gradient forwards with no optimizer step, such as + log-probability recomputation over several microbatches. It retains calibration attributes + while folded; a retained weight ``pre_quant_scale`` is inactive inside the context because its + value is already baked into the temporary weight. Sharing a weight or weight quantizer across + multiple :class:`QuantModule` instances and using :class:`SequentialQuantizer` are not + supported. + + Example:: + + with mtq.temporarily_fold_weights(model, snapshot_device="cpu"): + outputs = model(inputs) + + Args: + model: Quantized model whose weights will be temporarily folded. + snapshot_device: Device used to store parameter snapshots. ``None`` keeps each snapshot + on the parameter's device; ``"cpu"`` avoids the additional accelerator memory. + """ + fold_pairs = [] + for module in model.modules(): + if not isinstance(module, QuantModule): + continue + for weight, quantizer in module.iter_weights_for_calibration(): + if not isinstance(weight, torch.Tensor): + continue + if isinstance(quantizer, SequentialQuantizer): + raise NotImplementedError( + "temporarily_fold_weights does not support SequentialQuantizer" + ) + if not isinstance(quantizer, TensorQuantizer) or not quantizer.fake_quant: + continue + local_weight = weight.to_local() if hasattr(weight, "to_local") else weight + fold_pairs.append((local_weight, quantizer)) + + weight_snapshots = {} + quantizer_states = {} + for weight, quantizer in fold_pairs: + if ( + quantizer.is_enabled + or quantizer.pre_quant_scale is not None + or quantizer.rotate_is_enabled + ): + weight_id = id(weight) + if weight_id not in weight_snapshots: + weight_snapshots[weight_id] = ( + weight, + weight.detach().clone() + if snapshot_device is None + else weight.detach().to(snapshot_device, copy=True), + ) + quantizer_states.setdefault( + quantizer, + ( + quantizer._disabled, + quantizer._rotate, + quantizer._enable_pre_quant_scale, + quantizer._input_dtype, + ), + ) + try: + fold_weight(model, keep_attrs=True) + yield + finally: + with torch.no_grad(): + for weight, snapshot in weight_snapshots.values(): + weight.copy_(snapshot) + for quantizer, state in quantizer_states.items(): + ( + quantizer._disabled, + quantizer._rotate, + quantizer._enable_pre_quant_scale, + quantizer._input_dtype, + ) = state + + @torch.no_grad() def compute_quantization_mse( model: nn.Module, diff --git a/modelopt/torch/quantization/nn/modules/quant_linear.py b/modelopt/torch/quantization/nn/modules/quant_linear.py index da1b79a2f60..d0569ff5b71 100644 --- a/modelopt/torch/quantization/nn/modules/quant_linear.py +++ b/modelopt/torch/quantization/nn/modules/quant_linear.py @@ -164,14 +164,18 @@ def forward(self, input, *args, **kwargs): def fold_weight(self, keep_attrs: bool = False): """Fold the weight for faster eval.""" - super().fold_weight(keep_attrs) - if ( + should_fold = ( hasattr(self, "weight_quantizer") and hasattr(self, "weight") + and isinstance(self.weight, torch.Tensor) + and isinstance(self.weight_quantizer, TensorQuantizer) and self.weight_quantizer.fake_quant - ): + ) + super().fold_weight(keep_attrs) + if should_fold: if ( - self._not_sequential_quantizers() + not keep_attrs + and self._not_sequential_quantizers() and self.weight_quantizer.svdquant_lora_a is not None and self.weight_quantizer.svdquant_lora_b is not None ): diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py index 9c9aee478a8..83ee8940d19 100644 --- a/modelopt/torch/quantization/nn/modules/quant_module.py +++ b/modelopt/torch/quantization/nn/modules/quant_module.py @@ -141,22 +141,31 @@ def _fold_weight_quantizer( per-tensor quantizer over all experts of a fused MoE weight) can be folded view by view while disabling and dropping its calibration attrs exactly once. """ + if not quantizer.fake_quant: + return + for weight in weights: weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype)) quantizer.disable() quantizer.disable_rotate() - if not keep_attrs: - for attr_name in ("_pre_quant_scale", "_amax"): - if hasattr(quantizer, attr_name): - delattr(quantizer, attr_name) + if hasattr(quantizer, "_pre_quant_scale"): + if keep_attrs: + # The scale is already baked into the folded weight. + quantizer._enable_pre_quant_scale = False + else: + delattr(quantizer, "_pre_quant_scale") + if not keep_attrs and hasattr(quantizer, "_amax"): + delattr(quantizer, "_amax") def fold_weight(self, keep_attrs: bool = False): """Bake each fake-quant weight quantizer into its weight for faster eval. - Every fake-quant weight quantizer is folded regardless of its enabled state. The folded - transform is baked into the stored weight and then disabled, so subsequent forwards use - the stored weight directly. Calibration buffers (``_pre_quant_scale``, ``_amax``) are - dropped unless ``keep_attrs``. + Every fake-quant weight quantizer is folded, including disabled quantizers whose pre-quant + scale or rotation remains active. The folded transform is baked into the stored weight and + then disabled, so subsequent forwards use the stored weight directly. Calibration buffers + (``_pre_quant_scale``, ``_amax``) are dropped unless ``keep_attrs``. A retained pre-quant + scale remains stored but is made inactive because it has already been applied to the folded + weight. """ # Handle all attributes that end with _weight_quantizer for name in dir(self): @@ -173,7 +182,8 @@ def fold_weight(self, keep_attrs: bool = False): f"{name} doesn't have a corresponding {weight_name} in {self.__class__.__name__}" ) weight = getattr(self, weight_name) - self._fold_weight_quantizer(attr, (weight,), keep_attrs) + if isinstance(weight, torch.Tensor): + self._fold_weight_quantizer(attr, (weight,), keep_attrs) QuantModuleRegistry = _DMRegistryCls("Quant", QuantModule) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index dde24513086..4acb4d30dfa 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -524,8 +524,10 @@ def enable_weight_access_and_writeback(self): weight = self.weight # TODO: To support TP + FSDP, we need to redistribute the tensor with replicate instead of shard self.weight = nn.Parameter(weight.to_local()) - yield - self.weight = weight + try: + yield + finally: + self.weight = weight else: # transformers>=5.0: weights are already plain Parameters yield diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py index 95212435d87..c2ae02a7f93 100644 --- a/modelopt/torch/quantization/plugins/transformer_engine.py +++ b/modelopt/torch/quantization/plugins/transformer_engine.py @@ -245,6 +245,16 @@ def iter_weights_for_calibration(self): if weight_i is not None: yield weight_i, (self.weight_quantizer[i] if grouped else self.weight_quantizer) + def fold_weight(self, keep_attrs: bool = False): + """Fold each grouped weight with its corresponding fake-quant quantizer.""" + quantizer_weights: dict[TensorQuantizer, list[torch.Tensor]] = {} + for weight, quantizer in self.iter_weights_for_calibration(): + if isinstance(quantizer, TensorQuantizer) and quantizer.fake_quant: + quantizer_weights.setdefault(quantizer, []).append(weight) + + for quantizer, weights in quantizer_weights.items(): + self._fold_weight_quantizer(quantizer, weights, keep_attrs) + @staticmethod def te_grouped_quantized_linear_fn(package, func_name, self, *args): _assert_te_fp8_enabled() diff --git a/modelopt/torch/quantization/plugins/vllm.py b/modelopt/torch/quantization/plugins/vllm.py index 749d6b10ee5..67191a0f3c7 100644 --- a/modelopt/torch/quantization/plugins/vllm.py +++ b/modelopt/torch/quantization/plugins/vllm.py @@ -562,6 +562,11 @@ def _setup(self): ) self.parallel_state = create_parallel_state() + def iter_weights_for_calibration(self): + """Yield the fused MoE weights with their corresponding quantizers.""" + yield self.w13_weight, self.w13_weight_quantizer + yield self.w2_weight, self.w2_weight_quantizer + def invoke_fused_moe_quantized( self, A: torch.Tensor, # noqa: N803 diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py b/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py index 7671a1c28d1..13ecc941c76 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py @@ -25,7 +25,7 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.quantization.extensions import get_cuda_ext_mx -from modelopt.torch.quantization.nn import QuantModule +from modelopt.torch.quantization.nn import GroupedQuantizer, QuantModule class TELinear(nn.Module): @@ -118,6 +118,40 @@ def test_quantize(model_cls, config): quantize_model_and_forward(model, config, calib_data) +@pytest.mark.parametrize("share_weight_quantizer", [False, True]) +def test_fold_weight_grouped_linear(share_weight_quantizer): + model = TEGroupedLinear().cuda() + calib_data = [model.get_input().cuda()] + quantize_model_and_forward(model, mtq.INT8_DEFAULT_CFG, calib_data) + + grouped_linear = model.net + assert isinstance(grouped_linear.weight_quantizer, GroupedQuantizer) + if share_weight_quantizer: + grouped_linear.weight_quantizer = grouped_linear.weight_quantizer[0] + weights = [getattr(grouped_linear, f"weight{i}") for i in range(grouped_linear.num_gemms)] + quantizers = [quantizer for _, quantizer in grouped_linear.iter_weights_for_calibration()] + with torch.no_grad(): + output_before = model(calib_data[0]) + expected_weights = [ + quantizer(weight.float().contiguous()).to(weight.dtype) + for weight, quantizer in zip(weights, quantizers) + ] + + mtq.fold_weight(model) + + with torch.no_grad(): + output_after = model(calib_data[0]) + if isinstance(output_before, tuple): + output_before = output_before[0] + output_after = output_after[0] + assert torch.allclose(output_after, output_before) + + for weight, expected_weight, quantizer in zip(weights, expected_weights, quantizers): + assert torch.allclose(weight, expected_weight) + assert not quantizer.is_enabled + assert not hasattr(quantizer, "_amax") + + def test_quantize_forward_backward(): set_seed() model = TELinear().cuda() diff --git a/tests/unit/torch/quantization/plugins/test_huggingface.py b/tests/unit/torch/quantization/plugins/test_huggingface.py index 43a875d7336..0da6cfbb1fc 100644 --- a/tests/unit/torch/quantization/plugins/test_huggingface.py +++ b/tests/unit/torch/quantization/plugins/test_huggingface.py @@ -35,6 +35,7 @@ from modelopt.recipe.loader import load_recipe from modelopt.torch.quantization.nn import QuantLinear, QuantModuleRegistry, TensorQuantizer from modelopt.torch.quantization.plugins.huggingface import ( + _QuantHFParallelLinear, _TransposedExpertsCalibMixin, get_homogeneous_hf_decoder_layers, is_homogeneous_hf_model, @@ -77,6 +78,32 @@ def forward(self, x): return self.net(x) +def test_hf_parallel_weight_access_restores_dtensor_after_exception(monkeypatch): + class FakeDTensor: + placements = ("shard",) + + @staticmethod + def to_local(): + return torch.ones(2, 2) + + class ParallelLinear: + weight = FakeDTensor() + shard = FakeDTensor.placements + + monkeypatch.setattr(torch.distributed.tensor, "DTensor", FakeDTensor) + linear = ParallelLinear() + original_weight = linear.weight + + with ( + pytest.raises(RuntimeError, match="test error"), + _QuantHFParallelLinear.enable_weight_access_and_writeback(linear), + ): + assert isinstance(linear.weight, nn.Parameter) + raise RuntimeError("test error") + + assert linear.weight is original_weight + + def test_convert_conv1d(): set_seed() assert transformers.pytorch_utils.Conv1D in QuantModuleRegistry diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index b609761f12a..aa31d0b66d6 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -404,14 +404,17 @@ def test_postprocess_amax(): def test_svdquant_lora_weights(): model = _SimpleMLP(64, 64, 64, 64) + inputs = torch.randn(2, 64, 64) quant_config = mtq.INT8_SMOOTHQUANT_CFG.copy() quant_config["algorithm"] = "svdquant" - mtq.quantize(model, quant_config, partial(forward_loop, dataloader=[torch.randn(2, 64, 64)])) + mtq.quantize(model, quant_config, partial(forward_loop, dataloader=[inputs])) + original_weights = [] for module in model.modules(): if isinstance(module, torch.nn.Linear): + original_weights.append((module, module.weight.detach().clone())) assert module.weight_quantizer.svdquant_lora_a is not None assert module.weight_quantizer.svdquant_lora_b is not None @@ -420,6 +423,22 @@ def test_svdquant_lora_weights(): ) assert lora_residual.shape == module.weight.shape + output_before = model(inputs) + with mtq.temporarily_fold_weights(model): + output_folded = model(inputs) + for module, _ in original_weights: + assert module.weight_quantizer.svdquant_lora_a is not None + assert module.weight_quantizer.svdquant_lora_b is not None + assert not module.weight_quantizer.is_enabled + output_restored = model(inputs) + + assert torch.allclose(output_folded, output_before) + assert torch.allclose(output_restored, output_before) + for module, original_weight in original_weights: + assert torch.equal(module.weight, original_weight) + assert module.weight_quantizer.svdquant_lora_a is not None + assert module.weight_quantizer.svdquant_lora_b is not None + def test_layerwise_calibrate_support_gate(): class _UnsupportedModel(nn.Module): diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py index ba352ec2162..f8f3540a790 100644 --- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py @@ -317,21 +317,132 @@ def test_fold_weight_disables_quantizer_without_extra_transform( unregister_quant_backend(backend_name) -def test_fold_weight_keep_attrs_keeps_amax(monkeypatch): +def test_fold_weight_keep_attrs_keeps_calibration_attrs_inactive(monkeypatch): calls = [] backend_name = "test_fold_backend_keep" qlinear = _make_qlinear_with_backend(monkeypatch, calls, backend_name) try: qlinear.weight_quantizer.amax = torch.tensor(1.0) + qlinear.weight_quantizer.pre_quant_scale = torch.tensor([[0.5, 1.0, 1.5, 2.0]]) + inputs = torch.randn(2, 4) + output_before = qlinear(inputs) qlinear.fold_weight(keep_attrs=True) assert hasattr(qlinear.weight_quantizer, "_amax") + assert hasattr(qlinear.weight_quantizer, "_pre_quant_scale") + assert qlinear.weight_quantizer.pre_quant_scale is None assert not qlinear.weight_quantizer.is_enabled + assert torch.allclose(qlinear(inputs), output_before) finally: unregister_quant_backend(backend_name) +def test_temporarily_fold_weights_restores_after_exception(monkeypatch): + calls = [] + backend_name = "test_temporary_fold_backend" + qlinear = _make_qlinear_with_backend( + monkeypatch, + calls, + backend_name, + rotate={"enable": True}, + ) + try: + quantizer = qlinear.weight_quantizer + quantizer.amax = torch.tensor(1.0) + quantizer.pre_quant_scale = torch.tensor([[0.5, 1.0, 1.5, 2.0]]) + original_weight = qlinear.weight.detach().clone() + original_storage = qlinear.weight.data_ptr() + original_rotate = quantizer._rotate + original_pre_quant_scale = quantizer.pre_quant_scale.detach().clone() + original_input_dtype = quantizer._input_dtype + inputs = torch.randn(2, 4) + output_before = qlinear(inputs) + + with ( + pytest.raises(RuntimeError, match="test error"), + mtq.temporarily_fold_weights(qlinear, snapshot_device="cpu"), + ): + assert not torch.equal(qlinear.weight, original_weight) + assert qlinear.weight.data_ptr() == original_storage + assert not quantizer.is_enabled + assert not quantizer.rotate_is_enabled + assert hasattr(quantizer, "_pre_quant_scale") + assert quantizer.pre_quant_scale is None + assert torch.allclose(qlinear(inputs), output_before) + raise RuntimeError("test error") + + assert torch.equal(qlinear.weight, original_weight) + assert qlinear.weight.data_ptr() == original_storage + assert quantizer.is_enabled + assert quantizer._rotate == original_rotate + assert torch.equal(quantizer.pre_quant_scale, original_pre_quant_scale) + assert quantizer._input_dtype == original_input_dtype + assert torch.allclose(qlinear(inputs), output_before) + finally: + unregister_quant_backend(backend_name) + + +def test_temporarily_fold_weights_handles_disabled_transform_and_none_weight(): + disabled = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False)) + pre_quant_scale = torch.tensor([[0.5, 1.0, 1.5, 2.0]]) + disabled.weight_quantizer.pre_quant_scale = pre_quant_scale + disabled.weight_quantizer.disable() + original_weight = disabled.weight.detach().clone() + + tied_output = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False)) + tied_output.weight = None + + with mtq.temporarily_fold_weights(torch.nn.ModuleList([disabled, tied_output])): + assert torch.equal(disabled.weight, original_weight * pre_quant_scale) + assert not disabled.weight_quantizer.is_enabled + assert disabled.weight_quantizer.pre_quant_scale is None + assert tied_output.weight_quantizer.is_enabled + + assert torch.equal(disabled.weight, original_weight) + assert not disabled.weight_quantizer.is_enabled + assert torch.equal(disabled.weight_quantizer.pre_quant_scale, pre_quant_scale) + assert tied_output.weight_quantizer.is_enabled + + +def test_temporarily_fold_weights_restores_when_folding_fails(monkeypatch): + first_calls = [] + second_calls = [] + first_backend = "test_temporary_fold_first_backend" + second_backend = "test_temporary_fold_failing_backend" + first = _make_qlinear_with_backend(monkeypatch, first_calls, first_backend) + second = _make_qlinear_with_backend(monkeypatch, second_calls, second_backend) + + def failing_backend(inputs, _quantizer): + raise RuntimeError("fold failed") + + unregister_quant_backend(second_backend) + register_quant_backend(second_backend, failing_backend) + model = torch.nn.ModuleList([first, second]) + original_weights = [module.weight.detach().clone() for module in model] + try: + with pytest.raises(RuntimeError, match="fold failed"), mtq.temporarily_fold_weights(model): + pass + + for module, original_weight in zip(model, original_weights): + assert torch.equal(module.weight, original_weight) + assert module.weight_quantizer.is_enabled + finally: + unregister_quant_backend(first_backend) + unregister_quant_backend(second_backend) + + +def test_temporarily_fold_weights_rejects_sequential_quantizer(): + qlinear = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False)) + qlinear.weight_quantizer = SequentialQuantizer(TensorQuantizer(), TensorQuantizer()) + + with ( + pytest.raises(NotImplementedError, match="does not support SequentialQuantizer"), + mtq.temporarily_fold_weights(qlinear), + ): + pass + + WINT4INT8_CFG = { "quant_cfg": [ {"quantizer_name": "*", "enable": False},