NVIDIA · mxinO · Aug 11, 2026 · Aug 12, 2026 · Aug 12, 2026 · Aug 12, 2026
@@ -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)*

@@ -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,

@@ -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
             ):

@@ -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)

@@ -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
 

@@ -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()

@@ -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

@@ -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()

@@ -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

@@ -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):