From 8786941273a21dd0ea76da31607b2827a70e1b18 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Tue, 11 Aug 2026 06:12:51 +0000
Subject: [PATCH 01/16] Fix folding quantizer containers

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 .../quantization/nn/modules/quant_module.py   | 18 ++++++---
 .../plugins/transformer_engine.py             | 13 ++++++
 .../plugins/test_transformer_engine.py        | 40 ++++++++++++++++++-
 .../quantization/test_tensor_quant_cpu.py     | 34 ++++++++++++++++
 4 files changed, 98 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py
index 9c9aee478a8..829cd3f10b1 100644
--- a/modelopt/torch/quantization/nn/modules/quant_module.py
+++ b/modelopt/torch/quantization/nn/modules/quant_module.py
@@ -131,7 +131,7 @@ def iter_weights_for_calibration(self):
     @staticmethod
     @torch.no_grad()
     def _fold_weight_quantizer(
-        quantizer: TensorQuantizer,
+        quantizer: TensorQuantizer | SequentialQuantizer,
         weights: Iterable[torch.Tensor],
         keep_attrs: bool = False,
     ):
@@ -145,10 +145,16 @@ def _fold_weight_quantizer(
             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)
+        tensor_quantizers = (
+            quantizer if isinstance(quantizer, SequentialQuantizer) else (quantizer,)
+        )
+        for tensor_quantizer in tensor_quantizers:
+            # The pre-quant scale is already baked into the folded weight.
+            tensor_quantizer._enable_pre_quant_scale = False
+            if not keep_attrs:
+                for attr_name in ("_pre_quant_scale", "_amax"):
+                    if hasattr(tensor_quantizer, attr_name):
+                        delattr(tensor_quantizer, attr_name)
 
     def fold_weight(self, keep_attrs: bool = False):
         """Bake each fake-quant weight quantizer into its weight for faster eval.
@@ -163,7 +169,7 @@ def fold_weight(self, keep_attrs: bool = False):
             attr = getattr(self, name)
             if (
                 name.endswith("weight_quantizer")
-                and isinstance(attr, TensorQuantizer)
+                and isinstance(attr, (TensorQuantizer, SequentialQuantizer))
                 and attr.fake_quant
             ):
                 # Get the corresponding weight name by removing _weight_quantizer suffix
diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py
index 95212435d87..c7517eb559c 100644
--- a/modelopt/torch/quantization/plugins/transformer_engine.py
+++ b/modelopt/torch/quantization/plugins/transformer_engine.py
@@ -245,6 +245,19 @@ 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 | SequentialQuantizer, list[torch.Tensor]] = {}
+        for weight, quantizer in self.iter_weights_for_calibration():
+            if (
+                isinstance(quantizer, (TensorQuantizer, SequentialQuantizer))
+                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/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py b/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py
index 7671a1c28d1..adba5443ba4 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, SequentialQuantizer
 
 
 class TELinear(nn.Module):
@@ -118,6 +118,44 @@ def test_quantize(model_cls, config):
     quantize_model_and_forward(model, config, calib_data)
 
 
+@pytest.mark.parametrize("share_weight_quantizer", [False, True])
+@pytest.mark.parametrize("sequential_weight_quantizer", [False, True])
+def test_fold_weight_grouped_linear(share_weight_quantizer, sequential_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 sequential_weight_quantizer:
+        grouped_linear.weight_quantizer = GroupedQuantizer(
+            *(
+                SequentialQuantizer(quantizer, copy.deepcopy(quantizer))
+                for quantizer in grouped_linear.weight_quantizer
+            )
+        )
+    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():
+        expected_weights = [
+            quantizer(weight.float().contiguous()).to(weight.dtype)
+            for weight, quantizer in zip(weights, quantizers)
+        ]
+
+    mtq.fold_weight(model)
+
+    for weight, expected_weight, quantizer in zip(weights, expected_weights, quantizers):
+        assert torch.allclose(weight, expected_weight)
+        assert not quantizer.is_enabled
+        tensor_quantizers = (
+            quantizer if isinstance(quantizer, SequentialQuantizer) else (quantizer,)
+        )
+        for tensor_quantizer in tensor_quantizers:
+            assert not hasattr(tensor_quantizer, "_amax")
+
+
 def test_quantize_forward_backward():
     set_seed()
     model = TELinear().cuda()
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index ba352ec2162..eba48dfb4c3 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -332,6 +332,40 @@ def test_fold_weight_keep_attrs_keeps_amax(monkeypatch):
         unregister_quant_backend(backend_name)
 
 
+@pytest.mark.parametrize("keep_attrs", [False, True])
+def test_fold_weight_supports_sequential_quantizer(keep_attrs):
+    qlinear = QuantModuleRegistry.convert(torch.nn.Linear(4, 3))
+    qlinear.input_quantizer.disable()
+    qlinear.output_quantizer.disable()
+    qlinear.weight_quantizer = SequentialQuantizer(
+        TensorQuantizer(QuantizerAttributeConfig(num_bits=4)),
+        TensorQuantizer(QuantizerAttributeConfig(num_bits=8)),
+    )
+    pre_quant_scales = (
+        torch.tensor([0.5, 1.0, 1.5, 2.0]),
+        torch.tensor([2.0, 1.5, 1.0, 0.5]),
+    )
+    for quantizer, pre_quant_scale in zip(qlinear.weight_quantizer, pre_quant_scales):
+        quantizer.amax = torch.tensor(1.0)
+        quantizer.pre_quant_scale = pre_quant_scale
+    x = torch.randn(2, 4)
+    out_before = qlinear(x)
+    with torch.no_grad():
+        expected_weight = qlinear.weight_quantizer(qlinear.weight.float().contiguous()).to(
+            qlinear.weight.dtype
+        )
+
+    qlinear.fold_weight(keep_attrs=keep_attrs)
+
+    assert torch.allclose(qlinear.weight, expected_weight)
+    for quantizer in qlinear.weight_quantizer:
+        assert not quantizer.is_enabled
+        assert not quantizer._enable_pre_quant_scale
+        assert hasattr(quantizer, "_amax") is keep_attrs
+        assert hasattr(quantizer, "_pre_quant_scale") is keep_attrs
+    assert torch.allclose(qlinear(x), out_before)
+
+
 WINT4INT8_CFG = {
     "quant_cfg": [
         {"quantizer_name": "*", "enable": False},

From aeab1f6a34ee937f731d12b9a27b2451493ed093 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 12 Aug 2026 03:56:13 +0000
Subject: [PATCH 02/16] Narrow weight folding fix to TE grouped layers

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 .../quantization/nn/modules/quant_module.py   | 18 ++++------
 .../plugins/transformer_engine.py             |  7 ++--
 .../plugins/test_transformer_engine.py        | 26 ++++++--------
 .../quantization/test_tensor_quant_cpu.py     | 34 -------------------
 4 files changed, 19 insertions(+), 66 deletions(-)

diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py
index 829cd3f10b1..9c9aee478a8 100644
--- a/modelopt/torch/quantization/nn/modules/quant_module.py
+++ b/modelopt/torch/quantization/nn/modules/quant_module.py
@@ -131,7 +131,7 @@ def iter_weights_for_calibration(self):
     @staticmethod
     @torch.no_grad()
     def _fold_weight_quantizer(
-        quantizer: TensorQuantizer | SequentialQuantizer,
+        quantizer: TensorQuantizer,
         weights: Iterable[torch.Tensor],
         keep_attrs: bool = False,
     ):
@@ -145,16 +145,10 @@ def _fold_weight_quantizer(
             weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype))
         quantizer.disable()
         quantizer.disable_rotate()
-        tensor_quantizers = (
-            quantizer if isinstance(quantizer, SequentialQuantizer) else (quantizer,)
-        )
-        for tensor_quantizer in tensor_quantizers:
-            # The pre-quant scale is already baked into the folded weight.
-            tensor_quantizer._enable_pre_quant_scale = False
-            if not keep_attrs:
-                for attr_name in ("_pre_quant_scale", "_amax"):
-                    if hasattr(tensor_quantizer, attr_name):
-                        delattr(tensor_quantizer, attr_name)
+        if not keep_attrs:
+            for attr_name in ("_pre_quant_scale", "_amax"):
+                if hasattr(quantizer, attr_name):
+                    delattr(quantizer, attr_name)
 
     def fold_weight(self, keep_attrs: bool = False):
         """Bake each fake-quant weight quantizer into its weight for faster eval.
@@ -169,7 +163,7 @@ def fold_weight(self, keep_attrs: bool = False):
             attr = getattr(self, name)
             if (
                 name.endswith("weight_quantizer")
-                and isinstance(attr, (TensorQuantizer, SequentialQuantizer))
+                and isinstance(attr, TensorQuantizer)
                 and attr.fake_quant
             ):
                 # Get the corresponding weight name by removing _weight_quantizer suffix
diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py
index c7517eb559c..c2ae02a7f93 100644
--- a/modelopt/torch/quantization/plugins/transformer_engine.py
+++ b/modelopt/torch/quantization/plugins/transformer_engine.py
@@ -247,12 +247,9 @@ def iter_weights_for_calibration(self):
 
     def fold_weight(self, keep_attrs: bool = False):
         """Fold each grouped weight with its corresponding fake-quant quantizer."""
-        quantizer_weights: dict[TensorQuantizer | SequentialQuantizer, list[torch.Tensor]] = {}
+        quantizer_weights: dict[TensorQuantizer, list[torch.Tensor]] = {}
         for weight, quantizer in self.iter_weights_for_calibration():
-            if (
-                isinstance(quantizer, (TensorQuantizer, SequentialQuantizer))
-                and quantizer.fake_quant
-            ):
+            if isinstance(quantizer, TensorQuantizer) and quantizer.fake_quant:
                 quantizer_weights.setdefault(quantizer, []).append(weight)
 
         for quantizer, weights in quantizer_weights.items():
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 adba5443ba4..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 GroupedQuantizer, QuantModule, SequentialQuantizer
+from modelopt.torch.quantization.nn import GroupedQuantizer, QuantModule
 
 
 class TELinear(nn.Module):
@@ -119,26 +119,19 @@ def test_quantize(model_cls, config):
 
 
 @pytest.mark.parametrize("share_weight_quantizer", [False, True])
-@pytest.mark.parametrize("sequential_weight_quantizer", [False, True])
-def test_fold_weight_grouped_linear(share_weight_quantizer, sequential_weight_quantizer):
+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 sequential_weight_quantizer:
-        grouped_linear.weight_quantizer = GroupedQuantizer(
-            *(
-                SequentialQuantizer(quantizer, copy.deepcopy(quantizer))
-                for quantizer in grouped_linear.weight_quantizer
-            )
-        )
     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)
@@ -146,14 +139,17 @@ def test_fold_weight_grouped_linear(share_weight_quantizer, sequential_weight_qu
 
     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
-        tensor_quantizers = (
-            quantizer if isinstance(quantizer, SequentialQuantizer) else (quantizer,)
-        )
-        for tensor_quantizer in tensor_quantizers:
-            assert not hasattr(tensor_quantizer, "_amax")
+        assert not hasattr(quantizer, "_amax")
 
 
 def test_quantize_forward_backward():
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index eba48dfb4c3..ba352ec2162 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -332,40 +332,6 @@ def test_fold_weight_keep_attrs_keeps_amax(monkeypatch):
         unregister_quant_backend(backend_name)
 
 
-@pytest.mark.parametrize("keep_attrs", [False, True])
-def test_fold_weight_supports_sequential_quantizer(keep_attrs):
-    qlinear = QuantModuleRegistry.convert(torch.nn.Linear(4, 3))
-    qlinear.input_quantizer.disable()
-    qlinear.output_quantizer.disable()
-    qlinear.weight_quantizer = SequentialQuantizer(
-        TensorQuantizer(QuantizerAttributeConfig(num_bits=4)),
-        TensorQuantizer(QuantizerAttributeConfig(num_bits=8)),
-    )
-    pre_quant_scales = (
-        torch.tensor([0.5, 1.0, 1.5, 2.0]),
-        torch.tensor([2.0, 1.5, 1.0, 0.5]),
-    )
-    for quantizer, pre_quant_scale in zip(qlinear.weight_quantizer, pre_quant_scales):
-        quantizer.amax = torch.tensor(1.0)
-        quantizer.pre_quant_scale = pre_quant_scale
-    x = torch.randn(2, 4)
-    out_before = qlinear(x)
-    with torch.no_grad():
-        expected_weight = qlinear.weight_quantizer(qlinear.weight.float().contiguous()).to(
-            qlinear.weight.dtype
-        )
-
-    qlinear.fold_weight(keep_attrs=keep_attrs)
-
-    assert torch.allclose(qlinear.weight, expected_weight)
-    for quantizer in qlinear.weight_quantizer:
-        assert not quantizer.is_enabled
-        assert not quantizer._enable_pre_quant_scale
-        assert hasattr(quantizer, "_amax") is keep_attrs
-        assert hasattr(quantizer, "_pre_quant_scale") is keep_attrs
-    assert torch.allclose(qlinear(x), out_before)
-
-
 WINT4INT8_CFG = {
     "quant_cfg": [
         {"quantizer_name": "*", "enable": False},

From 10326e33489ce7ecba217d6e6b3b6dfbf904df78 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 12 Aug 2026 07:31:42 +0000
Subject: [PATCH 03/16] Add reversible weight folding context

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 CHANGELOG.rst                                 |   1 +
 modelopt/torch/quantization/model_quant.py    |  34 ++++
 .../quantization/nn/modules/quant_linear.py   |  13 +-
 .../quantization/nn/modules/quant_module.py   | 170 +++++++++++++++-
 .../torch/quantization/plugins/huggingface.py |   6 +-
 .../plugins/test_transformer_engine.py        |  53 +++++
 .../quantization/plugins/test_huggingface.py  |  57 ++++++
 tests/unit/torch/quantization/test_calib.py   |  38 +++-
 .../quantization/test_tensor_quant_cpu.py     | 187 +++++++++++++++++-
 9 files changed, 546 insertions(+), 13 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 78f8426e5f9..12fcdecb5cb 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -6,6 +6,7 @@ Changelog
 
 **New Features**
 
+- Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference: enabled ``TensorQuantizer`` weights are folded through each quantized module's native ``fold_weight`` implementation for the duration of a context, then weights and quantizer runtime state are restored in place. Retained weight pre-quant scales are made inactive while folded to avoid applying them twice. Parameters with storage tied across owners and ``SequentialQuantizer`` weight containers are not currently folded.
 - 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.
 - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
 - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 966a3643fe3..ac4db4a4835 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
@@ -41,6 +42,7 @@
 from .config import QuantizeAlgoCfgType
 from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg
 from .nn import QuantModule, TensorQuantizer
+from .nn.modules.quant_module import _FoldWeightState, _record_fold_weight_states
 from .utils import is_quantized
 
 __all__ = [
@@ -54,6 +56,7 @@
     "postprocess_amax",
     "print_quant_summary",
     "quantize",
+    "temporarily_fold_weights",
 ]
 
 
@@ -733,6 +736,37 @@ def fold_weight(model: nn.Module, keep_attrs: bool = False):
             module.fold_weight(keep_attrs)
 
 
+@contextmanager
+def temporarily_fold_weights(model: nn.Module):
+    """Temporarily fold enabled fake-quant weights for a frozen inference region.
+
+    Each :class:`QuantModule` performs its normal module-specific ``fold_weight`` operation. The
+    original weights and quantizer runtime state are restored on exit, including after an
+    exception. Parameters are restored in place so optimizer and distributed references remain
+    valid. Disabled and non-fake weight quantizers are left untouched. Parameters whose storage is
+    tied across owners are also left unfolded because mutating them could change a disabled owner.
+    Weight ``SequentialQuantizer`` containers are not currently folded.
+
+    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.
+
+    Example::
+
+        with mtq.temporarily_fold_weights(model):
+            outputs = model(inputs)
+    """
+    states: list[_FoldWeightState] = []
+    try:
+        with _record_fold_weight_states(model, states):
+            fold_weight(model, keep_attrs=True)
+        yield
+    finally:
+        for state in reversed(states):
+            state.restore()
+
+
 @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..0ae83c48b76 100644
--- a/modelopt/torch/quantization/nn/modules/quant_linear.py
+++ b/modelopt/torch/quantization/nn/modules/quant_linear.py
@@ -28,6 +28,7 @@
     QuantLinearConvBase,
     QuantModule,
     QuantModuleRegistry,
+    _is_temporary_weight_fold,
     _LegacyQuantLinearConvBaseMixin,
 )
 from .tensor_quantizer import TensorQuantizer
@@ -68,7 +69,9 @@ class SVDQuantTensorQuantizer(TensorQuantizer):
     @property
     def svdquant_lora_a(self):
         """Lora a weights for svdquant."""
-        if not hasattr(self, "_svdquant_lora_a"):
+        if not getattr(self, "_enable_svdquant_lora", True) or not hasattr(
+            self, "_svdquant_lora_a"
+        ):
             return None
         return self._svdquant_lora_a
 
@@ -92,7 +95,9 @@ def svdquant_lora_a(self, value):
     @property
     def svdquant_lora_b(self):
         """Lora b weights for svdquant."""
-        if not hasattr(self, "_svdquant_lora_b"):
+        if not getattr(self, "_enable_svdquant_lora", True) or not hasattr(
+            self, "_svdquant_lora_b"
+        ):
             return None
         return self._svdquant_lora_b
 
@@ -171,7 +176,8 @@ def fold_weight(self, keep_attrs: bool = False):
             and self.weight_quantizer.fake_quant
         ):
             if (
-                self._not_sequential_quantizers()
+                not _is_temporary_weight_fold()
+                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
             ):
@@ -179,6 +185,7 @@ def fold_weight(self, keep_attrs: bool = False):
                     self.weight
                     + self.weight_quantizer.svdquant_lora_b @ self.weight_quantizer.svdquant_lora_a
                 )
+                self.weight_quantizer._enable_svdquant_lora = False
             if not keep_attrs:
                 _attrs = [
                     "_svdquant_lora_a",
diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py
index 9c9aee478a8..90be50c0334 100644
--- a/modelopt/torch/quantization/nn/modules/quant_module.py
+++ b/modelopt/torch/quantization/nn/modules/quant_module.py
@@ -18,6 +18,8 @@
 import contextlib
 import warnings
 from collections.abc import Iterable
+from contextvars import ContextVar
+from dataclasses import dataclass
 from typing import Any
 
 import torch
@@ -38,6 +40,129 @@
 ]
 
 
+@dataclass
+class _FoldWeightState:
+    """State needed to undo one ``_fold_weight_quantizer`` call."""
+
+    quantizer: TensorQuantizer
+    weights: list[torch.Tensor]
+    original_weights: list[torch.Tensor]
+    weight_keys: set[tuple]
+    disabled: bool
+    rotate: Any
+    enable_pre_quant_scale: bool
+    input_dtype: torch.dtype | None
+
+    @classmethod
+    def capture(cls, quantizer: TensorQuantizer) -> "_FoldWeightState":
+        return cls(
+            quantizer=quantizer,
+            weights=[],
+            original_weights=[],
+            weight_keys=set(),
+            disabled=quantizer._disabled,
+            rotate=quantizer._rotate,
+            enable_pre_quant_scale=quantizer._enable_pre_quant_scale,
+            input_dtype=quantizer._input_dtype,
+        )
+
+    def add_weights(self, weights: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
+        """Snapshot and return weight views not already recorded for this quantizer."""
+        new_weights = []
+        for weight in weights:
+            key = _tensor_view_key(weight)
+            if key in self.weight_keys:
+                continue
+            self.weight_keys.add(key)
+            self.weights.append(weight)
+            self.original_weights.append(weight.detach().clone())
+            new_weights.append(weight)
+        return tuple(new_weights)
+
+    def prepare_quantizer(self):
+        """Restore the quantizer settings needed to fold another associated weight."""
+        self.quantizer._disabled = self.disabled
+        self.quantizer._rotate = self.rotate
+        self.quantizer._enable_pre_quant_scale = self.enable_pre_quant_scale
+        self.quantizer._input_dtype = self.input_dtype
+
+    @torch.no_grad()
+    def restore(self):
+        """Restore weights and quantizer runtime state in place."""
+        for weight, original_weight in zip(self.weights, self.original_weights):
+            weight.data.copy_(original_weight)
+        self.quantizer._disabled = self.disabled
+        self.quantizer._rotate = self.rotate
+        self.quantizer._enable_pre_quant_scale = self.enable_pre_quant_scale
+        self.quantizer._input_dtype = self.input_dtype
+
+
+@dataclass
+class _TemporaryFoldWeightContext:
+    """State shared by module-specific folds in one temporary transaction."""
+
+    states: list[_FoldWeightState]
+    states_by_quantizer: dict[TensorQuantizer, _FoldWeightState]
+    shared_parameter_storages: set[tuple]
+    blocked_quantizers: set[TensorQuantizer]
+
+
+_temporary_fold_weight_context: ContextVar[_TemporaryFoldWeightContext | None] = ContextVar(
+    "temporary_fold_weight_context", default=None
+)
+
+
+def _tensor_storage_key(tensor: torch.Tensor) -> tuple | None:
+    """Return a device-qualified storage key when the tensor exposes local storage."""
+    if tensor.is_meta or tensor.numel() == 0:
+        return None
+    try:
+        return ("storage", tensor.device, tensor.untyped_storage().data_ptr())
+    except (RuntimeError, NotImplementedError):
+        try:
+            local = tensor.to_local()
+            return ("storage", local.device, local.untyped_storage().data_ptr())
+        except (AttributeError, RuntimeError, NotImplementedError):
+            return ("object", id(tensor))
+
+
+def _tensor_view_key(tensor: torch.Tensor) -> tuple:
+    """Identify an exact tensor view while allowing distinct slices of one fused weight."""
+    return (
+        _tensor_storage_key(tensor),
+        tensor.storage_offset(),
+        tuple(tensor.shape),
+        tuple(tensor.stride()),
+    )
+
+
+def _shared_parameter_storages(model: nn.Module) -> set[tuple]:
+    """Find storage referenced by more than one parameter attribute in the model."""
+    storage_counts: dict[tuple, int] = {}
+    for module in model.modules():
+        for parameter in module._parameters.values():
+            if parameter is None or (key := _tensor_storage_key(parameter)) is None:
+                continue
+            storage_counts[key] = storage_counts.get(key, 0) + 1
+    return {key for key, count in storage_counts.items() if count > 1}
+
+
+@contextlib.contextmanager
+def _record_fold_weight_states(model: nn.Module, states: list[_FoldWeightState]):
+    """Record fold undo state without changing the public ``fold_weight`` contract."""
+    context = _TemporaryFoldWeightContext(states, {}, _shared_parameter_storages(model), set())
+    token = _temporary_fold_weight_context.set(context)
+    try:
+        yield
+    finally:
+        _temporary_fold_weight_context.reset(token)
+
+
+def _is_temporary_weight_fold() -> bool:
+    """Return whether folding is running inside ``temporarily_fold_weights``."""
+    return _temporary_fold_weight_context.get() is not None
+
+
 class QuantModule(DynamicModule):
     """A base class for quantized modules.
 
@@ -141,14 +266,45 @@ 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.
         """
+        weights = tuple(weights)
+        context = _temporary_fold_weight_context.get()
+        if context is not None:
+            if quantizer in context.blocked_quantizers:
+                return
+            if any(
+                _tensor_storage_key(weight) in context.shared_parameter_storages
+                for weight in weights
+            ):
+                context.blocked_quantizers.add(quantizer)
+                if state := context.states_by_quantizer.get(quantizer):
+                    state.restore()
+                return
+            if not quantizer.fake_quant:
+                return
+            state = context.states_by_quantizer.get(quantizer)
+            if state is None:
+                if not quantizer.is_enabled:
+                    return
+                state = _FoldWeightState.capture(quantizer)
+                context.states.append(state)
+                context.states_by_quantizer[quantizer] = state
+            weights = state.add_weights(weights)
+            if not weights:
+                return
+            state.prepare_quantizer()
+
         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.
@@ -156,7 +312,8 @@ def fold_weight(self, keep_attrs: bool = False):
         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``.
+        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 +330,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/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py b/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py
index 13ecc941c76..535c414e06c 100644
--- a/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py
+++ b/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py
@@ -152,6 +152,59 @@ def test_fold_weight_grouped_linear(share_weight_quantizer):
         assert not hasattr(quantizer, "_amax")
 
 
+@pytest.mark.parametrize("share_weight_quantizer", [False, True])
+def test_temporarily_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()]
+    original_weights = [weight.detach().clone() for weight in weights]
+    for weight, quantizer in zip(weights, quantizers):
+        if quantizer.pre_quant_scale is None:
+            quantizer.pre_quant_scale = torch.full(
+                (1, weight.shape[-1]),
+                0.5,
+                dtype=weight.dtype,
+                device=weight.device,
+            )
+    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)
+        ]
+
+    with mtq.temporarily_fold_weights(model):
+        with torch.no_grad():
+            output_folded = model(calib_data[0])
+        for weight, expected_weight, quantizer in zip(weights, expected_weights, quantizers):
+            assert torch.allclose(weight, expected_weight)
+            assert not quantizer.is_enabled
+            assert hasattr(quantizer, "_amax")
+            assert hasattr(quantizer, "_pre_quant_scale")
+            assert quantizer.pre_quant_scale is None
+
+    with torch.no_grad():
+        output_restored = model(calib_data[0])
+    if isinstance(output_before, tuple):
+        output_before = output_before[0]
+        output_folded = output_folded[0]
+        output_restored = output_restored[0]
+    assert torch.allclose(output_folded, output_before)
+    assert torch.allclose(output_restored, output_before)
+
+    for weight, original_weight, quantizer in zip(weights, original_weights, quantizers):
+        assert torch.equal(weight, original_weight)
+        assert quantizer.is_enabled
+        assert quantizer.pre_quant_scale is not None
+
+
 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..fd237e5743b 100644
--- a/tests/unit/torch/quantization/plugins/test_huggingface.py
+++ b/tests/unit/torch/quantization/plugins/test_huggingface.py
@@ -34,7 +34,9 @@
 import modelopt.torch.quantization as mtq
 from modelopt.recipe.loader import load_recipe
 from modelopt.torch.quantization.nn import QuantLinear, QuantModuleRegistry, TensorQuantizer
+from modelopt.torch.quantization.nn.modules.quant_module import _shared_parameter_storages
 from modelopt.torch.quantization.plugins.huggingface import (
+    _QuantHFParallelLinear,
     _TransposedExpertsCalibMixin,
     get_homogeneous_hf_decoder_layers,
     is_homogeneous_hf_model,
@@ -77,6 +79,61 @@ 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_temporary_fold_detects_tied_dtensor_local_storage():
+    class FakeDTensor:
+        def __init__(self, local):
+            self._local = local
+            self.device = local.device
+            self.is_meta = False
+
+        def numel(self):
+            return self._local.numel()
+
+        def untyped_storage(self):
+            raise RuntimeError("DTensor has no wrapper storage")
+
+        def to_local(self):
+            return self._local
+
+    local = torch.ones(2, 2)
+    first = FakeDTensor(local)
+    second = FakeDTensor(local)
+
+    class TiedModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self._parameters["first"] = first
+            self._parameters["second"] = second
+
+    assert len(_shared_parameter_storages(TiedModel())) == 1
+
+
 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..2155231e347 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,39 @@ 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)
+    expected_folded_weights = []
+    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
+        expected_folded_weights.append(
+            (
+                module,
+                module.weight_quantizer(module.weight.float().contiguous()).to(module.weight.dtype)
+                + module.weight_quantizer.svdquant_lora_b @ module.weight_quantizer.svdquant_lora_a,
+            )
+        )
+
+    mtq.fold_weight(model, keep_attrs=True)
+    for module, expected_weight in expected_folded_weights:
+        assert torch.allclose(module.weight, expected_weight)
+        assert not module.weight_quantizer.is_enabled
+        assert hasattr(module.weight_quantizer, "_svdquant_lora_a")
+        assert hasattr(module.weight_quantizer, "_svdquant_lora_b")
+        assert module.weight_quantizer.svdquant_lora_a is None
+        assert module.weight_quantizer.svdquant_lora_b is 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..14f9718884f 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -317,21 +317,206 @@ 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):
+            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_skips_disabled_and_none_weights():
+    disabled = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
+    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)
+        assert not disabled.weight_quantizer.is_enabled
+        assert tied_output.weight_quantizer.is_enabled
+
+    assert torch.equal(disabled.weight, original_weight)
+    assert not disabled.weight_quantizer.is_enabled
+    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_skips_non_fake_override():
+    qlinear = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
+    quantizer = qlinear.weight_quantizer
+    quantizer._fake_quant = False
+    original_weight = qlinear.weight.detach().clone()
+    fold_calls = []
+
+    def unconditional_fold(keep_attrs=False):
+        fold_calls.append(keep_attrs)
+        return qlinear._fold_weight_quantizer(quantizer, (qlinear.weight,), keep_attrs)
+
+    qlinear.fold_weight = unconditional_fold
+    with mtq.temporarily_fold_weights(qlinear):
+        assert torch.equal(qlinear.weight, original_weight)
+        assert quantizer.is_enabled
+
+    assert fold_calls == [True]
+    assert torch.equal(qlinear.weight, original_weight)
+    assert quantizer.is_enabled
+
+
+def test_temporarily_fold_weights_folds_all_weights_with_shared_quantizer(monkeypatch):
+    calls = []
+    backend_name = "test_temporary_fold_shared_quantizer_backend"
+    first = _make_qlinear_with_backend(monkeypatch, calls, backend_name)
+    second = QuantModuleRegistry.convert(torch.nn.Linear(4, 3))
+    second.input_quantizer.disable()
+    second.output_quantizer.disable()
+    second.weight_quantizer = first.weight_quantizer
+    quantizer = first.weight_quantizer
+    original_weights = [first.weight.detach().clone(), second.weight.detach().clone()]
+    inputs = torch.randn(2, 4)
+    outputs_before = [first(inputs), second(inputs)]
+    try:
+        with mtq.temporarily_fold_weights(torch.nn.ModuleList([first, second])):
+            assert not torch.equal(first.weight, original_weights[0])
+            assert not torch.equal(second.weight, original_weights[1])
+            assert not quantizer.is_enabled
+            assert torch.allclose(first(inputs), outputs_before[0])
+            assert torch.allclose(second(inputs), outputs_before[1])
+
+        assert quantizer.is_enabled
+        assert torch.equal(first.weight, original_weights[0])
+        assert torch.equal(second.weight, original_weights[1])
+    finally:
+        unregister_quant_backend(backend_name)
+
+
+def test_temporarily_fold_weights_skips_tied_parameter_storage():
+    first = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
+    second = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
+    second.weight = first.weight
+    second.weight_quantizer.disable()
+    original_weight = first.weight.detach().clone()
+    inputs = torch.randn(2, 4)
+    outputs_before = [first(inputs), second(inputs)]
+
+    with mtq.temporarily_fold_weights(torch.nn.ModuleList([first, second])):
+        assert torch.equal(first.weight, original_weight)
+        assert first.weight_quantizer.is_enabled
+        assert not second.weight_quantizer.is_enabled
+        assert torch.equal(first(inputs), outputs_before[0])
+        assert torch.equal(second(inputs), outputs_before[1])
+
+    assert torch.equal(first.weight, original_weight)
+    assert first.weight_quantizer.is_enabled
+    assert not second.weight_quantizer.is_enabled
+
+
+def test_temporarily_fold_weights_blocks_quantizer_shared_with_tied_weight():
+    independent = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
+    tied_first = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
+    tied_second = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
+    tied_second.weight = tied_first.weight
+    tied_first.weight_quantizer = independent.weight_quantizer
+    original_weights = [module.weight.detach().clone() for module in (independent, tied_first)]
+    inputs = torch.randn(2, 4)
+    outputs_before = [module(inputs) for module in (independent, tied_first, tied_second)]
+    model = torch.nn.ModuleList([independent, tied_first, tied_second])
+
+    with mtq.temporarily_fold_weights(model):
+        assert independent.weight_quantizer.is_enabled
+        assert torch.equal(independent.weight, original_weights[0])
+        assert torch.equal(tied_first.weight, original_weights[1])
+        for module, output_before in zip(model, outputs_before):
+            assert torch.equal(module(inputs), output_before)
+
+    assert independent.weight_quantizer.is_enabled
+    assert torch.equal(independent.weight, original_weights[0])
+    assert torch.equal(tied_first.weight, original_weights[1])
+
+
 WINT4INT8_CFG = {
     "quant_cfg": [
         {"quantizer_name": "*", "enable": False},

From f62e7abc5ba81062a3d6f6600af84678f2178a9f Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 12 Aug 2026 08:06:50 +0000
Subject: [PATCH 04/16] Simplify reversible weight folding

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 CHANGELOG.rst                                 |   2 +-
 modelopt/torch/quantization/model_quant.py    |  37 +++-
 .../quantization/nn/modules/quant_module.py   | 187 ++++++------------
 .../quantization/plugins/test_huggingface.py  |  30 ---
 .../quantization/test_tensor_quant_cpu.py     |  73 ++-----
 5 files changed, 105 insertions(+), 224 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 12fcdecb5cb..2ea9eff787d 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -6,7 +6,7 @@ Changelog
 
 **New Features**
 
-- Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference: enabled ``TensorQuantizer`` weights are folded through each quantized module's native ``fold_weight`` implementation for the duration of a context, then weights and quantizer runtime state are restored in place. Retained weight pre-quant scales are made inactive while folded to avoid applying them twice. Parameters with storage tied across owners and ``SequentialQuantizer`` weight containers are not currently folded.
+- Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference: ``TensorQuantizer`` weights are folded through each quantized module's native ``fold_weight`` implementation for the duration of a context, then weights and quantizer runtime state are restored in place. Retained weight pre-quant scales are made inactive while folded to avoid applying them twice. Weight ``SequentialQuantizer`` containers, quantizers shared across separate folding calls, and weights tied across quantized modules are not currently supported.
 - 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.
 - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
 - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index ac4db4a4835..a56c8a0c707 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -42,7 +42,7 @@
 from .config import QuantizeAlgoCfgType
 from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg
 from .nn import QuantModule, TensorQuantizer
-from .nn.modules.quant_module import _FoldWeightState, _record_fold_weight_states
+from .nn.modules.quant_module import _record_folded_weights
 from .utils import is_quantized
 
 __all__ = [
@@ -738,14 +738,13 @@ def fold_weight(model: nn.Module, keep_attrs: bool = False):
 
 @contextmanager
 def temporarily_fold_weights(model: nn.Module):
-    """Temporarily fold enabled fake-quant weights for a frozen inference region.
+    """Temporarily fold fake-quant weights for a frozen inference region.
 
     Each :class:`QuantModule` performs its normal module-specific ``fold_weight`` operation. The
     original weights and quantizer runtime state are restored on exit, including after an
     exception. Parameters are restored in place so optimizer and distributed references remain
-    valid. Disabled and non-fake weight quantizers are left untouched. Parameters whose storage is
-    tied across owners are also left unfolded because mutating them could change a disabled owner.
-    Weight ``SequentialQuantizer`` containers are not currently folded.
+    valid. Weight ``SequentialQuantizer`` containers, quantizers shared across separate folding
+    calls, and weights tied across quantized modules are not currently supported.
 
     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
@@ -757,14 +756,34 @@ def temporarily_fold_weights(model: nn.Module):
         with mtq.temporarily_fold_weights(model):
             outputs = model(inputs)
     """
-    states: list[_FoldWeightState] = []
+    weight_states: list[tuple[torch.Tensor, torch.Tensor]] = []
+    state_attrs = (
+        "_disabled",
+        "_rotate",
+        "_enable_pre_quant_scale",
+        "_input_dtype",
+    )
+    missing = object()
+    quantizer_states = [
+        (module, {name: getattr(module, name, missing) for name in state_attrs})
+        for module in model.modules()
+        if isinstance(module, TensorQuantizer)
+    ]
     try:
-        with _record_fold_weight_states(model, states):
+        with _record_folded_weights(model, weight_states):
             fold_weight(model, keep_attrs=True)
         yield
     finally:
-        for state in reversed(states):
-            state.restore()
+        with torch.no_grad():
+            for weight, original_weight in reversed(weight_states):
+                weight.copy_(original_weight)
+        for quantizer, state in quantizer_states:
+            for name, value in state.items():
+                if value is missing:
+                    if hasattr(quantizer, name):
+                        delattr(quantizer, name)
+                else:
+                    setattr(quantizer, name, value)
 
 
 @torch.no_grad()
diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py
index 90be50c0334..2b765d88fbf 100644
--- a/modelopt/torch/quantization/nn/modules/quant_module.py
+++ b/modelopt/torch/quantization/nn/modules/quant_module.py
@@ -19,7 +19,6 @@
 import warnings
 from collections.abc import Iterable
 from contextvars import ContextVar
-from dataclasses import dataclass
 from typing import Any
 
 import torch
@@ -40,127 +39,57 @@
 ]
 
 
-@dataclass
-class _FoldWeightState:
-    """State needed to undo one ``_fold_weight_quantizer`` call."""
+_fold_weight_recorder: ContextVar[
+    tuple[list[tuple[torch.Tensor, torch.Tensor]], set[TensorQuantizer], set[tuple], set[tuple]]
+    | None
+] = ContextVar("fold_weight_recorder", default=None)
 
-    quantizer: TensorQuantizer
-    weights: list[torch.Tensor]
-    original_weights: list[torch.Tensor]
-    weight_keys: set[tuple]
-    disabled: bool
-    rotate: Any
-    enable_pre_quant_scale: bool
-    input_dtype: torch.dtype | None
 
-    @classmethod
-    def capture(cls, quantizer: TensorQuantizer) -> "_FoldWeightState":
-        return cls(
-            quantizer=quantizer,
-            weights=[],
-            original_weights=[],
-            weight_keys=set(),
-            disabled=quantizer._disabled,
-            rotate=quantizer._rotate,
-            enable_pre_quant_scale=quantizer._enable_pre_quant_scale,
-            input_dtype=quantizer._input_dtype,
-        )
-
-    def add_weights(self, weights: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
-        """Snapshot and return weight views not already recorded for this quantizer."""
-        new_weights = []
-        for weight in weights:
-            key = _tensor_view_key(weight)
-            if key in self.weight_keys:
-                continue
-            self.weight_keys.add(key)
-            self.weights.append(weight)
-            self.original_weights.append(weight.detach().clone())
-            new_weights.append(weight)
-        return tuple(new_weights)
-
-    def prepare_quantizer(self):
-        """Restore the quantizer settings needed to fold another associated weight."""
-        self.quantizer._disabled = self.disabled
-        self.quantizer._rotate = self.rotate
-        self.quantizer._enable_pre_quant_scale = self.enable_pre_quant_scale
-        self.quantizer._input_dtype = self.input_dtype
-
-    @torch.no_grad()
-    def restore(self):
-        """Restore weights and quantizer runtime state in place."""
-        for weight, original_weight in zip(self.weights, self.original_weights):
-            weight.data.copy_(original_weight)
-        self.quantizer._disabled = self.disabled
-        self.quantizer._rotate = self.rotate
-        self.quantizer._enable_pre_quant_scale = self.enable_pre_quant_scale
-        self.quantizer._input_dtype = self.input_dtype
-
-
-@dataclass
-class _TemporaryFoldWeightContext:
-    """State shared by module-specific folds in one temporary transaction."""
-
-    states: list[_FoldWeightState]
-    states_by_quantizer: dict[TensorQuantizer, _FoldWeightState]
-    shared_parameter_storages: set[tuple]
-    blocked_quantizers: set[TensorQuantizer]
-
-
-_temporary_fold_weight_context: ContextVar[_TemporaryFoldWeightContext | None] = ContextVar(
-    "temporary_fold_weight_context", default=None
-)
-
-
-def _tensor_storage_key(tensor: torch.Tensor) -> tuple | None:
-    """Return a device-qualified storage key when the tensor exposes local storage."""
-    if tensor.is_meta or tensor.numel() == 0:
-        return None
+def _tensor_view_key(tensor: torch.Tensor) -> tuple:
+    """Identify an exact tensor view, including aliases backed by distinct tensor objects."""
     try:
-        return ("storage", tensor.device, tensor.untyped_storage().data_ptr())
-    except (RuntimeError, NotImplementedError):
+        return (
+            tensor.device,
+            tensor.untyped_storage().data_ptr(),
+            tensor.storage_offset(),
+            tuple(tensor.shape),
+            tuple(tensor.stride()),
+        )
+    except (AttributeError, RuntimeError, NotImplementedError):
         try:
-            local = tensor.to_local()
-            return ("storage", local.device, local.untyped_storage().data_ptr())
+            local_tensor = tensor.to_local()
+            if local_tensor is not tensor:
+                return _tensor_view_key(local_tensor)
         except (AttributeError, RuntimeError, NotImplementedError):
-            return ("object", id(tensor))
+            pass
+        return (id(tensor),)
 
 
-def _tensor_view_key(tensor: torch.Tensor) -> tuple:
-    """Identify an exact tensor view while allowing distinct slices of one fused weight."""
-    return (
-        _tensor_storage_key(tensor),
-        tensor.storage_offset(),
-        tuple(tensor.shape),
-        tuple(tensor.stride()),
-    )
-
-
-def _shared_parameter_storages(model: nn.Module) -> set[tuple]:
-    """Find storage referenced by more than one parameter attribute in the model."""
-    storage_counts: dict[tuple, int] = {}
-    for module in model.modules():
-        for parameter in module._parameters.values():
-            if parameter is None or (key := _tensor_storage_key(parameter)) is None:
-                continue
-            storage_counts[key] = storage_counts.get(key, 0) + 1
-    return {key for key, count in storage_counts.items() if count > 1}
+def _tensor_storage_key(tensor: torch.Tensor) -> tuple:
+    view_key = _tensor_view_key(tensor)
+    return view_key[:2] if len(view_key) > 1 else view_key
+
+
+def _tied_parameter_storages(model: nn.Module) -> set[tuple]:
+    counts: dict[tuple, int] = {}
+    for _, parameter in model.named_parameters(remove_duplicate=False):
+        key = _tensor_storage_key(parameter)
+        counts[key] = counts.get(key, 0) + 1
+    return {key for key, count in counts.items() if count > 1}
 
 
 @contextlib.contextmanager
-def _record_fold_weight_states(model: nn.Module, states: list[_FoldWeightState]):
-    """Record fold undo state without changing the public ``fold_weight`` contract."""
-    context = _TemporaryFoldWeightContext(states, {}, _shared_parameter_storages(model), set())
-    token = _temporary_fold_weight_context.set(context)
+def _record_folded_weights(model: nn.Module, states: list[tuple[torch.Tensor, torch.Tensor]]):
+    """Record weights immediately before module-specific folding mutates them."""
+    token = _fold_weight_recorder.set((states, set(), set(), _tied_parameter_storages(model)))
     try:
         yield
     finally:
-        _temporary_fold_weight_context.reset(token)
+        _fold_weight_recorder.reset(token)
 
 
 def _is_temporary_weight_fold() -> bool:
-    """Return whether folding is running inside ``temporarily_fold_weights``."""
-    return _temporary_fold_weight_context.get() is not None
+    return _fold_weight_recorder.get() is not None
 
 
 class QuantModule(DynamicModule):
@@ -267,31 +196,31 @@ def _fold_weight_quantizer(
         view while disabling and dropping its calibration attrs exactly once.
         """
         weights = tuple(weights)
-        context = _temporary_fold_weight_context.get()
-        if context is not None:
-            if quantizer in context.blocked_quantizers:
-                return
-            if any(
-                _tensor_storage_key(weight) in context.shared_parameter_storages
-                for weight in weights
-            ):
-                context.blocked_quantizers.add(quantizer)
-                if state := context.states_by_quantizer.get(quantizer):
-                    state.restore()
-                return
+        if recorder := _fold_weight_recorder.get():
             if not quantizer.fake_quant:
                 return
-            state = context.states_by_quantizer.get(quantizer)
-            if state is None:
-                if not quantizer.is_enabled:
-                    return
-                state = _FoldWeightState.capture(quantizer)
-                context.states.append(state)
-                context.states_by_quantizer[quantizer] = state
-            weights = state.add_weights(weights)
-            if not weights:
-                return
-            state.prepare_quantizer()
+            states, folded_quantizers, folded_weight_keys, tied_weight_storages = recorder
+            if quantizer in folded_quantizers:
+                raise RuntimeError(
+                    "temporarily_fold_weights does not support a weight quantizer shared across "
+                    "multiple fold_weight calls"
+                )
+            unique_weights = {}
+            for weight in weights:
+                unique_weights.setdefault(_tensor_view_key(weight), weight)
+            weights = tuple(unique_weights.values())
+            weight_keys = set(unique_weights)
+            if any(_tensor_storage_key(weight) in tied_weight_storages for weight in weights):
+                raise RuntimeError(
+                    "temporarily_fold_weights does not support a weight tied across modules"
+                )
+            if not weight_keys.isdisjoint(folded_weight_keys):
+                raise RuntimeError(
+                    "temporarily_fold_weights does not support a weight shared across quantized modules"
+                )
+            folded_quantizers.add(quantizer)
+            folded_weight_keys.update(weight_keys)
+            states.extend((weight, weight.detach().clone()) for weight in weights)
 
         for weight in weights:
             weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype))
diff --git a/tests/unit/torch/quantization/plugins/test_huggingface.py b/tests/unit/torch/quantization/plugins/test_huggingface.py
index fd237e5743b..0da6cfbb1fc 100644
--- a/tests/unit/torch/quantization/plugins/test_huggingface.py
+++ b/tests/unit/torch/quantization/plugins/test_huggingface.py
@@ -34,7 +34,6 @@
 import modelopt.torch.quantization as mtq
 from modelopt.recipe.loader import load_recipe
 from modelopt.torch.quantization.nn import QuantLinear, QuantModuleRegistry, TensorQuantizer
-from modelopt.torch.quantization.nn.modules.quant_module import _shared_parameter_storages
 from modelopt.torch.quantization.plugins.huggingface import (
     _QuantHFParallelLinear,
     _TransposedExpertsCalibMixin,
@@ -105,35 +104,6 @@ class ParallelLinear:
     assert linear.weight is original_weight
 
 
-def test_temporary_fold_detects_tied_dtensor_local_storage():
-    class FakeDTensor:
-        def __init__(self, local):
-            self._local = local
-            self.device = local.device
-            self.is_meta = False
-
-        def numel(self):
-            return self._local.numel()
-
-        def untyped_storage(self):
-            raise RuntimeError("DTensor has no wrapper storage")
-
-        def to_local(self):
-            return self._local
-
-    local = torch.ones(2, 2)
-    first = FakeDTensor(local)
-    second = FakeDTensor(local)
-
-    class TiedModel(nn.Module):
-        def __init__(self):
-            super().__init__()
-            self._parameters["first"] = first
-            self._parameters["second"] = second
-
-    assert len(_shared_parameter_storages(TiedModel())) == 1
-
-
 def test_convert_conv1d():
     set_seed()
     assert transformers.pytorch_utils.Conv1D in QuantModuleRegistry
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index 14f9718884f..7ef81c63b7a 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -430,23 +430,16 @@ def test_temporarily_fold_weights_skips_non_fake_override():
     quantizer = qlinear.weight_quantizer
     quantizer._fake_quant = False
     original_weight = qlinear.weight.detach().clone()
-    fold_calls = []
-
-    def unconditional_fold(keep_attrs=False):
-        fold_calls.append(keep_attrs)
-        return qlinear._fold_weight_quantizer(quantizer, (qlinear.weight,), keep_attrs)
+    qlinear.fold_weight = lambda keep_attrs=False: qlinear._fold_weight_quantizer(
+        quantizer, (qlinear.weight,), keep_attrs
+    )
 
-    qlinear.fold_weight = unconditional_fold
     with mtq.temporarily_fold_weights(qlinear):
         assert torch.equal(qlinear.weight, original_weight)
         assert quantizer.is_enabled
 
-    assert fold_calls == [True]
-    assert torch.equal(qlinear.weight, original_weight)
-    assert quantizer.is_enabled
-
 
-def test_temporarily_fold_weights_folds_all_weights_with_shared_quantizer(monkeypatch):
+def test_temporarily_fold_weights_rejects_quantizer_shared_across_modules(monkeypatch):
     calls = []
     backend_name = "test_temporary_fold_shared_quantizer_backend"
     first = _make_qlinear_with_backend(monkeypatch, calls, backend_name)
@@ -456,15 +449,12 @@ def test_temporarily_fold_weights_folds_all_weights_with_shared_quantizer(monkey
     second.weight_quantizer = first.weight_quantizer
     quantizer = first.weight_quantizer
     original_weights = [first.weight.detach().clone(), second.weight.detach().clone()]
-    inputs = torch.randn(2, 4)
-    outputs_before = [first(inputs), second(inputs)]
     try:
-        with mtq.temporarily_fold_weights(torch.nn.ModuleList([first, second])):
-            assert not torch.equal(first.weight, original_weights[0])
-            assert not torch.equal(second.weight, original_weights[1])
-            assert not quantizer.is_enabled
-            assert torch.allclose(first(inputs), outputs_before[0])
-            assert torch.allclose(second(inputs), outputs_before[1])
+        with (
+            pytest.raises(RuntimeError, match="weight quantizer shared"),
+            mtq.temporarily_fold_weights(torch.nn.ModuleList([first, second])),
+        ):
+            pass
 
         assert quantizer.is_enabled
         assert torch.equal(first.weight, original_weights[0])
@@ -473,48 +463,21 @@ def test_temporarily_fold_weights_folds_all_weights_with_shared_quantizer(monkey
         unregister_quant_backend(backend_name)
 
 
-def test_temporarily_fold_weights_skips_tied_parameter_storage():
+def test_temporarily_fold_weights_rejects_tied_parameter():
     first = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
     second = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
-    second.weight = first.weight
-    second.weight_quantizer.disable()
+    second.weight = torch.nn.Parameter(first.weight.detach().T)
+    second.weight_quantizer._fake_quant = False
     original_weight = first.weight.detach().clone()
-    inputs = torch.randn(2, 4)
-    outputs_before = [first(inputs), second(inputs)]
-
-    with mtq.temporarily_fold_weights(torch.nn.ModuleList([first, second])):
-        assert torch.equal(first.weight, original_weight)
-        assert first.weight_quantizer.is_enabled
-        assert not second.weight_quantizer.is_enabled
-        assert torch.equal(first(inputs), outputs_before[0])
-        assert torch.equal(second(inputs), outputs_before[1])
+    with (
+        pytest.raises(RuntimeError, match="weight tied across modules"),
+        mtq.temporarily_fold_weights(torch.nn.ModuleList([first, second])),
+    ):
+        pass
 
     assert torch.equal(first.weight, original_weight)
     assert first.weight_quantizer.is_enabled
-    assert not second.weight_quantizer.is_enabled
-
-
-def test_temporarily_fold_weights_blocks_quantizer_shared_with_tied_weight():
-    independent = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
-    tied_first = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
-    tied_second = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
-    tied_second.weight = tied_first.weight
-    tied_first.weight_quantizer = independent.weight_quantizer
-    original_weights = [module.weight.detach().clone() for module in (independent, tied_first)]
-    inputs = torch.randn(2, 4)
-    outputs_before = [module(inputs) for module in (independent, tied_first, tied_second)]
-    model = torch.nn.ModuleList([independent, tied_first, tied_second])
-
-    with mtq.temporarily_fold_weights(model):
-        assert independent.weight_quantizer.is_enabled
-        assert torch.equal(independent.weight, original_weights[0])
-        assert torch.equal(tied_first.weight, original_weights[1])
-        for module, output_before in zip(model, outputs_before):
-            assert torch.equal(module(inputs), output_before)
-
-    assert independent.weight_quantizer.is_enabled
-    assert torch.equal(independent.weight, original_weights[0])
-    assert torch.equal(tied_first.weight, original_weights[1])
+    assert second.weight_quantizer.is_enabled
 
 
 WINT4INT8_CFG = {

From 9a5eed712960eac4bbd2082530fdd0e73fd5a7f7 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 12 Aug 2026 08:11:42 +0000
Subject: [PATCH 05/16] Keep temporary folding scope narrow

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 CHANGELOG.rst                                   |  2 +-
 modelopt/torch/quantization/model_quant.py      |  2 +-
 .../quantization/nn/modules/quant_linear.py     |  9 ++-------
 .../quantization/nn/modules/quant_module.py     |  2 +-
 tests/unit/torch/quantization/test_calib.py     | 17 -----------------
 .../torch/quantization/test_tensor_quant_cpu.py |  4 ++--
 6 files changed, 7 insertions(+), 29 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 2ea9eff787d..ad12f30240c 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -6,7 +6,7 @@ Changelog
 
 **New Features**
 
-- Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference: ``TensorQuantizer`` weights are folded through each quantized module's native ``fold_weight`` implementation for the duration of a context, then weights and quantizer runtime state are restored in place. Retained weight pre-quant scales are made inactive while folded to avoid applying them twice. Weight ``SequentialQuantizer`` containers, quantizers shared across separate folding calls, and weights tied across quantized modules are not currently supported.
+- Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference: ``TensorQuantizer`` weights are folded through each quantized module's native ``fold_weight`` implementation for the duration of a context, then weights and quantizer runtime state are restored in place. Retained weight pre-quant scales are made inactive while folded to avoid applying them twice. Weight ``SequentialQuantizer`` containers, quantizers shared across separate folding calls, and parameters sharing storage across modules are not currently supported.
 - 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.
 - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
 - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index a56c8a0c707..f2d04921204 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -744,7 +744,7 @@ def temporarily_fold_weights(model: nn.Module):
     original weights and quantizer runtime state are restored on exit, including after an
     exception. Parameters are restored in place so optimizer and distributed references remain
     valid. Weight ``SequentialQuantizer`` containers, quantizers shared across separate folding
-    calls, and weights tied across quantized modules are not currently supported.
+    calls, and parameters sharing storage across modules are not currently supported.
 
     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
diff --git a/modelopt/torch/quantization/nn/modules/quant_linear.py b/modelopt/torch/quantization/nn/modules/quant_linear.py
index 0ae83c48b76..457b898a110 100644
--- a/modelopt/torch/quantization/nn/modules/quant_linear.py
+++ b/modelopt/torch/quantization/nn/modules/quant_linear.py
@@ -69,9 +69,7 @@ class SVDQuantTensorQuantizer(TensorQuantizer):
     @property
     def svdquant_lora_a(self):
         """Lora a weights for svdquant."""
-        if not getattr(self, "_enable_svdquant_lora", True) or not hasattr(
-            self, "_svdquant_lora_a"
-        ):
+        if not hasattr(self, "_svdquant_lora_a"):
             return None
         return self._svdquant_lora_a
 
@@ -95,9 +93,7 @@ def svdquant_lora_a(self, value):
     @property
     def svdquant_lora_b(self):
         """Lora b weights for svdquant."""
-        if not getattr(self, "_enable_svdquant_lora", True) or not hasattr(
-            self, "_svdquant_lora_b"
-        ):
+        if not hasattr(self, "_svdquant_lora_b"):
             return None
         return self._svdquant_lora_b
 
@@ -185,7 +181,6 @@ def fold_weight(self, keep_attrs: bool = False):
                     self.weight
                     + self.weight_quantizer.svdquant_lora_b @ self.weight_quantizer.svdquant_lora_a
                 )
-                self.weight_quantizer._enable_svdquant_lora = False
             if not keep_attrs:
                 _attrs = [
                     "_svdquant_lora_a",
diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py
index 2b765d88fbf..51e1a280d08 100644
--- a/modelopt/torch/quantization/nn/modules/quant_module.py
+++ b/modelopt/torch/quantization/nn/modules/quant_module.py
@@ -212,7 +212,7 @@ def _fold_weight_quantizer(
             weight_keys = set(unique_weights)
             if any(_tensor_storage_key(weight) in tied_weight_storages for weight in weights):
                 raise RuntimeError(
-                    "temporarily_fold_weights does not support a weight tied across modules"
+                    "temporarily_fold_weights does not support parameters sharing storage across modules"
                 )
             if not weight_keys.isdisjoint(folded_weight_keys):
                 raise RuntimeError(
diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py
index 2155231e347..aa31d0b66d6 100644
--- a/tests/unit/torch/quantization/test_calib.py
+++ b/tests/unit/torch/quantization/test_calib.py
@@ -434,27 +434,10 @@ def test_svdquant_lora_weights():
 
     assert torch.allclose(output_folded, output_before)
     assert torch.allclose(output_restored, output_before)
-    expected_folded_weights = []
     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
-        expected_folded_weights.append(
-            (
-                module,
-                module.weight_quantizer(module.weight.float().contiguous()).to(module.weight.dtype)
-                + module.weight_quantizer.svdquant_lora_b @ module.weight_quantizer.svdquant_lora_a,
-            )
-        )
-
-    mtq.fold_weight(model, keep_attrs=True)
-    for module, expected_weight in expected_folded_weights:
-        assert torch.allclose(module.weight, expected_weight)
-        assert not module.weight_quantizer.is_enabled
-        assert hasattr(module.weight_quantizer, "_svdquant_lora_a")
-        assert hasattr(module.weight_quantizer, "_svdquant_lora_b")
-        assert module.weight_quantizer.svdquant_lora_a is None
-        assert module.weight_quantizer.svdquant_lora_b is None
 
 
 def test_layerwise_calibrate_support_gate():
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index 7ef81c63b7a..c0e67c4da41 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -463,14 +463,14 @@ def test_temporarily_fold_weights_rejects_quantizer_shared_across_modules(monkey
         unregister_quant_backend(backend_name)
 
 
-def test_temporarily_fold_weights_rejects_tied_parameter():
+def test_temporarily_fold_weights_rejects_shared_parameter_storage():
     first = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
     second = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
     second.weight = torch.nn.Parameter(first.weight.detach().T)
     second.weight_quantizer._fake_quant = False
     original_weight = first.weight.detach().clone()
     with (
-        pytest.raises(RuntimeError, match="weight tied across modules"),
+        pytest.raises(RuntimeError, match="parameters sharing storage across modules"),
         mtq.temporarily_fold_weights(torch.nn.ModuleList([first, second])),
     ):
         pass

From 9243b023e51ba3c629062ca01b07be229072de6a Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 12 Aug 2026 08:47:39 +0000
Subject: [PATCH 06/16] Simplify temporary weight snapshots

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 CHANGELOG.rst                                 |   2 +-
 modelopt/torch/quantization/model_quant.py    | 112 +++++++++++++-----
 .../quantization/nn/modules/quant_module.py   |  75 ++----------
 .../torch/quantization/plugins/huggingface.py |   6 +
 .../quantization/test_tensor_quant_cpu.py     |  48 ++------
 5 files changed, 107 insertions(+), 136 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index ad12f30240c..31250b3a4a7 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -6,7 +6,7 @@ Changelog
 
 **New Features**
 
-- Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference: ``TensorQuantizer`` weights are folded through each quantized module's native ``fold_weight`` implementation for the duration of a context, then weights and quantizer runtime state are restored in place. Retained weight pre-quant scales are made inactive while folded to avoid applying them twice. Weight ``SequentialQuantizer`` containers, quantizers shared across separate folding calls, and parameters sharing storage across modules are not currently supported.
+- Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference: enabled fake-quant weights exposed by each quantized module's ``iter_weights_for_calibration()`` 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. Set ``snapshot_weights=False`` to leave the model folded without allocating snapshots or providing rollback on folding errors. Retained weight pre-quant scales are inactive while folded to avoid applying them twice. Reusing one weight or weight quantizer across modules is rejected before folding.
 - 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.
 - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
 - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index f2d04921204..89551e65261 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -42,7 +42,7 @@
 from .config import QuantizeAlgoCfgType
 from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg
 from .nn import QuantModule, TensorQuantizer
-from .nn.modules.quant_module import _record_folded_weights
+from .nn.modules.quant_module import _temporary_weight_fold_context
 from .utils import is_quantized
 
 __all__ = [
@@ -737,14 +737,18 @@ def fold_weight(model: nn.Module, keep_attrs: bool = False):
 
 
 @contextmanager
-def temporarily_fold_weights(model: nn.Module):
+def temporarily_fold_weights(
+    model: nn.Module,
+    snapshot_device: torch.device | str | None = None,
+    snapshot_weights: bool = True,
+):
     """Temporarily fold fake-quant weights for a frozen inference region.
 
-    Each :class:`QuantModule` performs its normal module-specific ``fold_weight`` operation. The
-    original weights and quantizer runtime state are restored on exit, including after an
-    exception. Parameters are restored in place so optimizer and distributed references remain
-    valid. Weight ``SequentialQuantizer`` containers, quantizers shared across separate folding
-    calls, and parameters sharing storage across modules are not currently supported.
+    Each :class:`QuantModule` performs its normal module-specific ``fold_weight`` operation. When
+    ``snapshot_weights=True``, enabled fake-quant weights and their quantizer runtime state are
+    restored on exit, including after an exception. Weights are restored in place so optimizer and
+    distributed references remain valid. Reusing one weight or weight quantizer across modules is
+    rejected before folding.
 
     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
@@ -753,37 +757,81 @@ def temporarily_fold_weights(model: nn.Module):
 
     Example::
 
-        with mtq.temporarily_fold_weights(model):
+        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.
+        snapshot_weights: Snapshot and restore folded weights and quantizer state. If ``False``,
+            the model remains folded after the context exits and folding errors are not rolled back.
     """
-    weight_states: list[tuple[torch.Tensor, torch.Tensor]] = []
-    state_attrs = (
-        "_disabled",
-        "_rotate",
-        "_enable_pre_quant_scale",
-        "_input_dtype",
-    )
-    missing = object()
-    quantizer_states = [
-        (module, {name: getattr(module, name, missing) for name in state_attrs})
-        for module in model.modules()
-        if isinstance(module, TensorQuantizer)
-    ]
+    active_pairs = []
+    quantizer_owners = {}
+    weight_owners = {}
+    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
+            local_weight = weight.to_local() if hasattr(weight, "to_local") else weight
+            try:
+                storage_id = (local_weight.device, local_weight.untyped_storage().data_ptr())
+            except (RuntimeError, NotImplementedError):
+                storage_id = id(local_weight)
+            weight_owner = weight_owners.setdefault(storage_id, module)
+            if weight_owner is not module:
+                raise ValueError("A weight is shared across quantized modules")
+            if not (
+                isinstance(quantizer, TensorQuantizer)
+                and quantizer.fake_quant
+                and quantizer.is_enabled
+            ):
+                continue
+            quantizer_owner = quantizer_owners.setdefault(id(quantizer), module)
+            if quantizer_owner is not module:
+                raise ValueError("A weight quantizer is shared across quantized modules")
+            active_pairs.append((local_weight, quantizer))
+
+    weight_snapshots = {}
+    quantizer_states = {}
+    if snapshot_weights:
+        for weight, quantizer in active_pairs:
+            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:
-        with _record_folded_weights(model, weight_states):
+        with _temporary_weight_fold_context():
             fold_weight(model, keep_attrs=True)
         yield
     finally:
-        with torch.no_grad():
-            for weight, original_weight in reversed(weight_states):
-                weight.copy_(original_weight)
-        for quantizer, state in quantizer_states:
-            for name, value in state.items():
-                if value is missing:
-                    if hasattr(quantizer, name):
-                        delattr(quantizer, name)
-                else:
-                    setattr(quantizer, name, value)
+        if snapshot_weights:
+            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()
diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py
index 51e1a280d08..7cca8a6c1d8 100644
--- a/modelopt/torch/quantization/nn/modules/quant_module.py
+++ b/modelopt/torch/quantization/nn/modules/quant_module.py
@@ -39,57 +39,20 @@
 ]
 
 
-_fold_weight_recorder: ContextVar[
-    tuple[list[tuple[torch.Tensor, torch.Tensor]], set[TensorQuantizer], set[tuple], set[tuple]]
-    | None
-] = ContextVar("fold_weight_recorder", default=None)
-
-
-def _tensor_view_key(tensor: torch.Tensor) -> tuple:
-    """Identify an exact tensor view, including aliases backed by distinct tensor objects."""
-    try:
-        return (
-            tensor.device,
-            tensor.untyped_storage().data_ptr(),
-            tensor.storage_offset(),
-            tuple(tensor.shape),
-            tuple(tensor.stride()),
-        )
-    except (AttributeError, RuntimeError, NotImplementedError):
-        try:
-            local_tensor = tensor.to_local()
-            if local_tensor is not tensor:
-                return _tensor_view_key(local_tensor)
-        except (AttributeError, RuntimeError, NotImplementedError):
-            pass
-        return (id(tensor),)
-
-
-def _tensor_storage_key(tensor: torch.Tensor) -> tuple:
-    view_key = _tensor_view_key(tensor)
-    return view_key[:2] if len(view_key) > 1 else view_key
-
-
-def _tied_parameter_storages(model: nn.Module) -> set[tuple]:
-    counts: dict[tuple, int] = {}
-    for _, parameter in model.named_parameters(remove_duplicate=False):
-        key = _tensor_storage_key(parameter)
-        counts[key] = counts.get(key, 0) + 1
-    return {key for key, count in counts.items() if count > 1}
+_temporary_weight_fold_active = ContextVar("temporary_weight_fold_active", default=False)
 
 
 @contextlib.contextmanager
-def _record_folded_weights(model: nn.Module, states: list[tuple[torch.Tensor, torch.Tensor]]):
-    """Record weights immediately before module-specific folding mutates them."""
-    token = _fold_weight_recorder.set((states, set(), set(), _tied_parameter_storages(model)))
+def _temporary_weight_fold_context():
+    token = _temporary_weight_fold_active.set(True)
     try:
         yield
     finally:
-        _fold_weight_recorder.reset(token)
+        _temporary_weight_fold_active.reset(token)
 
 
 def _is_temporary_weight_fold() -> bool:
-    return _fold_weight_recorder.get() is not None
+    return _temporary_weight_fold_active.get()
 
 
 class QuantModule(DynamicModule):
@@ -195,32 +158,8 @@ 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.
         """
-        weights = tuple(weights)
-        if recorder := _fold_weight_recorder.get():
-            if not quantizer.fake_quant:
-                return
-            states, folded_quantizers, folded_weight_keys, tied_weight_storages = recorder
-            if quantizer in folded_quantizers:
-                raise RuntimeError(
-                    "temporarily_fold_weights does not support a weight quantizer shared across "
-                    "multiple fold_weight calls"
-                )
-            unique_weights = {}
-            for weight in weights:
-                unique_weights.setdefault(_tensor_view_key(weight), weight)
-            weights = tuple(unique_weights.values())
-            weight_keys = set(unique_weights)
-            if any(_tensor_storage_key(weight) in tied_weight_storages for weight in weights):
-                raise RuntimeError(
-                    "temporarily_fold_weights does not support parameters sharing storage across modules"
-                )
-            if not weight_keys.isdisjoint(folded_weight_keys):
-                raise RuntimeError(
-                    "temporarily_fold_weights does not support a weight shared across quantized modules"
-                )
-            folded_quantizers.add(quantizer)
-            folded_weight_keys.update(weight_keys)
-            states.extend((weight, weight.detach().clone()) for weight in weights)
+        if _is_temporary_weight_fold() and (not quantizer.fake_quant or not quantizer.is_enabled):
+            return
 
         for weight in weights:
             weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype))
diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 4acb4d30dfa..2d7205c7bef 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -619,6 +619,12 @@ def iter_weights_for_calibration(self):
             weight = getattr(self, weight_name)
             yield weight.transpose(-1, -2), getattr(self, f"{weight_name}_weight_quantizer")
 
+    def fold_weight(self, keep_attrs: bool = False):
+        """Fold expert weights in the same transposed orientation used by the forward."""
+        for weight, quantizer in self.iter_weights_for_calibration():
+            if isinstance(quantizer, TensorQuantizer) and quantizer.fake_quant:
+                QuantModule._fold_weight_quantizer(quantizer, (weight,), keep_attrs)
+
 
 class _QuantSparseSequentialMoe(QuantModule):
     """Quantization wrapper for HuggingFace sparse MoE blocks.
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index c0e67c4da41..104a165a119 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -359,7 +359,10 @@ def test_temporarily_fold_weights_restores_after_exception(monkeypatch):
         inputs = torch.randn(2, 4)
         output_before = qlinear(inputs)
 
-        with pytest.raises(RuntimeError, match="test error"), mtq.temporarily_fold_weights(qlinear):
+        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
@@ -439,47 +442,22 @@ def test_temporarily_fold_weights_skips_non_fake_override():
         assert quantizer.is_enabled
 
 
-def test_temporarily_fold_weights_rejects_quantizer_shared_across_modules(monkeypatch):
+def test_temporarily_fold_weights_without_snapshot_leaves_model_folded(monkeypatch):
     calls = []
-    backend_name = "test_temporary_fold_shared_quantizer_backend"
-    first = _make_qlinear_with_backend(monkeypatch, calls, backend_name)
-    second = QuantModuleRegistry.convert(torch.nn.Linear(4, 3))
-    second.input_quantizer.disable()
-    second.output_quantizer.disable()
-    second.weight_quantizer = first.weight_quantizer
-    quantizer = first.weight_quantizer
-    original_weights = [first.weight.detach().clone(), second.weight.detach().clone()]
+    backend_name = "test_temporary_fold_without_snapshot"
+    qlinear = _make_qlinear_with_backend(monkeypatch, calls, backend_name)
+    original_weight = qlinear.weight.detach().clone()
     try:
-        with (
-            pytest.raises(RuntimeError, match="weight quantizer shared"),
-            mtq.temporarily_fold_weights(torch.nn.ModuleList([first, second])),
-        ):
-            pass
+        with mtq.temporarily_fold_weights(qlinear, snapshot_weights=False):
+            assert not torch.equal(qlinear.weight, original_weight)
+            assert not qlinear.weight_quantizer.is_enabled
 
-        assert quantizer.is_enabled
-        assert torch.equal(first.weight, original_weights[0])
-        assert torch.equal(second.weight, original_weights[1])
+        assert not torch.equal(qlinear.weight, original_weight)
+        assert not qlinear.weight_quantizer.is_enabled
     finally:
         unregister_quant_backend(backend_name)
 
 
-def test_temporarily_fold_weights_rejects_shared_parameter_storage():
-    first = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
-    second = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
-    second.weight = torch.nn.Parameter(first.weight.detach().T)
-    second.weight_quantizer._fake_quant = False
-    original_weight = first.weight.detach().clone()
-    with (
-        pytest.raises(RuntimeError, match="parameters sharing storage across modules"),
-        mtq.temporarily_fold_weights(torch.nn.ModuleList([first, second])),
-    ):
-        pass
-
-    assert torch.equal(first.weight, original_weight)
-    assert first.weight_quantizer.is_enabled
-    assert second.weight_quantizer.is_enabled
-
-
 WINT4INT8_CFG = {
     "quant_cfg": [
         {"quantizer_name": "*", "enable": False},

From fb09ca9930abe7ca4bf556b18c50e9a2c2b65a8d Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 12 Aug 2026 09:24:08 +0000
Subject: [PATCH 07/16] Simplify temporary weight folding

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 modelopt/torch/quantization/model_quant.py    | 22 +++-----------
 .../quantization/nn/modules/quant_linear.py   | 12 ++++----
 .../quantization/nn/modules/quant_module.py   | 29 ++++---------------
 3 files changed, 17 insertions(+), 46 deletions(-)

diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 89551e65261..26140e1b961 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -42,7 +42,6 @@
 from .config import QuantizeAlgoCfgType
 from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg
 from .nn import QuantModule, TensorQuantizer
-from .nn.modules.quant_module import _temporary_weight_fold_context
 from .utils import is_quantized
 
 __all__ = [
@@ -747,13 +746,13 @@ def temporarily_fold_weights(
     Each :class:`QuantModule` performs its normal module-specific ``fold_weight`` operation. When
     ``snapshot_weights=True``, enabled fake-quant weights and their quantizer runtime state are
     restored on exit, including after an exception. Weights are restored in place so optimizer and
-    distributed references remain valid. Reusing one weight or weight quantizer across modules is
-    rejected before folding.
+    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.
+    value is already baked into the temporary weight. Sharing a weight or weight quantizer across
+    multiple :class:`QuantModule` instances is not supported.
 
     Example::
 
@@ -768,8 +767,6 @@ def temporarily_fold_weights(
             the model remains folded after the context exits and folding errors are not rolled back.
     """
     active_pairs = []
-    quantizer_owners = {}
-    weight_owners = {}
     for module in model.modules():
         if not isinstance(module, QuantModule):
             continue
@@ -777,22 +774,12 @@ def temporarily_fold_weights(
             if not isinstance(weight, torch.Tensor):
                 continue
             local_weight = weight.to_local() if hasattr(weight, "to_local") else weight
-            try:
-                storage_id = (local_weight.device, local_weight.untyped_storage().data_ptr())
-            except (RuntimeError, NotImplementedError):
-                storage_id = id(local_weight)
-            weight_owner = weight_owners.setdefault(storage_id, module)
-            if weight_owner is not module:
-                raise ValueError("A weight is shared across quantized modules")
             if not (
                 isinstance(quantizer, TensorQuantizer)
                 and quantizer.fake_quant
                 and quantizer.is_enabled
             ):
                 continue
-            quantizer_owner = quantizer_owners.setdefault(id(quantizer), module)
-            if quantizer_owner is not module:
-                raise ValueError("A weight quantizer is shared across quantized modules")
             active_pairs.append((local_weight, quantizer))
 
     weight_snapshots = {}
@@ -817,8 +804,7 @@ def temporarily_fold_weights(
                 ),
             )
     try:
-        with _temporary_weight_fold_context():
-            fold_weight(model, keep_attrs=True)
+        fold_weight(model, keep_attrs=True)
         yield
     finally:
         if snapshot_weights:
diff --git a/modelopt/torch/quantization/nn/modules/quant_linear.py b/modelopt/torch/quantization/nn/modules/quant_linear.py
index 457b898a110..b4b911a11bc 100644
--- a/modelopt/torch/quantization/nn/modules/quant_linear.py
+++ b/modelopt/torch/quantization/nn/modules/quant_linear.py
@@ -28,7 +28,6 @@
     QuantLinearConvBase,
     QuantModule,
     QuantModuleRegistry,
-    _is_temporary_weight_fold,
     _LegacyQuantLinearConvBaseMixin,
 )
 from .tensor_quantizer import TensorQuantizer
@@ -165,14 +164,17 @@ 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_quantizer, TensorQuantizer)
             and self.weight_quantizer.fake_quant
-        ):
+            and self.weight_quantizer.is_enabled
+        )
+        super().fold_weight(keep_attrs)
+        if should_fold:
             if (
-                not _is_temporary_weight_fold()
+                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 7cca8a6c1d8..bee90f85ab2 100644
--- a/modelopt/torch/quantization/nn/modules/quant_module.py
+++ b/modelopt/torch/quantization/nn/modules/quant_module.py
@@ -18,7 +18,6 @@
 import contextlib
 import warnings
 from collections.abc import Iterable
-from contextvars import ContextVar
 from typing import Any
 
 import torch
@@ -39,22 +38,6 @@
 ]
 
 
-_temporary_weight_fold_active = ContextVar("temporary_weight_fold_active", default=False)
-
-
-@contextlib.contextmanager
-def _temporary_weight_fold_context():
-    token = _temporary_weight_fold_active.set(True)
-    try:
-        yield
-    finally:
-        _temporary_weight_fold_active.reset(token)
-
-
-def _is_temporary_weight_fold() -> bool:
-    return _temporary_weight_fold_active.get()
-
-
 class QuantModule(DynamicModule):
     """A base class for quantized modules.
 
@@ -158,7 +141,7 @@ 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 _is_temporary_weight_fold() and (not quantizer.fake_quant or not quantizer.is_enabled):
+        if not quantizer.fake_quant or not quantizer.is_enabled:
             return
 
         for weight in weights:
@@ -177,11 +160,11 @@ def _fold_weight_quantizer(
     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``. A retained pre-quant scale remains stored but is made
-        inactive because it has already been applied to the folded weight.
+        Every enabled fake-quant weight quantizer is folded. 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):

From 14208316e0e1f11889923c82ec63fce6dff1d4f2 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 12 Aug 2026 09:46:36 +0000
Subject: [PATCH 08/16] Always restore temporarily folded weights

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 CHANGELOG.rst                                 |  2 +-
 modelopt/torch/quantization/model_quant.py    | 75 +++++++++----------
 .../quantization/nn/modules/quant_linear.py   |  1 +
 .../quantization/test_tensor_quant_cpu.py     | 19 ++---
 4 files changed, 45 insertions(+), 52 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 31250b3a4a7..95f95cc1a5d 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -6,7 +6,7 @@ Changelog
 
 **New Features**
 
-- Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference: enabled fake-quant weights exposed by each quantized module's ``iter_weights_for_calibration()`` 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. Set ``snapshot_weights=False`` to leave the model folded without allocating snapshots or providing rollback on folding errors. Retained weight pre-quant scales are inactive while folded to avoid applying them twice. Reusing one weight or weight quantizer across modules is rejected before folding.
+- Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference: enabled fake-quant weights exposed by each quantized module's ``iter_weights_for_calibration()`` 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. 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.
 - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
 - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 26140e1b961..21cfd961492 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -41,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__ = [
@@ -739,20 +739,19 @@ def fold_weight(model: nn.Module, keep_attrs: bool = False):
 def temporarily_fold_weights(
     model: nn.Module,
     snapshot_device: torch.device | str | None = None,
-    snapshot_weights: bool = True,
 ):
     """Temporarily fold fake-quant weights for a frozen inference region.
 
-    Each :class:`QuantModule` performs its normal module-specific ``fold_weight`` operation. When
-    ``snapshot_weights=True``, enabled fake-quant weights and their quantizer runtime state are
-    restored on exit, including after an exception. Weights are restored in place so optimizer and
-    distributed references remain valid.
+    Each :class:`QuantModule` performs its normal module-specific ``fold_weight`` operation. Enabled
+    fake-quant weights and their quantizer runtime state 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 is not supported.
+    multiple :class:`QuantModule` instances and using :class:`SequentialQuantizer` are not
+    supported.
 
     Example::
 
@@ -763,8 +762,6 @@ def temporarily_fold_weights(
         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.
-        snapshot_weights: Snapshot and restore folded weights and quantizer state. If ``False``,
-            the model remains folded after the context exits and folding errors are not rolled back.
     """
     active_pairs = []
     for module in model.modules():
@@ -773,6 +770,10 @@ def temporarily_fold_weights(
         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"
+                )
             local_weight = weight.to_local() if hasattr(weight, "to_local") else weight
             if not (
                 isinstance(quantizer, TensorQuantizer)
@@ -784,40 +785,38 @@ def temporarily_fold_weights(
 
     weight_snapshots = {}
     quantizer_states = {}
-    if snapshot_weights:
-        for weight, quantizer in active_pairs:
-            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,
-                ),
+    for weight, quantizer in active_pairs:
+        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:
-        if snapshot_weights:
-            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
+        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()
diff --git a/modelopt/torch/quantization/nn/modules/quant_linear.py b/modelopt/torch/quantization/nn/modules/quant_linear.py
index b4b911a11bc..009a3c8da29 100644
--- a/modelopt/torch/quantization/nn/modules/quant_linear.py
+++ b/modelopt/torch/quantization/nn/modules/quant_linear.py
@@ -167,6 +167,7 @@ def fold_weight(self, keep_attrs: bool = False):
         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
             and self.weight_quantizer.is_enabled
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index 104a165a119..b5fbfbe7ed6 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -442,20 +442,13 @@ def test_temporarily_fold_weights_skips_non_fake_override():
         assert quantizer.is_enabled
 
 
-def test_temporarily_fold_weights_without_snapshot_leaves_model_folded(monkeypatch):
-    calls = []
-    backend_name = "test_temporary_fold_without_snapshot"
-    qlinear = _make_qlinear_with_backend(monkeypatch, calls, backend_name)
-    original_weight = qlinear.weight.detach().clone()
-    try:
-        with mtq.temporarily_fold_weights(qlinear, snapshot_weights=False):
-            assert not torch.equal(qlinear.weight, original_weight)
-            assert not qlinear.weight_quantizer.is_enabled
+def test_temporarily_fold_weights_rejects_sequential_quantizer():
+    qlinear = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
+    qlinear.weight_quantizer = SequentialQuantizer(TensorQuantizer(), TensorQuantizer())
 
-        assert not torch.equal(qlinear.weight, original_weight)
-        assert not qlinear.weight_quantizer.is_enabled
-    finally:
-        unregister_quant_backend(backend_name)
+    with pytest.raises(NotImplementedError, match="does not support SequentialQuantizer"):
+        with mtq.temporarily_fold_weights(qlinear):
+            pass
 
 
 WINT4INT8_CFG = {

From d09d4e500f99c8c86c9725d827dcafb2ff370810 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 12 Aug 2026 10:03:09 +0000
Subject: [PATCH 09/16] Fold active transforms on disabled quantizers

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 CHANGELOG.rst                                 |  2 +-
 modelopt/torch/quantization/model_quant.py    | 42 ++++++++++---------
 .../quantization/nn/modules/quant_linear.py   |  1 -
 .../quantization/nn/modules/quant_module.py   | 22 ++++++----
 .../quantization/test_tensor_quant_cpu.py     | 18 ++++++--
 5 files changed, 52 insertions(+), 33 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 95f95cc1a5d..a10c612c7b5 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -6,7 +6,7 @@ Changelog
 
 **New Features**
 
-- Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference: enabled fake-quant weights exposed by each quantized module's ``iter_weights_for_calibration()`` 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 ``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. 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.
 - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
 - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 21cfd961492..315b163da88 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -742,9 +742,10 @@ def temporarily_fold_weights(
 ):
     """Temporarily fold fake-quant weights for a frozen inference region.
 
-    Each :class:`QuantModule` performs its normal module-specific ``fold_weight`` operation. Enabled
-    fake-quant weights and their quantizer runtime state are restored on exit, including after an
-    exception. Weights are restored in place so optimizer and distributed references remain valid.
+    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
@@ -763,7 +764,7 @@ def temporarily_fold_weights(
         snapshot_device: Device used to store parameter snapshots. ``None`` keeps each snapshot
             on the parameter's device; ``"cpu"`` avoids the additional accelerator memory.
     """
-    active_pairs = []
+    fold_pairs = []
     for module in model.modules():
         if not isinstance(module, QuantModule):
             continue
@@ -774,26 +775,27 @@ def temporarily_fold_weights(
                 raise NotImplementedError(
                     "temporarily_fold_weights does not support SequentialQuantizer"
                 )
-            local_weight = weight.to_local() if hasattr(weight, "to_local") else weight
-            if not (
-                isinstance(quantizer, TensorQuantizer)
-                and quantizer.fake_quant
-                and quantizer.is_enabled
-            ):
+            if not isinstance(quantizer, TensorQuantizer) or not quantizer.fake_quant:
                 continue
-            active_pairs.append((local_weight, quantizer))
+            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 active_pairs:
-        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),
-            )
+    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,
             (
diff --git a/modelopt/torch/quantization/nn/modules/quant_linear.py b/modelopt/torch/quantization/nn/modules/quant_linear.py
index 009a3c8da29..d0569ff5b71 100644
--- a/modelopt/torch/quantization/nn/modules/quant_linear.py
+++ b/modelopt/torch/quantization/nn/modules/quant_linear.py
@@ -170,7 +170,6 @@ def fold_weight(self, keep_attrs: bool = False):
             and isinstance(self.weight, torch.Tensor)
             and isinstance(self.weight_quantizer, TensorQuantizer)
             and self.weight_quantizer.fake_quant
-            and self.weight_quantizer.is_enabled
         )
         super().fold_weight(keep_attrs)
         if should_fold:
diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py
index bee90f85ab2..cebfe5af29d 100644
--- a/modelopt/torch/quantization/nn/modules/quant_module.py
+++ b/modelopt/torch/quantization/nn/modules/quant_module.py
@@ -141,11 +141,16 @@ 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 or not quantizer.is_enabled:
+        if not quantizer.fake_quant:
             return
 
-        for weight in weights:
-            weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype))
+        if (
+            quantizer.is_enabled
+            or quantizer.pre_quant_scale is not None
+            or quantizer.rotate_is_enabled
+        ):
+            for weight in weights:
+                weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype))
         quantizer.disable()
         quantizer.disable_rotate()
         if hasattr(quantizer, "_pre_quant_scale"):
@@ -160,11 +165,12 @@ def _fold_weight_quantizer(
     def fold_weight(self, keep_attrs: bool = False):
         """Bake each fake-quant weight quantizer into its weight for faster eval.
 
-        Every enabled fake-quant weight quantizer is folded. 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.
+        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):
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index b5fbfbe7ed6..b9487ca9077 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -383,22 +383,34 @@ def test_temporarily_fold_weights_restores_after_exception(monkeypatch):
         unregister_quant_backend(backend_name)
 
 
-def test_temporarily_fold_weights_skips_disabled_and_none_weights():
+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)
+    inactive = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False).double())
+    inactive.weight_quantizer.disable()
+    with torch.no_grad():
+        inactive.weight[0, 0] = 1.000000000001
+    inactive_weight = inactive.weight.detach().clone()
+
+    with mtq.temporarily_fold_weights(torch.nn.ModuleList([disabled, tied_output, inactive])):
+        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(inactive.weight, inactive_weight)
 
     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
+    assert torch.equal(inactive.weight, inactive_weight)
 
 
 def test_temporarily_fold_weights_restores_when_folding_fails(monkeypatch):

From 286e2e236cb596808df16ee2ae86bbcf6ce94ca1 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 12 Aug 2026 10:08:59 +0000
Subject: [PATCH 10/16] Keep weight folding behavior unchanged

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 modelopt/torch/quantization/nn/modules/quant_module.py |  9 ++-------
 tests/unit/torch/quantization/test_tensor_quant_cpu.py | 10 +---------
 2 files changed, 3 insertions(+), 16 deletions(-)

diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py
index cebfe5af29d..83ee8940d19 100644
--- a/modelopt/torch/quantization/nn/modules/quant_module.py
+++ b/modelopt/torch/quantization/nn/modules/quant_module.py
@@ -144,13 +144,8 @@ def _fold_weight_quantizer(
         if not quantizer.fake_quant:
             return
 
-        if (
-            quantizer.is_enabled
-            or quantizer.pre_quant_scale is not None
-            or quantizer.rotate_is_enabled
-        ):
-            for weight in weights:
-                weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype))
+        for weight in weights:
+            weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype))
         quantizer.disable()
         quantizer.disable_rotate()
         if hasattr(quantizer, "_pre_quant_scale"):
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index b9487ca9077..3edb8d636bc 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -393,24 +393,16 @@ def test_temporarily_fold_weights_handles_disabled_transform_and_none_weight():
     tied_output = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
     tied_output.weight = None
 
-    inactive = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False).double())
-    inactive.weight_quantizer.disable()
-    with torch.no_grad():
-        inactive.weight[0, 0] = 1.000000000001
-    inactive_weight = inactive.weight.detach().clone()
-
-    with mtq.temporarily_fold_weights(torch.nn.ModuleList([disabled, tied_output, inactive])):
+    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(inactive.weight, inactive_weight)
 
     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
-    assert torch.equal(inactive.weight, inactive_weight)
 
 
 def test_temporarily_fold_weights_restores_when_folding_fails(monkeypatch):

From 0f56f9869b65bf9b95ff18012a1560ccd95c1fad Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 12 Aug 2026 10:13:41 +0000
Subject: [PATCH 11/16] Trim redundant weight folding tests

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 .../plugins/test_transformer_engine.py        | 53 -------------------
 .../quantization/test_tensor_quant_cpu.py     | 14 -----
 2 files changed, 67 deletions(-)

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 535c414e06c..13ecc941c76 100644
--- a/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py
+++ b/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py
@@ -152,59 +152,6 @@ def test_fold_weight_grouped_linear(share_weight_quantizer):
         assert not hasattr(quantizer, "_amax")
 
 
-@pytest.mark.parametrize("share_weight_quantizer", [False, True])
-def test_temporarily_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()]
-    original_weights = [weight.detach().clone() for weight in weights]
-    for weight, quantizer in zip(weights, quantizers):
-        if quantizer.pre_quant_scale is None:
-            quantizer.pre_quant_scale = torch.full(
-                (1, weight.shape[-1]),
-                0.5,
-                dtype=weight.dtype,
-                device=weight.device,
-            )
-    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)
-        ]
-
-    with mtq.temporarily_fold_weights(model):
-        with torch.no_grad():
-            output_folded = model(calib_data[0])
-        for weight, expected_weight, quantizer in zip(weights, expected_weights, quantizers):
-            assert torch.allclose(weight, expected_weight)
-            assert not quantizer.is_enabled
-            assert hasattr(quantizer, "_amax")
-            assert hasattr(quantizer, "_pre_quant_scale")
-            assert quantizer.pre_quant_scale is None
-
-    with torch.no_grad():
-        output_restored = model(calib_data[0])
-    if isinstance(output_before, tuple):
-        output_before = output_before[0]
-        output_folded = output_folded[0]
-        output_restored = output_restored[0]
-    assert torch.allclose(output_folded, output_before)
-    assert torch.allclose(output_restored, output_before)
-
-    for weight, original_weight, quantizer in zip(weights, original_weights, quantizers):
-        assert torch.equal(weight, original_weight)
-        assert quantizer.is_enabled
-        assert quantizer.pre_quant_scale is not None
-
-
 def test_quantize_forward_backward():
     set_seed()
     model = TELinear().cuda()
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index 3edb8d636bc..248065a1f91 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -432,20 +432,6 @@ def failing_backend(inputs, _quantizer):
         unregister_quant_backend(second_backend)
 
 
-def test_temporarily_fold_weights_skips_non_fake_override():
-    qlinear = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
-    quantizer = qlinear.weight_quantizer
-    quantizer._fake_quant = False
-    original_weight = qlinear.weight.detach().clone()
-    qlinear.fold_weight = lambda keep_attrs=False: qlinear._fold_weight_quantizer(
-        quantizer, (qlinear.weight,), keep_attrs
-    )
-
-    with mtq.temporarily_fold_weights(qlinear):
-        assert torch.equal(qlinear.weight, original_weight)
-        assert quantizer.is_enabled
-
-
 def test_temporarily_fold_weights_rejects_sequential_quantizer():
     qlinear = QuantModuleRegistry.convert(torch.nn.Linear(4, 3, bias=False))
     qlinear.weight_quantizer = SequentialQuantizer(TensorQuantizer(), TensorQuantizer())

From ebd407e2ce9548cacb9b9fcde32bd13957853294 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 12 Aug 2026 10:22:30 +0000
Subject: [PATCH 12/16] Remove unused transposed expert folding

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 modelopt/torch/quantization/plugins/huggingface.py | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 2d7205c7bef..7e2a1f7ab3d 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -619,13 +619,6 @@ def iter_weights_for_calibration(self):
             weight = getattr(self, weight_name)
             yield weight.transpose(-1, -2), getattr(self, f"{weight_name}_weight_quantizer")
 
-    def fold_weight(self, keep_attrs: bool = False):
-        """Fold expert weights in the same transposed orientation used by the forward."""
-        for weight, quantizer in self.iter_weights_for_calibration():
-            if isinstance(quantizer, TensorQuantizer) and quantizer.fake_quant:
-                QuantModule._fold_weight_quantizer(quantizer, (weight,), keep_attrs)
-
-
 class _QuantSparseSequentialMoe(QuantModule):
     """Quantization wrapper for HuggingFace sparse MoE blocks.
 

From cd2f7a48482d61ff284f5508a6c82ffc0eb915a1 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 12 Aug 2026 10:35:16 +0000
Subject: [PATCH 13/16] Fix code quality checks

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 .claude/skills/benchmark-model-kernels                 | 1 +
 modelopt/torch/quantization/plugins/huggingface.py     | 1 +
 tests/unit/torch/quantization/test_tensor_quant_cpu.py | 8 +++++---
 3 files changed, 7 insertions(+), 3 deletions(-)
 create mode 120000 .claude/skills/benchmark-model-kernels

diff --git a/.claude/skills/benchmark-model-kernels b/.claude/skills/benchmark-model-kernels
new file mode 120000
index 00000000000..1bfd1fefe85
--- /dev/null
+++ b/.claude/skills/benchmark-model-kernels
@@ -0,0 +1 @@
+../../.agents/skills/benchmark-model-kernels
\ No newline at end of file
diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py
index 7e2a1f7ab3d..4acb4d30dfa 100644
--- a/modelopt/torch/quantization/plugins/huggingface.py
+++ b/modelopt/torch/quantization/plugins/huggingface.py
@@ -619,6 +619,7 @@ def iter_weights_for_calibration(self):
             weight = getattr(self, weight_name)
             yield weight.transpose(-1, -2), getattr(self, f"{weight_name}_weight_quantizer")
 
+
 class _QuantSparseSequentialMoe(QuantModule):
     """Quantization wrapper for HuggingFace sparse MoE blocks.
 
diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
index 248065a1f91..f8f3540a790 100644
--- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py
+++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py
@@ -436,9 +436,11 @@ 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"):
-        with mtq.temporarily_fold_weights(qlinear):
-            pass
+    with (
+        pytest.raises(NotImplementedError, match="does not support SequentialQuantizer"),
+        mtq.temporarily_fold_weights(qlinear),
+    ):
+        pass
 
 
 WINT4INT8_CFG = {

From 4e02ab2d4cee8e3521034ae7b0ec539f0e93e8f5 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Wed, 12 Aug 2026 10:35:41 +0000
Subject: [PATCH 14/16] Remove unrelated generated symlink

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 .claude/skills/benchmark-model-kernels | 1 -
 1 file changed, 1 deletion(-)
 delete mode 120000 .claude/skills/benchmark-model-kernels

diff --git a/.claude/skills/benchmark-model-kernels b/.claude/skills/benchmark-model-kernels
deleted file mode 120000
index 1bfd1fefe85..00000000000
--- a/.claude/skills/benchmark-model-kernels
+++ /dev/null
@@ -1 +0,0 @@
-../../.agents/skills/benchmark-model-kernels
\ No newline at end of file

From 0c623e07c7f4352c321d8abd308d88ad39503ef6 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Thu, 13 Aug 2026 01:48:12 +0000
Subject: [PATCH 15/16] Expose vLLM fused MoE weights for calibration

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 modelopt/torch/quantization/plugins/vllm.py | 5 +++++
 1 file changed, 5 insertions(+)

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

From f6eedf870aa33b2ba2c50c85ce6e7a1b6d7257e8 Mon Sep 17 00:00:00 2001
From: Meng Xin <mxin@nvidia.com>
Date: Thu, 13 Aug 2026 03:22:58 +0000
Subject: [PATCH 16/16] Preserve changelog entry placement

Signed-off-by: Meng Xin <mxin@nvidia.com>
---
 CHANGELOG.rst | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index ef8c0ecda36..fe45b45daa5 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -10,12 +10,12 @@ Changelog
 
 - 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.
-- Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
-- Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
 
 *Megatron Framework (M-LM / M-Bridge)*
 
 - Add SFT-masked data support to ``examples/megatron_bridge/distill.py``: ``--sft --sft_dataset_root <dir>`` distills on raw prompt-completion JSONL (``{"input", "output"}`` records) with the loss masked to the response tokens, using Megatron-Bridge's ``FinetuningDatasetConfig`` and the model's own HuggingFace tokenizer instead of the pretraining ``GPTDataset`` and ``NullTokenizer``.
+- Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
+- Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
 
 *Misc*