From d225e85d54710340cdf7833ce57c9c13020ab1e8 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Fri, 4 Sep 2026 20:34:05 +0000
Subject: [PATCH 01/11] [Fix] Calibrate non-decoder modules during layerwise
 quantization

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 modelopt/torch/quantization/model_calib.py    |  37 ++++
 .../quantization/utils/layerwise_calib.py     |  45 +++-
 .../quantization/test_layerwise_calibrate.py  | 194 +++++++++++++++++-
 3 files changed, 274 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index 5ca1e548294..31fef826768 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -20,6 +20,7 @@
 import time
 import warnings
 from collections.abc import Callable, Mapping, Sequence
+from contextlib import ExitStack
 from functools import partial
 from typing import Any, TypeAlias
 
@@ -34,6 +35,7 @@
 from modelopt.torch.quantization.utils.layerwise_calib import (
     LayerActivationCollector,
     _CheckpointState,
+    _hide_modules_from_traversal,
     _reconcile_export_with_resume,
 )
 from modelopt.torch.utils import print_rank_0, warn_rank_0
@@ -2084,6 +2086,20 @@ def layerwise_calibrate(
             "Layerwise calibration requires a model with identifiable transformer layers."
         )
 
+    decoder_owned_ids = {id(module) for layer in transformer_layers for module in layer.modules()}
+    has_enabled_outside_quantizer = any(
+        isinstance(module, TensorQuantizer)
+        and module.is_enabled
+        and id(module) not in decoder_owned_ids
+        for module in model.modules()
+    )
+
+    if export_dir is not None and has_enabled_outside_quantizer:
+        raise ValueError(
+            "Layerwise export does not support enabled quantizers outside transformer layers. "
+            "Calibrate without export_dir, then export the completed model separately."
+        )
+
     num_layers = len(transformer_layers)
     print_rank_0(f"Layerwise calibration: Found {num_layers} transformer layers")
 
@@ -2215,6 +2231,27 @@ def _layer_forward_loop(m, _inputs=layer_inputs):
     if ckpt:
         ckpt.full_restore(transformer_layers, model)
 
+    if has_enabled_outside_quantizer:
+        if any(device == "disk" for device in getattr(model, "hf_device_map", {}).values()):
+            warn_rank_0(
+                "Layerwise calibration found enabled quantizers outside transformer layers. "
+                "The required full-model calibration pass may be slow because disk-offloaded "
+                "decoder weights can be streamed for every batch."
+            )
+
+        with _hide_modules_from_traversal(model, transformer_layers):
+            if qdq_from_prev:
+                calib_func(model, forward_loop, **calib_kwargs)
+            else:
+                with ExitStack() as stack:
+                    for layer in transformer_layers:
+                        stack.enter_context(
+                            set_quantizer_by_cfg_context(
+                                layer, [{"quantizer_name": "*", "enable": False}]
+                            )
+                        )
+                    calib_func(model, forward_loop, **calib_kwargs)
+
     if exporter is not None:
         warn_rank_0(
             f"Layerwise export: wrote every layer shard to {exporter.export_dir}.{finalize_hint}"
diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py
index 1000caa301a..c02ecc3bd67 100644
--- a/modelopt/torch/quantization/utils/layerwise_calib.py
+++ b/modelopt/torch/quantization/utils/layerwise_calib.py
@@ -27,6 +27,7 @@
 import os
 import shutil
 from collections import deque
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
@@ -42,7 +43,7 @@
 )
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import Callable, Sequence
 
     from modelopt.torch.opt.searcher import ForwardLoop
 
@@ -101,6 +102,48 @@ def forward(self, *args, **kwargs):
         )
 
 
+class _ForwardOnlyLayer(nn.Module):
+    """Hide a layer from module traversal while preserving its forward execution."""
+
+    _PROXY_BLOCKLIST = _SkipLayer._PROXY_BLOCKLIST
+
+    def __init__(self, original: nn.Module):
+        super().__init__()
+        object.__setattr__(self, "_original", original)
+
+    def __getattr__(self, name: str):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            if name in self._PROXY_BLOCKLIST:
+                raise
+            return getattr(object.__getattribute__(self, "_original"), name)
+
+    def forward(self, *args, **kwargs):
+        return self._original(*args, **kwargs)
+
+
+@contextmanager
+def _hide_modules_from_traversal(model: nn.Module, modules: Sequence[nn.Module]):
+    """Temporarily hide registered modules while retaining their forward behavior."""
+    target_ids = {id(module) for module in modules}
+    slots = [
+        (parent, child_name, child)
+        for parent in tuple(model.modules())
+        for child_name, child in tuple(parent._modules.items())
+        if child is not None and id(child) in target_ids
+    ]
+    proxies = {id(child): _ForwardOnlyLayer(child) for _, _, child in slots}
+
+    try:
+        for parent, child_name, child in slots:
+            parent._modules[child_name] = proxies[id(child)]
+        yield
+    finally:
+        for parent, child_name, child in slots:
+            parent._modules[child_name] = child
+
+
 class LayerActivationCollector:
     """Collects layer activations for layerwise (layer-by-layer) calibration.
 
diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py
index bda8c6029b1..a9b4ab00471 100644
--- a/tests/unit/torch/quantization/test_layerwise_calibrate.py
+++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py
@@ -18,6 +18,7 @@
 import copy
 import json
 from collections import deque
+from contextlib import nullcontext
 
 import pytest
 import torch
@@ -26,7 +27,12 @@
 import modelopt.torch.quantization as mtq
 from modelopt.torch.quantization.model_calib import layerwise_calibrate
 from modelopt.torch.quantization.nn import TensorQuantizer
-from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector, _SkipLayer
+from modelopt.torch.quantization.utils.layerwise_calib import (
+    LayerActivationCollector,
+    _ForwardOnlyLayer,
+    _hide_modules_from_traversal,
+    _SkipLayer,
+)
 
 
 class _DecoderBlock(nn.Module):
@@ -63,6 +69,15 @@ def forward(self, x, **kwargs):
         return x
 
 
+class _TransformerWithLMHead(_SimpleTransformerModel):
+    def __init__(self, n_layers=3, dim=16):
+        super().__init__(n_layers=n_layers, dim=dim)
+        self.lm_head = nn.Linear(dim, 32, bias=False)
+
+    def forward(self, x, **kwargs):
+        return self.lm_head(super().forward(x, **kwargs))
+
+
 class _FlatMLP(nn.Module):
     """No decoder-layer structure -- should be rejected by layerwise_calibrate."""
 
@@ -74,6 +89,48 @@ def forward(self, x):
         return self.net(x)
 
 
+class _TrackingQuantizer(TensorQuantizer):
+    """Deterministic quantizer used to distinguish QDQ and FP activations."""
+
+    def __init__(self):
+        super().__init__(amax=3.0)
+        self.calls = 0
+
+    def forward(self, x):
+        self.calls += 1
+        return x / 2 if self.is_enabled else x
+
+
+class _QuantizedLayer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.quantizer = _TrackingQuantizer()
+
+    def forward(self, x):
+        return self.quantizer(x)
+
+
+class _QuantizedTail(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.quantizer = _TrackingQuantizer()
+        self.inputs = []
+
+    def forward(self, x):
+        self.inputs.append(x.detach().clone())
+        return self.quantizer(x)
+
+
+class _ModelWithQuantizedTail(nn.Module):
+    def __init__(self, with_tail=True):
+        super().__init__()
+        self.layers = nn.ModuleList([_QuantizedLayer()])
+        self.tail = _QuantizedTail() if with_tail else nn.Identity()
+
+    def forward(self, x):
+        return self.tail(self.layers[0](x))
+
+
 class _SimpleTwoLayerModel(nn.Module):
     """Minimal model with explicit layers for activation-collection tests."""
 
@@ -243,6 +300,122 @@ def test_layerwise_calib_empty_forward_loop_raises(monkeypatch):
         )
 
 
+@pytest.mark.parametrize("raises", [False, True])
+def test_hide_modules_from_traversal_restores_aliases(raises):
+    model = _ModelWithQuantizedTail()
+    original = model.layers[0]
+    model.layer_alias = original
+
+    error_context = pytest.raises(RuntimeError, match="injected") if raises else nullcontext()
+    with error_context, _hide_modules_from_traversal(model, [original]):
+        assert isinstance(model.layers[0], _ForwardOnlyLayer)
+        assert model.layer_alias is model.layers[0]
+        assert original not in model.modules()
+        assert original.quantizer not in model.modules()
+        if raises:
+            raise RuntimeError("injected")
+
+    assert model.layers[0] is original
+    assert model.layer_alias is original
+
+
+@pytest.mark.parametrize(
+    ("qdq_from_prev", "expected_tail_input"),
+    [(True, 1.0), (False, 2.0)],
+)
+def test_layerwise_calibrates_only_outside_quantizers_with_full_model_forward(
+    monkeypatch, qdq_from_prev, expected_tail_input
+):
+    monkeypatch.setattr(
+        LayerActivationCollector,
+        "_decoder_layer_support",
+        [(lambda m: hasattr(m, "layers"), lambda m: list(m.layers))],
+    )
+    model = _ModelWithQuantizedTail()
+    decoder_quantizer = model.layers[0].quantizer
+    calibrated_quantizers = []
+    calibrated_targets = []
+    decoder_amax_before_extra_pass = []
+
+    def calib_func(target, target_forward_loop):
+        calibrated_targets.append(target)
+        if target is model:
+            decoder_amax_before_extra_pass.append(decoder_quantizer._amax.clone())
+        calibrated_quantizers.append(
+            {id(module) for module in target.modules() if isinstance(module, TensorQuantizer)}
+        )
+        target_forward_loop(target)
+
+    layerwise_calibrate(
+        model,
+        lambda m: m(torch.tensor([2.0])),
+        calib_func,
+        get_qdq_activations_from_prev_layer=qdq_from_prev,
+    )
+
+    assert calibrated_targets == [model.layers[0], model]
+    assert calibrated_quantizers == [{id(decoder_quantizer)}, {id(model.tail.quantizer)}]
+    torch.testing.assert_close(model.tail.inputs[-1], torch.tensor([expected_tail_input]))
+    torch.testing.assert_close(decoder_quantizer._amax, decoder_amax_before_extra_pass[0])
+    assert decoder_quantizer.calls == (2 if qdq_from_prev else 1)
+    assert decoder_quantizer.is_enabled
+
+
+def test_layerwise_skips_full_model_pass_without_outside_quantizer(monkeypatch):
+    _register_test_discoverer(monkeypatch)
+    model = _ModelWithQuantizedTail(with_tail=False)
+    calibrated_targets = []
+
+    def calib_func(target, target_forward_loop):
+        calibrated_targets.append(target)
+        target_forward_loop(target)
+
+    layerwise_calibrate(model, lambda m: m(torch.tensor([2.0])), calib_func)
+
+    assert calibrated_targets == [model.layers[0]]
+
+
+@pytest.mark.parametrize(
+    ("device_map", "with_tail", "warns"),
+    [
+        ({"layers.0": "disk"}, True, True),
+        ({"layers.0": "cpu"}, True, False),
+        ({}, True, False),
+        ({"layers.0": "disk"}, False, False),
+    ],
+)
+def test_layerwise_disk_offload_warning_gating(monkeypatch, device_map, with_tail, warns):
+    _register_test_discoverer(monkeypatch)
+    model = _ModelWithQuantizedTail(with_tail=with_tail)
+    model.hf_device_map = device_map
+    warnings = []
+    monkeypatch.setattr("modelopt.torch.quantization.model_calib.warn_rank_0", warnings.append)
+
+    def calib_func(target, target_forward_loop):
+        target_forward_loop(target)
+
+    layerwise_calibrate(model, lambda m: m(torch.tensor([2.0])), calib_func)
+
+    assert bool(warnings) is warns
+    if warns:
+        assert "disk-offloaded" in warnings[0]
+
+
+def test_layerwise_export_rejects_enabled_outside_quantizer(monkeypatch, tmp_path):
+    _register_test_discoverer(monkeypatch)
+    model = _ModelWithQuantizedTail()
+
+    with pytest.raises(ValueError, match="outside transformer layers"):
+        layerwise_calibrate(
+            model,
+            lambda m: m(torch.tensor([2.0])),
+            lambda *_args, **_kwargs: None,
+            export_dir=str(tmp_path / "export"),
+        )
+
+    assert not (tmp_path / "export").exists()
+
+
 # ---------------------------------------------------------------------------
 # Skip / run / capture path verification tests
 # ---------------------------------------------------------------------------
@@ -725,6 +898,25 @@ def forward_loop(m):
         model(calib_data[0])
 
 
+def test_mtq_quantize_layerwise_calibrates_lm_head(monkeypatch):
+    _register_test_discoverer(monkeypatch)
+    config = _int8_cfg_with_algorithm(
+        {
+            "method": "max",
+            "layerwise": {"enable": True, "get_qdq_activations_from_prev_layer": True},
+        }
+    )
+    config["quant_cfg"].append({"quantizer_name": "*lm_head*", "enable": True})
+    model = _TransformerWithLMHead(n_layers=2, dim=16)
+    calib_data = [torch.randint(0, 32, (2, 8))]
+
+    mtq.quantize(model, config, forward_loop=lambda m: [m(batch) for batch in calib_data])
+
+    assert model.lm_head.input_quantizer._amax is not None
+    with torch.no_grad():
+        model(calib_data[0])
+
+
 @pytest.mark.parametrize(
     "algorithm",
     ["gptq", "awq_lite", "smoothquant", "mse"],

From 59d0ef9ac0808aa47be16922a7516eb42837e33f Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Wed, 9 Sep 2026 14:03:22 +0000
Subject: [PATCH 02/11] Avoid unnecessary layerwise calibration forwards

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 modelopt/torch/quantization/mode.py           |   2 +
 modelopt/torch/quantization/model_calib.py    |  62 ++++---
 .../quantization/test_layerwise_calibrate.py  | 168 +++++++++++++++++-
 3 files changed, 203 insertions(+), 29 deletions(-)

diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py
index db7704e89b5..f7c98df2da5 100644
--- a/modelopt/torch/quantization/mode.py
+++ b/modelopt/torch/quantization/mode.py
@@ -226,6 +226,8 @@ def wrapped_calib_func(
     So lets wrap them to be compatible with the ModelOpt convert entrypoint.
     """
     kwargs = config.model_dump()
+    if "skip_forward_without_activation_calib" not in config.model_fields_set:
+        kwargs.pop("skip_forward_without_activation_calib", None)
     method = kwargs.pop("method")
     layerwise_cfg = kwargs.pop("layerwise", None) or {}
     layerwise = layerwise_cfg.get("enable", False)
diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index 31fef826768..b2859ec1879 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -62,6 +62,7 @@
     enable_fake_quant,
     enable_quant,
     enable_weight_access_and_writeback,
+    has_accelerate_offload,
     is_quantized_column_parallel_linear,
     is_quantized_linear,
     is_quantized_row_parallel_linear,
@@ -2094,12 +2095,49 @@ def layerwise_calibrate(
         for module in model.modules()
     )
 
-    if export_dir is not None and has_enabled_outside_quantizer:
+    outside_calib_needs_forward = has_enabled_outside_quantizer
+    outside_calib_runs_forward = has_enabled_outside_quantizer
+    outside_calib_kwargs = calib_kwargs
+    if has_enabled_outside_quantizer and calib_func is max_calibrate:
+        with _hide_modules_from_traversal(model, transformer_layers):
+            outside_calib_needs_forward = _needs_activation_forward_for_max_calib(model)
+        outside_calib_kwargs = dict(calib_kwargs)
+        outside_calib_kwargs.setdefault("skip_forward_without_activation_calib", True)
+        outside_calib_runs_forward = (
+            outside_calib_needs_forward
+            or not outside_calib_kwargs["skip_forward_without_activation_calib"]
+        )
+
+    if export_dir is not None and outside_calib_runs_forward:
         raise ValueError(
             "Layerwise export does not support enabled quantizers outside transformer layers. "
             "Calibrate without export_dir, then export the completed model separately."
         )
 
+    def _calibrate_outside_quantizers():
+        if not has_enabled_outside_quantizer:
+            return
+
+        if outside_calib_runs_forward and has_accelerate_offload(model):
+            warn_rank_0(
+                "Layerwise calibration found enabled quantizers outside transformer layers. "
+                "The required full-model calibration pass may be slow because CPU- or "
+                "disk-offloaded decoder weights can be transferred for every batch."
+            )
+
+        with _hide_modules_from_traversal(model, transformer_layers):
+            if qdq_from_prev:
+                calib_func(model, forward_loop, **outside_calib_kwargs)
+            else:
+                with ExitStack() as stack:
+                    for layer in transformer_layers:
+                        stack.enter_context(
+                            set_quantizer_by_cfg_context(
+                                layer, [{"quantizer_name": "*", "enable": False}]
+                            )
+                        )
+                    calib_func(model, forward_loop, **outside_calib_kwargs)
+
     num_layers = len(transformer_layers)
     print_rank_0(f"Layerwise calibration: Found {num_layers} transformer layers")
 
@@ -2130,6 +2168,7 @@ def layerwise_calibrate(
     if exporter is not None and _reconcile_export_with_resume(
         exporter, checkpoint_dir, start_layer, num_layers
     ):
+        _calibrate_outside_quantizers()
         warn_rank_0(
             f"Layerwise export: every layer shard in {exporter.export_dir} is already "
             f"written.{finalize_hint}"
@@ -2231,26 +2270,7 @@ def _layer_forward_loop(m, _inputs=layer_inputs):
     if ckpt:
         ckpt.full_restore(transformer_layers, model)
 
-    if has_enabled_outside_quantizer:
-        if any(device == "disk" for device in getattr(model, "hf_device_map", {}).values()):
-            warn_rank_0(
-                "Layerwise calibration found enabled quantizers outside transformer layers. "
-                "The required full-model calibration pass may be slow because disk-offloaded "
-                "decoder weights can be streamed for every batch."
-            )
-
-        with _hide_modules_from_traversal(model, transformer_layers):
-            if qdq_from_prev:
-                calib_func(model, forward_loop, **calib_kwargs)
-            else:
-                with ExitStack() as stack:
-                    for layer in transformer_layers:
-                        stack.enter_context(
-                            set_quantizer_by_cfg_context(
-                                layer, [{"quantizer_name": "*", "enable": False}]
-                            )
-                        )
-                    calib_func(model, forward_loop, **calib_kwargs)
+    _calibrate_outside_quantizers()
 
     if exporter is not None:
         warn_rank_0(
diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py
index a9b4ab00471..541240c50ce 100644
--- a/tests/unit/torch/quantization/test_layerwise_calibrate.py
+++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py
@@ -376,19 +376,21 @@ def calib_func(target, target_forward_loop):
 
 
 @pytest.mark.parametrize(
-    ("device_map", "with_tail", "warns"),
+    ("offloaded", "with_tail", "warns"),
     [
-        ({"layers.0": "disk"}, True, True),
-        ({"layers.0": "cpu"}, True, False),
-        ({}, True, False),
-        ({"layers.0": "disk"}, False, False),
+        (True, True, True),
+        (False, True, False),
+        (True, False, False),
     ],
 )
-def test_layerwise_disk_offload_warning_gating(monkeypatch, device_map, with_tail, warns):
+def test_layerwise_offload_warning_gating(monkeypatch, offloaded, with_tail, warns):
     _register_test_discoverer(monkeypatch)
     model = _ModelWithQuantizedTail(with_tail=with_tail)
-    model.hf_device_map = device_map
     warnings = []
+    monkeypatch.setattr(
+        "modelopt.torch.quantization.model_calib.has_accelerate_offload",
+        lambda target: target is model and offloaded,
+    )
     monkeypatch.setattr("modelopt.torch.quantization.model_calib.warn_rank_0", warnings.append)
 
     def calib_func(target, target_forward_loop):
@@ -398,7 +400,40 @@ def calib_func(target, target_forward_loop):
 
     assert bool(warnings) is warns
     if warns:
-        assert "disk-offloaded" in warnings[0]
+        assert "CPU- or disk-offloaded" in warnings[0]
+
+
+@pytest.mark.parametrize(
+    ("skip_forward_without_activation_calib", "expected_forward_calls", "warns"),
+    [(None, 1, False), (False, 2, True), (True, 1, False)],
+)
+def test_layerwise_max_offload_warning_matches_outside_forward(
+    monkeypatch, skip_forward_without_activation_calib, expected_forward_calls, warns
+):
+    _register_test_discoverer(monkeypatch)
+    monkeypatch.setattr(
+        "modelopt.torch.quantization.model_calib.has_accelerate_offload", lambda _model: True
+    )
+    warnings = []
+    monkeypatch.setattr("modelopt.torch.quantization.model_calib.warn_rank_0", warnings.append)
+    config = copy.deepcopy(mtq.INT8_WEIGHT_ONLY_CFG)
+    config["quant_cfg"].append({"quantizer_name": "*lm_head*weight_quantizer", "enable": True})
+    algorithm = {"method": "max", "layerwise": {"enable": True}}
+    if skip_forward_without_activation_calib is not None:
+        algorithm["skip_forward_without_activation_calib"] = skip_forward_without_activation_calib
+    config["algorithm"] = algorithm
+    model = _TransformerWithLMHead(n_layers=1, dim=16)
+    forward_calls = 0
+
+    def forward_loop(target):
+        nonlocal forward_calls
+        forward_calls += 1
+        target(torch.randint(0, 32, (2, 8)))
+
+    mtq.quantize(model, config, forward_loop=forward_loop)
+
+    assert forward_calls == expected_forward_calls
+    assert bool(warnings) is warns
 
 
 def test_layerwise_export_rejects_enabled_outside_quantizer(monkeypatch, tmp_path):
@@ -416,6 +451,123 @@ def test_layerwise_export_rejects_enabled_outside_quantizer(monkeypatch, tmp_pat
     assert not (tmp_path / "export").exists()
 
 
+def test_layerwise_export_rejects_explicit_outside_forward(monkeypatch, tmp_path):
+    _register_test_discoverer(monkeypatch)
+    config = copy.deepcopy(mtq.INT8_WEIGHT_ONLY_CFG)
+    config["quant_cfg"].append({"quantizer_name": "*lm_head*weight_quantizer", "enable": True})
+    config["algorithm"] = {
+        "method": "max",
+        "skip_forward_without_activation_calib": False,
+        "layerwise": {"enable": True, "export_dir": str(tmp_path / "export")},
+    }
+
+    with pytest.raises(ValueError, match="outside transformer layers"):
+        mtq.quantize(
+            _TransformerWithLMHead(n_layers=1, dim=16),
+            config,
+            forward_loop=lambda model: model(torch.randint(0, 32, (2, 8))),
+        )
+
+    assert not (tmp_path / "export").exists()
+
+
+def test_layerwise_export_allows_weight_only_outside_quantizer(monkeypatch, tmp_path):
+    _register_test_discoverer(monkeypatch)
+
+    class _FakeExporter:
+        instances = []
+
+        def __init__(self, model, export_dir):
+            self.exported_layers = []
+            self.finalized = False
+            self.instances.append(self)
+
+        def export_layer(self, layer_idx, layer, layer_inputs):
+            self.exported_layers.append(layer_idx)
+
+        def finalize(self):
+            self.finalized = True
+
+    monkeypatch.setattr("modelopt.torch.export.layerwise_export.LayerwiseExporter", _FakeExporter)
+    config = copy.deepcopy(mtq.INT8_WEIGHT_ONLY_CFG)
+    config["quant_cfg"].append({"quantizer_name": "*lm_head*weight_quantizer", "enable": True})
+    config["algorithm"] = {
+        "method": "max",
+        "layerwise": {"enable": True, "export_dir": str(tmp_path / "export")},
+    }
+    model = _TransformerWithLMHead(n_layers=1, dim=16)
+    calib_data = torch.randint(0, 32, (2, 8))
+    forward_calls = 0
+
+    def forward_loop(target):
+        nonlocal forward_calls
+        forward_calls += 1
+        target(calib_data)
+
+    mtq.quantize(model, config, forward_loop=forward_loop)
+
+    assert forward_calls == 1
+    assert model.lm_head.weight_quantizer._amax is not None
+    assert _FakeExporter.instances[0].exported_layers == [0]
+    assert _FakeExporter.instances[0].finalized
+
+
+@pytest.mark.parametrize(
+    "skip_forward_without_activation_calib",
+    [None, True],
+)
+def test_layerwise_export_completed_resume_calibrates_weight_only_tail(
+    monkeypatch, tmp_path, skip_forward_without_activation_calib
+):
+    _register_test_discoverer(monkeypatch)
+    checkpoint_dir = tmp_path / "checkpoint"
+    checkpoint_dir.mkdir()
+    (checkpoint_dir / "manifest.json").write_text(
+        json.dumps({"last_completed_layer": 0, "num_layers": 1})
+    )
+
+    class _FakeExporter:
+        instances = []
+
+        def __init__(self, model, export_dir):
+            self.model = model
+            self.finalized_tail_amax = None
+            self.instances.append(self)
+
+        def assert_shards_present(self, num_layers):
+            assert num_layers == 1
+
+        def finalize(self):
+            self.finalized_tail_amax = self.model.lm_head.weight_quantizer._amax
+
+    monkeypatch.setattr("modelopt.torch.export.layerwise_export.LayerwiseExporter", _FakeExporter)
+    config = copy.deepcopy(mtq.INT8_WEIGHT_ONLY_CFG)
+    config["quant_cfg"].append({"quantizer_name": "*lm_head*weight_quantizer", "enable": True})
+    algorithm = {
+        "method": "max",
+        "layerwise": {
+            "enable": True,
+            "checkpoint_dir": str(checkpoint_dir),
+            "export_dir": str(tmp_path / "export"),
+        },
+    }
+    if skip_forward_without_activation_calib is not None:
+        algorithm["skip_forward_without_activation_calib"] = skip_forward_without_activation_calib
+    config["algorithm"] = algorithm
+    model = _TransformerWithLMHead(n_layers=1, dim=16)
+    forward_calls = 0
+
+    def forward_loop(target):
+        nonlocal forward_calls
+        forward_calls += 1
+        target(torch.randint(0, 32, (2, 8)))
+
+    mtq.quantize(model, config, forward_loop=forward_loop)
+
+    assert forward_calls == 0
+    assert _FakeExporter.instances[0].finalized_tail_amax is not None
+
+
 # ---------------------------------------------------------------------------
 # Skip / run / capture path verification tests
 # ---------------------------------------------------------------------------

From afc81e95151733d74b4c9eb67798dfd468451719 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Wed, 9 Sep 2026 14:36:30 +0000
Subject: [PATCH 03/11] Preserve automatic layerwise calibration defaults

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 modelopt/torch/quantization/config.py                | 12 ++++++------
 modelopt/torch/quantization/mode.py                  |  2 +-
 tests/unit/torch/quantization/test_calib.py          | 12 +++++++++---
 .../torch/quantization/test_layerwise_calibrate.py   |  8 ++++++--
 4 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py
index 541cf5cefdf..0b3525e9c23 100644
--- a/modelopt/torch/quantization/config.py
+++ b/modelopt/torch/quantization/config.py
@@ -928,8 +928,8 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig):
         ),
     )
 
-    skip_forward_without_activation_calib: bool = ModeloptField(
-        default=False,
+    skip_forward_without_activation_calib: bool | None = ModeloptField(
+        default=None,
         title="Skip the calibration forward when no activation quantizer needs data.",
         description=(
             "If True, max calibration skips the ``forward_loop`` entirely when no enabled "
@@ -938,10 +938,10 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig):
             "dynamic, or MX (MXFP4/MXFP8) quantization. Weight calibration still runs on the "
             "weight tensors directly, so the quantized weights are unchanged; only the wasted "
             "forward is avoided. "
-            "Opt-in (default False) because the provided ``forward_loop`` can carry side "
-            "effects the caller relies on — most notably materializing sharded parameters under "
-            "DeepSpeed ZeRO-3 — so enable it per-recipe when the calibration data is known to be "
-            "unnecessary."
+            "Set False when the provided ``forward_loop`` carries side effects the caller relies "
+            "on — most notably materializing sharded parameters under DeepSpeed ZeRO-3. The "
+            "default None behaves as False for whole-model and per-layer calibration, and as True "
+            "for ModelOpt's generated layerwise non-decoder pass."
         ),
     )
 
diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py
index f7c98df2da5..b4fe0f38a67 100644
--- a/modelopt/torch/quantization/mode.py
+++ b/modelopt/torch/quantization/mode.py
@@ -226,7 +226,7 @@ def wrapped_calib_func(
     So lets wrap them to be compatible with the ModelOpt convert entrypoint.
     """
     kwargs = config.model_dump()
-    if "skip_forward_without_activation_calib" not in config.model_fields_set:
+    if kwargs.get("skip_forward_without_activation_calib") is None:
         kwargs.pop("skip_forward_without_activation_calib", None)
     method = kwargs.pop("method")
     layerwise_cfg = kwargs.pop("layerwise", None) or {}
diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py
index aa31d0b66d6..8a3a9ff77e2 100644
--- a/tests/unit/torch/quantization/test_calib.py
+++ b/tests/unit/torch/quantization/test_calib.py
@@ -714,9 +714,15 @@ def test_needs_activation_forward_for_static_bias_calibrator():
     assert not _needs_activation_forward_for_max_calib(const_static_bias)
 
 
-def test_max_calib_config_skip_is_opt_in():
-    """The flag is opt-in (default False) so it does not change behavior for direct callers."""
-    assert MaxCalibConfig().skip_forward_without_activation_calib is False
+def test_max_calib_config_skip_default_is_automatic():
+    """The automatic default preserves forward behavior for direct callers."""
+    assert MaxCalibConfig().skip_forward_without_activation_calib is None
+    assert (
+        MaxCalibConfig(
+            skip_forward_without_activation_calib=False
+        ).skip_forward_without_activation_calib
+        is False
+    )
     assert MaxCalibConfig(
         skip_forward_without_activation_calib=True
     ).skip_forward_without_activation_calib
diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py
index 541240c50ce..7635a0fe734 100644
--- a/tests/unit/torch/quantization/test_layerwise_calibrate.py
+++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py
@@ -471,7 +471,10 @@ def test_layerwise_export_rejects_explicit_outside_forward(monkeypatch, tmp_path
     assert not (tmp_path / "export").exists()
 
 
-def test_layerwise_export_allows_weight_only_outside_quantizer(monkeypatch, tmp_path):
+@pytest.mark.parametrize("algorithm_as_config", [False, True])
+def test_layerwise_export_allows_weight_only_outside_quantizer(
+    monkeypatch, tmp_path, algorithm_as_config
+):
     _register_test_discoverer(monkeypatch)
 
     class _FakeExporter:
@@ -491,10 +494,11 @@ def finalize(self):
     monkeypatch.setattr("modelopt.torch.export.layerwise_export.LayerwiseExporter", _FakeExporter)
     config = copy.deepcopy(mtq.INT8_WEIGHT_ONLY_CFG)
     config["quant_cfg"].append({"quantizer_name": "*lm_head*weight_quantizer", "enable": True})
-    config["algorithm"] = {
+    algorithm = {
         "method": "max",
         "layerwise": {"enable": True, "export_dir": str(tmp_path / "export")},
     }
+    config["algorithm"] = mtq.MaxCalibConfig(**algorithm) if algorithm_as_config else algorithm
     model = _TransformerWithLMHead(n_layers=1, dim=16)
     calib_data = torch.randint(0, 32, (2, 8))
     forward_calls = 0

From a776897329facf3663f3785d11ebb4c00ef2e870 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Thu, 10 Sep 2026 18:42:50 +0000
Subject: [PATCH 04/11] Refactor layerwise outside quantizer calibration

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 modelopt/torch/quantization/model_calib.py    | 62 +++------------
 .../quantization/utils/layerwise_calib.py     | 78 ++++++++++++++++++-
 .../quantization/test_layerwise_calibrate.py  | 32 ++++----
 3 files changed, 106 insertions(+), 66 deletions(-)

diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index b2859ec1879..9a72037723e 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -20,7 +20,6 @@
 import time
 import warnings
 from collections.abc import Callable, Mapping, Sequence
-from contextlib import ExitStack
 from functools import partial
 from typing import Any, TypeAlias
 
@@ -35,7 +34,7 @@
 from modelopt.torch.quantization.utils.layerwise_calib import (
     LayerActivationCollector,
     _CheckpointState,
-    _hide_modules_from_traversal,
+    _OutsideQuantizerCalibrator,
     _reconcile_export_with_resume,
 )
 from modelopt.torch.utils import print_rank_0, warn_rank_0
@@ -62,7 +61,6 @@
     enable_fake_quant,
     enable_quant,
     enable_weight_access_and_writeback,
-    has_accelerate_offload,
     is_quantized_column_parallel_linear,
     is_quantized_linear,
     is_quantized_row_parallel_linear,
@@ -2087,57 +2085,21 @@ def layerwise_calibrate(
             "Layerwise calibration requires a model with identifiable transformer layers."
         )
 
-    decoder_owned_ids = {id(module) for layer in transformer_layers for module in layer.modules()}
-    has_enabled_outside_quantizer = any(
-        isinstance(module, TensorQuantizer)
-        and module.is_enabled
-        and id(module) not in decoder_owned_ids
-        for module in model.modules()
+    outside_calibrator = _OutsideQuantizerCalibrator(
+        model,
+        transformer_layers,
+        forward_loop,
+        calib_func,
+        calib_kwargs,
+        qdq_from_prev,
+        _needs_activation_forward_for_max_calib if calib_func is max_calibrate else None,
     )
-
-    outside_calib_needs_forward = has_enabled_outside_quantizer
-    outside_calib_runs_forward = has_enabled_outside_quantizer
-    outside_calib_kwargs = calib_kwargs
-    if has_enabled_outside_quantizer and calib_func is max_calibrate:
-        with _hide_modules_from_traversal(model, transformer_layers):
-            outside_calib_needs_forward = _needs_activation_forward_for_max_calib(model)
-        outside_calib_kwargs = dict(calib_kwargs)
-        outside_calib_kwargs.setdefault("skip_forward_without_activation_calib", True)
-        outside_calib_runs_forward = (
-            outside_calib_needs_forward
-            or not outside_calib_kwargs["skip_forward_without_activation_calib"]
-        )
-
-    if export_dir is not None and outside_calib_runs_forward:
+    if export_dir is not None and outside_calibrator.runs_forward:
         raise ValueError(
             "Layerwise export does not support enabled quantizers outside transformer layers. "
             "Calibrate without export_dir, then export the completed model separately."
         )
 
-    def _calibrate_outside_quantizers():
-        if not has_enabled_outside_quantizer:
-            return
-
-        if outside_calib_runs_forward and has_accelerate_offload(model):
-            warn_rank_0(
-                "Layerwise calibration found enabled quantizers outside transformer layers. "
-                "The required full-model calibration pass may be slow because CPU- or "
-                "disk-offloaded decoder weights can be transferred for every batch."
-            )
-
-        with _hide_modules_from_traversal(model, transformer_layers):
-            if qdq_from_prev:
-                calib_func(model, forward_loop, **outside_calib_kwargs)
-            else:
-                with ExitStack() as stack:
-                    for layer in transformer_layers:
-                        stack.enter_context(
-                            set_quantizer_by_cfg_context(
-                                layer, [{"quantizer_name": "*", "enable": False}]
-                            )
-                        )
-                    calib_func(model, forward_loop, **outside_calib_kwargs)
-
     num_layers = len(transformer_layers)
     print_rank_0(f"Layerwise calibration: Found {num_layers} transformer layers")
 
@@ -2168,7 +2130,7 @@ def _calibrate_outside_quantizers():
     if exporter is not None and _reconcile_export_with_resume(
         exporter, checkpoint_dir, start_layer, num_layers
     ):
-        _calibrate_outside_quantizers()
+        outside_calibrator.calibrate()
         warn_rank_0(
             f"Layerwise export: every layer shard in {exporter.export_dir} is already "
             f"written.{finalize_hint}"
@@ -2270,7 +2232,7 @@ def _layer_forward_loop(m, _inputs=layer_inputs):
     if ckpt:
         ckpt.full_restore(transformer_layers, model)
 
-    _calibrate_outside_quantizers()
+    outside_calibrator.calibrate()
 
     if exporter is not None:
         warn_rank_0(
diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py
index c02ecc3bd67..fa9418921f9 100644
--- a/modelopt/torch/quantization/utils/layerwise_calib.py
+++ b/modelopt/torch/quantization/utils/layerwise_calib.py
@@ -27,7 +27,7 @@
 import os
 import shutil
 from collections import deque
-from contextlib import contextmanager
+from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
@@ -35,13 +35,15 @@
 import torch.nn as nn
 
 from modelopt.torch.utils import distributed as dist
-from modelopt.torch.utils import print_rank_0
+from modelopt.torch.utils import print_rank_0, warn_rank_0
 from modelopt.torch.utils.network import (
     bind_forward_method,
     get_module_device,
     unpatch_forward_method,
 )
 
+from .core_utils import has_accelerate_offload
+
 if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
 
@@ -144,6 +146,78 @@ def _hide_modules_from_traversal(model: nn.Module, modules: Sequence[nn.Module])
             parent._modules[child_name] = child
 
 
+class _OutsideQuantizerCalibrator:
+    """Calibrate enabled quantizers outside the layerwise decoder subtrees."""
+
+    def __init__(
+        self,
+        model: nn.Module,
+        transformer_layers: Sequence[nn.Module],
+        forward_loop: ForwardLoop,
+        calib_func: Callable,
+        calib_kwargs: dict[str, Any],
+        qdq_from_prev: bool,
+        activation_forward_predicate: Callable[[nn.Module], bool] | None = None,
+    ):
+        self.model = model
+        self.transformer_layers = transformer_layers
+        self.forward_loop = forward_loop
+        self.calib_func = calib_func
+        self.calib_kwargs = calib_kwargs
+        self.qdq_from_prev = qdq_from_prev
+
+        # Inline import breaks nn -> qtensor -> utils -> layerwise_calib import cycle.
+        from ..nn import TensorQuantizer
+
+        decoder_owned_ids = {
+            id(module) for layer in transformer_layers for module in layer.modules()
+        }
+        self.enabled = any(
+            isinstance(module, TensorQuantizer)
+            and module.is_enabled
+            and id(module) not in decoder_owned_ids
+            for module in model.modules()
+        )
+        self.runs_forward = self.enabled
+        if self.enabled and activation_forward_predicate is not None:
+            with _hide_modules_from_traversal(model, transformer_layers):
+                needs_forward = activation_forward_predicate(model)
+            self.calib_kwargs = dict(calib_kwargs)
+            self.calib_kwargs.setdefault("skip_forward_without_activation_calib", True)
+            self.runs_forward = (
+                needs_forward or not self.calib_kwargs["skip_forward_without_activation_calib"]
+            )
+
+    def calibrate(self):
+        """Run calibration while excluding decoder-owned quantizers from traversal."""
+        if not self.enabled:
+            return
+
+        if self.runs_forward and has_accelerate_offload(self.model):
+            warn_rank_0(
+                "Layerwise calibration found enabled quantizers outside transformer layers. "
+                "The required full-model calibration pass may be slow because CPU- or "
+                "disk-offloaded decoder weights can be transferred for every batch."
+            )
+
+        with _hide_modules_from_traversal(self.model, self.transformer_layers):
+            if self.qdq_from_prev:
+                self.calib_func(self.model, self.forward_loop, **self.calib_kwargs)
+                return
+
+            # Inline import breaks conversion -> utils -> layerwise_calib import cycle.
+            from ..conversion import set_quantizer_by_cfg_context
+
+            with ExitStack() as stack:
+                for layer in self.transformer_layers:
+                    stack.enter_context(
+                        set_quantizer_by_cfg_context(
+                            layer, [{"quantizer_name": "*", "enable": False}]
+                        )
+                    )
+                self.calib_func(self.model, self.forward_loop, **self.calib_kwargs)
+
+
 class LayerActivationCollector:
     """Collects layer activations for layerwise (layer-by-layer) calibration.
 
diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py
index 7635a0fe734..628f0be7680 100644
--- a/tests/unit/torch/quantization/test_layerwise_calibrate.py
+++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py
@@ -388,10 +388,12 @@ def test_layerwise_offload_warning_gating(monkeypatch, offloaded, with_tail, war
     model = _ModelWithQuantizedTail(with_tail=with_tail)
     warnings = []
     monkeypatch.setattr(
-        "modelopt.torch.quantization.model_calib.has_accelerate_offload",
+        "modelopt.torch.quantization.utils.layerwise_calib.has_accelerate_offload",
         lambda target: target is model and offloaded,
     )
-    monkeypatch.setattr("modelopt.torch.quantization.model_calib.warn_rank_0", warnings.append)
+    monkeypatch.setattr(
+        "modelopt.torch.quantization.utils.layerwise_calib.warn_rank_0", warnings.append
+    )
 
     def calib_func(target, target_forward_loop):
         target_forward_loop(target)
@@ -412,10 +414,13 @@ def test_layerwise_max_offload_warning_matches_outside_forward(
 ):
     _register_test_discoverer(monkeypatch)
     monkeypatch.setattr(
-        "modelopt.torch.quantization.model_calib.has_accelerate_offload", lambda _model: True
+        "modelopt.torch.quantization.utils.layerwise_calib.has_accelerate_offload",
+        lambda _model: True,
     )
     warnings = []
-    monkeypatch.setattr("modelopt.torch.quantization.model_calib.warn_rank_0", warnings.append)
+    monkeypatch.setattr(
+        "modelopt.torch.quantization.utils.layerwise_calib.warn_rank_0", warnings.append
+    )
     config = copy.deepcopy(mtq.INT8_WEIGHT_ONLY_CFG)
     config["quant_cfg"].append({"quantizer_name": "*lm_head*weight_quantizer", "enable": True})
     algorithm = {"method": "max", "layerwise": {"enable": True}}
@@ -481,16 +486,16 @@ class _FakeExporter:
         instances = []
 
         def __init__(self, model, export_dir):
+            self.export_dir = export_dir
             self.exported_layers = []
-            self.finalized = False
             self.instances.append(self)
 
+        def bind(self, calibrated_layers):
+            self.calibrated_layers = calibrated_layers
+
         def export_layer(self, layer_idx, layer, layer_inputs):
             self.exported_layers.append(layer_idx)
 
-        def finalize(self):
-            self.finalized = True
-
     monkeypatch.setattr("modelopt.torch.export.layerwise_export.LayerwiseExporter", _FakeExporter)
     config = copy.deepcopy(mtq.INT8_WEIGHT_ONLY_CFG)
     config["quant_cfg"].append({"quantizer_name": "*lm_head*weight_quantizer", "enable": True})
@@ -513,7 +518,6 @@ def forward_loop(target):
     assert forward_calls == 1
     assert model.lm_head.weight_quantizer._amax is not None
     assert _FakeExporter.instances[0].exported_layers == [0]
-    assert _FakeExporter.instances[0].finalized
 
 
 @pytest.mark.parametrize(
@@ -535,15 +539,15 @@ class _FakeExporter:
 
         def __init__(self, model, export_dir):
             self.model = model
-            self.finalized_tail_amax = None
+            self.export_dir = export_dir
             self.instances.append(self)
 
+        def bind(self, calibrated_layers):
+            self.calibrated_layers = calibrated_layers
+
         def assert_shards_present(self, num_layers):
             assert num_layers == 1
 
-        def finalize(self):
-            self.finalized_tail_amax = self.model.lm_head.weight_quantizer._amax
-
     monkeypatch.setattr("modelopt.torch.export.layerwise_export.LayerwiseExporter", _FakeExporter)
     config = copy.deepcopy(mtq.INT8_WEIGHT_ONLY_CFG)
     config["quant_cfg"].append({"quantizer_name": "*lm_head*weight_quantizer", "enable": True})
@@ -569,7 +573,7 @@ def forward_loop(target):
     mtq.quantize(model, config, forward_loop=forward_loop)
 
     assert forward_calls == 0
-    assert _FakeExporter.instances[0].finalized_tail_amax is not None
+    assert model.lm_head.weight_quantizer._amax is not None
 
 
 # ---------------------------------------------------------------------------

From dfde7960e9c2aa0b665e1e9cdb39403d7ea14a8d Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Thu, 10 Sep 2026 20:14:44 +0000
Subject: [PATCH 05/11] Simplify outside layerwise quantizer calibration

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 modelopt/torch/quantization/config.py         |  12 +-
 modelopt/torch/quantization/mode.py           |   2 +-
 modelopt/torch/quantization/model_calib.py    |   3 +-
 .../quantization/utils/layerwise_calib.py     |  72 ++++-----
 tests/unit/torch/quantization/test_calib.py   |  18 +--
 .../quantization/test_layerwise_calibrate.py  | 152 +++++-------------
 6 files changed, 89 insertions(+), 170 deletions(-)

diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py
index 0b3525e9c23..541cf5cefdf 100644
--- a/modelopt/torch/quantization/config.py
+++ b/modelopt/torch/quantization/config.py
@@ -928,8 +928,8 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig):
         ),
     )
 
-    skip_forward_without_activation_calib: bool | None = ModeloptField(
-        default=None,
+    skip_forward_without_activation_calib: bool = ModeloptField(
+        default=False,
         title="Skip the calibration forward when no activation quantizer needs data.",
         description=(
             "If True, max calibration skips the ``forward_loop`` entirely when no enabled "
@@ -938,10 +938,10 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig):
             "dynamic, or MX (MXFP4/MXFP8) quantization. Weight calibration still runs on the "
             "weight tensors directly, so the quantized weights are unchanged; only the wasted "
             "forward is avoided. "
-            "Set False when the provided ``forward_loop`` carries side effects the caller relies "
-            "on — most notably materializing sharded parameters under DeepSpeed ZeRO-3. The "
-            "default None behaves as False for whole-model and per-layer calibration, and as True "
-            "for ModelOpt's generated layerwise non-decoder pass."
+            "Opt-in (default False) because the provided ``forward_loop`` can carry side "
+            "effects the caller relies on — most notably materializing sharded parameters under "
+            "DeepSpeed ZeRO-3 — so enable it per-recipe when the calibration data is known to be "
+            "unnecessary."
         ),
     )
 
diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py
index b4fe0f38a67..f7c98df2da5 100644
--- a/modelopt/torch/quantization/mode.py
+++ b/modelopt/torch/quantization/mode.py
@@ -226,7 +226,7 @@ def wrapped_calib_func(
     So lets wrap them to be compatible with the ModelOpt convert entrypoint.
     """
     kwargs = config.model_dump()
-    if kwargs.get("skip_forward_without_activation_calib") is None:
+    if "skip_forward_without_activation_calib" not in config.model_fields_set:
         kwargs.pop("skip_forward_without_activation_calib", None)
     method = kwargs.pop("method")
     layerwise_cfg = kwargs.pop("layerwise", None) or {}
diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index 9a72037723e..999d5be6476 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -2092,9 +2092,8 @@ def layerwise_calibrate(
         calib_func,
         calib_kwargs,
         qdq_from_prev,
-        _needs_activation_forward_for_max_calib if calib_func is max_calibrate else None,
     )
-    if export_dir is not None and outside_calibrator.runs_forward:
+    if export_dir is not None and outside_calibrator.enabled:
         raise ValueError(
             "Layerwise export does not support enabled quantizers outside transformer layers. "
             "Calibrate without export_dir, then export the completed model separately."
diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py
index fa9418921f9..55a82b362a2 100644
--- a/modelopt/torch/quantization/utils/layerwise_calib.py
+++ b/modelopt/torch/quantization/utils/layerwise_calib.py
@@ -105,7 +105,7 @@ def forward(self, *args, **kwargs):
 
 
 class _ForwardOnlyLayer(nn.Module):
-    """Hide a layer from module traversal while preserving its forward execution."""
+    """Preserve forward execution while hiding a layer's modules, parameters, and buffers."""
 
     _PROXY_BLOCKLIST = _SkipLayer._PROXY_BLOCKLIST
 
@@ -126,15 +126,8 @@ def forward(self, *args, **kwargs):
 
 
 @contextmanager
-def _hide_modules_from_traversal(model: nn.Module, modules: Sequence[nn.Module]):
-    """Temporarily hide registered modules while retaining their forward behavior."""
-    target_ids = {id(module) for module in modules}
-    slots = [
-        (parent, child_name, child)
-        for parent in tuple(model.modules())
-        for child_name, child in tuple(parent._modules.items())
-        if child is not None and id(child) in target_ids
-    ]
+def _hide_modules_from_traversal(slots: Sequence[tuple[nn.Module, str, nn.Module]]):
+    """Retain forward behavior while hiding registered modules from traversal and state dicts."""
     proxies = {id(child): _ForwardOnlyLayer(child) for _, _, child in slots}
 
     try:
@@ -157,7 +150,6 @@ def __init__(
         calib_func: Callable,
         calib_kwargs: dict[str, Any],
         qdq_from_prev: bool,
-        activation_forward_predicate: Callable[[nn.Module], bool] | None = None,
     ):
         self.model = model
         self.transformer_layers = transformer_layers
@@ -169,40 +161,44 @@ def __init__(
         # Inline import breaks nn -> qtensor -> utils -> layerwise_calib import cycle.
         from ..nn import TensorQuantizer
 
-        decoder_owned_ids = {
-            id(module) for layer in transformer_layers for module in layer.modules()
-        }
-        self.enabled = any(
-            isinstance(module, TensorQuantizer)
-            and module.is_enabled
-            and id(module) not in decoder_owned_ids
-            for module in model.modules()
-        )
-        self.runs_forward = self.enabled
-        if self.enabled and activation_forward_predicate is not None:
-            with _hide_modules_from_traversal(model, transformer_layers):
-                needs_forward = activation_forward_predicate(model)
-            self.calib_kwargs = dict(calib_kwargs)
-            self.calib_kwargs.setdefault("skip_forward_without_activation_calib", True)
-            self.runs_forward = (
-                needs_forward or not self.calib_kwargs["skip_forward_without_activation_calib"]
+        layer_ids = {id(layer) for layer in transformer_layers}
+        self.transformer_layer_slots = [
+            (parent, child_name, child)
+            for parent in tuple(model.modules())
+            for child_name, child in parent._modules.items()
+            if child is not None and id(child) in layer_ids
+        ]
+
+        with _hide_modules_from_traversal(self.transformer_layer_slots):
+            self.enabled = any(
+                isinstance(module, TensorQuantizer) and module.is_enabled
+                for module in model.modules()
             )
 
     def calibrate(self):
-        """Run calibration while excluding decoder-owned quantizers from traversal."""
+        """Calibrate outside quantizers while decoder state is hidden from caller traversal."""
         if not self.enabled:
             return
 
-        if self.runs_forward and has_accelerate_offload(self.model):
-            warn_rank_0(
-                "Layerwise calibration found enabled quantizers outside transformer layers. "
-                "The required full-model calibration pass may be slow because CPU- or "
-                "disk-offloaded decoder weights can be transferred for every batch."
-            )
+        has_offload = has_accelerate_offload(self.model)
+        warned_for_offload = False
+
+        def forward_loop(model):
+            nonlocal warned_for_offload
+            if not warned_for_offload:
+                warned_for_offload = True
+                if has_offload:
+                    warn_rank_0(
+                        "Layerwise calibration found enabled quantizers outside transformer "
+                        "layers. The required full-model calibration pass may be slow because "
+                        "CPU- or disk-offloaded decoder weights can be transferred for every "
+                        "batch."
+                    )
+            return self.forward_loop(model)
 
-        with _hide_modules_from_traversal(self.model, self.transformer_layers):
+        with _hide_modules_from_traversal(self.transformer_layer_slots):
             if self.qdq_from_prev:
-                self.calib_func(self.model, self.forward_loop, **self.calib_kwargs)
+                self.calib_func(self.model, forward_loop, **self.calib_kwargs)
                 return
 
             # Inline import breaks conversion -> utils -> layerwise_calib import cycle.
@@ -215,7 +211,7 @@ def calibrate(self):
                             layer, [{"quantizer_name": "*", "enable": False}]
                         )
                     )
-                self.calib_func(self.model, self.forward_loop, **self.calib_kwargs)
+                self.calib_func(self.model, forward_loop, **self.calib_kwargs)
 
 
 class LayerActivationCollector:
diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py
index 8a3a9ff77e2..9f555be22f3 100644
--- a/tests/unit/torch/quantization/test_calib.py
+++ b/tests/unit/torch/quantization/test_calib.py
@@ -714,15 +714,15 @@ def test_needs_activation_forward_for_static_bias_calibrator():
     assert not _needs_activation_forward_for_max_calib(const_static_bias)
 
 
-def test_max_calib_config_skip_default_is_automatic():
-    """The automatic default preserves forward behavior for direct callers."""
-    assert MaxCalibConfig().skip_forward_without_activation_calib is None
-    assert (
-        MaxCalibConfig(
-            skip_forward_without_activation_calib=False
-        ).skip_forward_without_activation_calib
-        is False
-    )
+def test_max_calib_config_skip_is_opt_in():
+    """The flag is opt-in (default False) so it does not change behavior for direct callers."""
+    default_config = MaxCalibConfig()
+    assert default_config.skip_forward_without_activation_calib is False
+    assert "skip_forward_without_activation_calib" not in default_config.model_fields_set
+
+    explicit_config = MaxCalibConfig(skip_forward_without_activation_calib=False)
+    assert explicit_config.skip_forward_without_activation_calib is False
+    assert "skip_forward_without_activation_calib" in explicit_config.model_fields_set
     assert MaxCalibConfig(
         skip_forward_without_activation_calib=True
     ).skip_forward_without_activation_calib
diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py
index 628f0be7680..c4af9cdbfb9 100644
--- a/tests/unit/torch/quantization/test_layerwise_calibrate.py
+++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py
@@ -18,7 +18,6 @@
 import copy
 import json
 from collections import deque
-from contextlib import nullcontext
 
 import pytest
 import torch
@@ -30,7 +29,7 @@
 from modelopt.torch.quantization.utils.layerwise_calib import (
     LayerActivationCollector,
     _ForwardOnlyLayer,
-    _hide_modules_from_traversal,
+    _OutsideQuantizerCalibrator,
     _SkipLayer,
 )
 
@@ -301,21 +300,41 @@ def test_layerwise_calib_empty_forward_loop_raises(monkeypatch):
 
 
 @pytest.mark.parametrize("raises", [False, True])
-def test_hide_modules_from_traversal_restores_aliases(raises):
+def test_outside_calibrator_hides_and_restores_layer_aliases(raises):
     model = _ModelWithQuantizedTail()
     original = model.layers[0]
+    model.layers_alias = model.layers
     model.layer_alias = original
 
-    error_context = pytest.raises(RuntimeError, match="injected") if raises else nullcontext()
-    with error_context, _hide_modules_from_traversal(model, [original]):
-        assert isinstance(model.layers[0], _ForwardOnlyLayer)
-        assert model.layer_alias is model.layers[0]
-        assert original not in model.modules()
-        assert original.quantizer not in model.modules()
+    def calib_func(target, _forward_loop):
+        assert isinstance(target.layers[0], _ForwardOnlyLayer)
+        assert target.layers_alias[0] is target.layers[0]
+        assert target.layer_alias is target.layers[0]
+        assert original not in target.modules()
         if raises:
             raise RuntimeError("injected")
 
+    calibrator = _OutsideQuantizerCalibrator(
+        model,
+        [original],
+        lambda target: target(torch.tensor([2.0])),
+        calib_func,
+        {},
+        qdq_from_prev=True,
+    )
+    assert {(id(parent), name) for parent, name, _ in calibrator.transformer_layer_slots} == {
+        (id(model), "layer_alias"),
+        (id(model.layers), "0"),
+    }
+
+    if raises:
+        with pytest.raises(RuntimeError, match="injected"):
+            calibrator.calibrate()
+    else:
+        calibrator.calibrate()
+
     assert model.layers[0] is original
+    assert model.layers_alias[0] is original
     assert model.layer_alias is original
 
 
@@ -407,9 +426,9 @@ def calib_func(target, target_forward_loop):
 
 @pytest.mark.parametrize(
     ("skip_forward_without_activation_calib", "expected_forward_calls", "warns"),
-    [(None, 1, False), (False, 2, True), (True, 1, False)],
+    [(None, 2, True), (False, 2, True), (True, 1, False)],
 )
-def test_layerwise_max_offload_warning_matches_outside_forward(
+def test_layerwise_max_outside_calibration_uses_configured_forward_behavior(
     monkeypatch, skip_forward_without_activation_calib, expected_forward_calls, warns
 ):
     _register_test_discoverer(monkeypatch)
@@ -456,15 +475,20 @@ def test_layerwise_export_rejects_enabled_outside_quantizer(monkeypatch, tmp_pat
     assert not (tmp_path / "export").exists()
 
 
-def test_layerwise_export_rejects_explicit_outside_forward(monkeypatch, tmp_path):
+@pytest.mark.parametrize("skip_forward_without_activation_calib", [None, False, True])
+def test_layerwise_export_rejects_weight_only_outside_quantizer(
+    monkeypatch, tmp_path, skip_forward_without_activation_calib
+):
     _register_test_discoverer(monkeypatch)
     config = copy.deepcopy(mtq.INT8_WEIGHT_ONLY_CFG)
     config["quant_cfg"].append({"quantizer_name": "*lm_head*weight_quantizer", "enable": True})
-    config["algorithm"] = {
+    algorithm = {
         "method": "max",
-        "skip_forward_without_activation_calib": False,
         "layerwise": {"enable": True, "export_dir": str(tmp_path / "export")},
     }
+    if skip_forward_without_activation_calib is not None:
+        algorithm["skip_forward_without_activation_calib"] = skip_forward_without_activation_calib
+    config["algorithm"] = algorithm
 
     with pytest.raises(ValueError, match="outside transformer layers"):
         mtq.quantize(
@@ -476,106 +500,6 @@ def test_layerwise_export_rejects_explicit_outside_forward(monkeypatch, tmp_path
     assert not (tmp_path / "export").exists()
 
 
-@pytest.mark.parametrize("algorithm_as_config", [False, True])
-def test_layerwise_export_allows_weight_only_outside_quantizer(
-    monkeypatch, tmp_path, algorithm_as_config
-):
-    _register_test_discoverer(monkeypatch)
-
-    class _FakeExporter:
-        instances = []
-
-        def __init__(self, model, export_dir):
-            self.export_dir = export_dir
-            self.exported_layers = []
-            self.instances.append(self)
-
-        def bind(self, calibrated_layers):
-            self.calibrated_layers = calibrated_layers
-
-        def export_layer(self, layer_idx, layer, layer_inputs):
-            self.exported_layers.append(layer_idx)
-
-    monkeypatch.setattr("modelopt.torch.export.layerwise_export.LayerwiseExporter", _FakeExporter)
-    config = copy.deepcopy(mtq.INT8_WEIGHT_ONLY_CFG)
-    config["quant_cfg"].append({"quantizer_name": "*lm_head*weight_quantizer", "enable": True})
-    algorithm = {
-        "method": "max",
-        "layerwise": {"enable": True, "export_dir": str(tmp_path / "export")},
-    }
-    config["algorithm"] = mtq.MaxCalibConfig(**algorithm) if algorithm_as_config else algorithm
-    model = _TransformerWithLMHead(n_layers=1, dim=16)
-    calib_data = torch.randint(0, 32, (2, 8))
-    forward_calls = 0
-
-    def forward_loop(target):
-        nonlocal forward_calls
-        forward_calls += 1
-        target(calib_data)
-
-    mtq.quantize(model, config, forward_loop=forward_loop)
-
-    assert forward_calls == 1
-    assert model.lm_head.weight_quantizer._amax is not None
-    assert _FakeExporter.instances[0].exported_layers == [0]
-
-
-@pytest.mark.parametrize(
-    "skip_forward_without_activation_calib",
-    [None, True],
-)
-def test_layerwise_export_completed_resume_calibrates_weight_only_tail(
-    monkeypatch, tmp_path, skip_forward_without_activation_calib
-):
-    _register_test_discoverer(monkeypatch)
-    checkpoint_dir = tmp_path / "checkpoint"
-    checkpoint_dir.mkdir()
-    (checkpoint_dir / "manifest.json").write_text(
-        json.dumps({"last_completed_layer": 0, "num_layers": 1})
-    )
-
-    class _FakeExporter:
-        instances = []
-
-        def __init__(self, model, export_dir):
-            self.model = model
-            self.export_dir = export_dir
-            self.instances.append(self)
-
-        def bind(self, calibrated_layers):
-            self.calibrated_layers = calibrated_layers
-
-        def assert_shards_present(self, num_layers):
-            assert num_layers == 1
-
-    monkeypatch.setattr("modelopt.torch.export.layerwise_export.LayerwiseExporter", _FakeExporter)
-    config = copy.deepcopy(mtq.INT8_WEIGHT_ONLY_CFG)
-    config["quant_cfg"].append({"quantizer_name": "*lm_head*weight_quantizer", "enable": True})
-    algorithm = {
-        "method": "max",
-        "layerwise": {
-            "enable": True,
-            "checkpoint_dir": str(checkpoint_dir),
-            "export_dir": str(tmp_path / "export"),
-        },
-    }
-    if skip_forward_without_activation_calib is not None:
-        algorithm["skip_forward_without_activation_calib"] = skip_forward_without_activation_calib
-    config["algorithm"] = algorithm
-    model = _TransformerWithLMHead(n_layers=1, dim=16)
-    forward_calls = 0
-
-    def forward_loop(target):
-        nonlocal forward_calls
-        forward_calls += 1
-        target(torch.randint(0, 32, (2, 8)))
-
-    mtq.quantize(model, config, forward_loop=forward_loop)
-
-    assert forward_calls == 0
-    assert model.lm_head.weight_quantizer._amax is not None
-
-
 # ---------------------------------------------------------------------------
 # Skip / run / capture path verification tests
 # ---------------------------------------------------------------------------

From 76106e237c9c5dec1d75a1e891effc6dd953fe29 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Thu, 10 Sep 2026 20:32:39 +0000
Subject: [PATCH 06/11] Remove redundant calibration default handling

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 modelopt/torch/quantization/mode.py         | 2 --
 tests/unit/torch/quantization/test_calib.py | 5 -----
 2 files changed, 7 deletions(-)

diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py
index f7c98df2da5..db7704e89b5 100644
--- a/modelopt/torch/quantization/mode.py
+++ b/modelopt/torch/quantization/mode.py
@@ -226,8 +226,6 @@ def wrapped_calib_func(
     So lets wrap them to be compatible with the ModelOpt convert entrypoint.
     """
     kwargs = config.model_dump()
-    if "skip_forward_without_activation_calib" not in config.model_fields_set:
-        kwargs.pop("skip_forward_without_activation_calib", None)
     method = kwargs.pop("method")
     layerwise_cfg = kwargs.pop("layerwise", None) or {}
     layerwise = layerwise_cfg.get("enable", False)
diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py
index 9f555be22f3..22047efb73a 100644
--- a/tests/unit/torch/quantization/test_calib.py
+++ b/tests/unit/torch/quantization/test_calib.py
@@ -718,11 +718,6 @@ def test_max_calib_config_skip_is_opt_in():
     """The flag is opt-in (default False) so it does not change behavior for direct callers."""
     default_config = MaxCalibConfig()
     assert default_config.skip_forward_without_activation_calib is False
-    assert "skip_forward_without_activation_calib" not in default_config.model_fields_set
-
-    explicit_config = MaxCalibConfig(skip_forward_without_activation_calib=False)
-    assert explicit_config.skip_forward_without_activation_calib is False
-    assert "skip_forward_without_activation_calib" in explicit_config.model_fields_set
     assert MaxCalibConfig(
         skip_forward_without_activation_calib=True
     ).skip_forward_without_activation_calib

From dded6ad1c20d35b70d2fda7ed37e62ef67f41781 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Mon, 14 Sep 2026 12:46:24 +0000
Subject: [PATCH 07/11] Address layerwise calibration review feedback

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 modelopt/torch/quantization/model_calib.py        |  1 -
 .../quantization/plugins/test_accelerate_gpu.py   | 15 +++++++++++++++
 tests/unit/torch/quantization/test_calib.py       |  3 +--
 3 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index 999d5be6476..f3d87b0c966 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -2129,7 +2129,6 @@ def layerwise_calibrate(
     if exporter is not None and _reconcile_export_with_resume(
         exporter, checkpoint_dir, start_layer, num_layers
     ):
-        outside_calibrator.calibrate()
         warn_rank_0(
             f"Layerwise export: every layer shard in {exporter.export_dir} is already "
             f"written.{finalize_hint}"
diff --git a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py
index f585307958a..6aab8fe48c8 100644
--- a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py
+++ b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py
@@ -193,6 +193,21 @@ def test_layerwise_calibrate_cpu_offloaded(tmp_path, use_checkpoint):
         assert manifest["num_layers"] == num_layers
 
 
+def test_layerwise_calibrates_lm_head_with_accelerate_offload(tmp_path):
+    config = copy.deepcopy(mtq.INT8_DEFAULT_CFG)
+    config["quant_cfg"].append({"quantizer_name": "*lm_head*", "enable": True})
+    config = make_layerwise_cfg(config)
+    model, _, _, inputs = make_cpu_offloaded_model(tmp_path, num_hidden_layers=1)
+    offloaded_layer = model.model.layers[0]
+
+    assert isinstance(offloaded_layer._hf_hook, AlignDevicesHook)
+    assert all(parameter.device.type == "meta" for parameter in offloaded_layer.parameters())
+
+    mtq.quantize(model, config, lambda target: target(inputs))
+
+    assert model.lm_head.input_quantizer.amax is not None
+
+
 def test_sequential_checkpoint_resume_cpu_offloaded(tmp_path):
     """Resume from a partial checkpoint on a CPU-offloaded model matches a full run."""
     quant_cfg = mtq.NVFP4_AWQ_LITE_CFG
diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py
index 22047efb73a..aa31d0b66d6 100644
--- a/tests/unit/torch/quantization/test_calib.py
+++ b/tests/unit/torch/quantization/test_calib.py
@@ -716,8 +716,7 @@ def test_needs_activation_forward_for_static_bias_calibrator():
 
 def test_max_calib_config_skip_is_opt_in():
     """The flag is opt-in (default False) so it does not change behavior for direct callers."""
-    default_config = MaxCalibConfig()
-    assert default_config.skip_forward_without_activation_calib is False
+    assert MaxCalibConfig().skip_forward_without_activation_calib is False
     assert MaxCalibConfig(
         skip_forward_without_activation_calib=True
     ).skip_forward_without_activation_calib

From 9a0794588af67ecaad2c08b042f910f75b48e018 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Mon, 14 Sep 2026 13:41:56 +0000
Subject: [PATCH 08/11] Reuse skip layer for traversal-only forwarding

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 .../quantization/utils/layerwise_calib.py     | 30 ++++---------------
 .../quantization/test_layerwise_calibrate.py  |  3 +-
 2 files changed, 7 insertions(+), 26 deletions(-)

diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py
index 55a82b362a2..6702d327c9d 100644
--- a/modelopt/torch/quantization/utils/layerwise_calib.py
+++ b/modelopt/torch/quantization/utils/layerwise_calib.py
@@ -70,17 +70,18 @@ class _LayerCalibState:
 
 
 class _SkipLayer(nn.Module):
-    """Parameter-free stand-in for a fully calibrated decoder layer.
+    """Parameter-free stand-in that skips or forwards through a decoder layer.
 
     Replaces the real layer in the ModuleList so that framework hooks
     (accelerate, FSDP2, etc.) have no parameters to transfer. Holds a
     reference to the original layer for restoration during cleanup.
     """
 
-    def __init__(self, original: nn.Module):
+    def __init__(self, original: nn.Module, *, forward_original: bool = False):
         super().__init__()
         # Bypass nn.Module.__setattr__ to avoid registering original as a submodule.
         object.__setattr__(self, "_original", original)
+        self._forward_original = forward_original
         self._layerwise_calib = _LayerCalibState(mode="skip")
 
     _PROXY_BLOCKLIST = frozenset({"_hf_hook", "_old_forward"})
@@ -99,36 +100,17 @@ def __getattr__(self, name: str):
             return getattr(object.__getattribute__(self, "_original"), name)
 
     def forward(self, *args, **kwargs):
+        if self._forward_original:
+            return self._original(*args, **kwargs)
         return LayerActivationCollector._zeros_from_meta(
             self._original._layerwise_calib.output_meta
         )
 
 
-class _ForwardOnlyLayer(nn.Module):
-    """Preserve forward execution while hiding a layer's modules, parameters, and buffers."""
-
-    _PROXY_BLOCKLIST = _SkipLayer._PROXY_BLOCKLIST
-
-    def __init__(self, original: nn.Module):
-        super().__init__()
-        object.__setattr__(self, "_original", original)
-
-    def __getattr__(self, name: str):
-        try:
-            return super().__getattr__(name)
-        except AttributeError:
-            if name in self._PROXY_BLOCKLIST:
-                raise
-            return getattr(object.__getattribute__(self, "_original"), name)
-
-    def forward(self, *args, **kwargs):
-        return self._original(*args, **kwargs)
-
-
 @contextmanager
 def _hide_modules_from_traversal(slots: Sequence[tuple[nn.Module, str, nn.Module]]):
     """Retain forward behavior while hiding registered modules from traversal and state dicts."""
-    proxies = {id(child): _ForwardOnlyLayer(child) for _, _, child in slots}
+    proxies = {id(child): _SkipLayer(child, forward_original=True) for _, _, child in slots}
 
     try:
         for parent, child_name, child in slots:
diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py
index c4af9cdbfb9..664071496b5 100644
--- a/tests/unit/torch/quantization/test_layerwise_calibrate.py
+++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py
@@ -28,7 +28,6 @@
 from modelopt.torch.quantization.nn import TensorQuantizer
 from modelopt.torch.quantization.utils.layerwise_calib import (
     LayerActivationCollector,
-    _ForwardOnlyLayer,
     _OutsideQuantizerCalibrator,
     _SkipLayer,
 )
@@ -307,7 +306,7 @@ def test_outside_calibrator_hides_and_restores_layer_aliases(raises):
     model.layer_alias = original
 
     def calib_func(target, _forward_loop):
-        assert isinstance(target.layers[0], _ForwardOnlyLayer)
+        assert isinstance(target.layers[0], _SkipLayer)
         assert target.layers_alias[0] is target.layers[0]
         assert target.layer_alias is target.layers[0]
         assert original not in target.modules()

From b39586a1fd3eea3a994a1418f06f271849c8d0dc Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Mon, 14 Sep 2026 14:04:18 +0000
Subject: [PATCH 09/11] Simplify layerwise offload warning

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 .../quantization/utils/layerwise_calib.py     | 25 ++++++-----------
 .../quantization/test_layerwise_calibrate.py  | 27 ++++---------------
 2 files changed, 13 insertions(+), 39 deletions(-)

diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py
index 6702d327c9d..56e7554f522 100644
--- a/modelopt/torch/quantization/utils/layerwise_calib.py
+++ b/modelopt/torch/quantization/utils/layerwise_calib.py
@@ -162,25 +162,16 @@ def calibrate(self):
         if not self.enabled:
             return
 
-        has_offload = has_accelerate_offload(self.model)
-        warned_for_offload = False
-
-        def forward_loop(model):
-            nonlocal warned_for_offload
-            if not warned_for_offload:
-                warned_for_offload = True
-                if has_offload:
-                    warn_rank_0(
-                        "Layerwise calibration found enabled quantizers outside transformer "
-                        "layers. The required full-model calibration pass may be slow because "
-                        "CPU- or disk-offloaded decoder weights can be transferred for every "
-                        "batch."
-                    )
-            return self.forward_loop(model)
+        if has_accelerate_offload(self.model):
+            warn_rank_0(
+                "Layerwise calibration found enabled quantizers outside transformer layers. "
+                "Calibrating them may be slow because CPU- or disk-offloaded decoder weights "
+                "can be transferred for every batch."
+            )
 
         with _hide_modules_from_traversal(self.transformer_layer_slots):
             if self.qdq_from_prev:
-                self.calib_func(self.model, forward_loop, **self.calib_kwargs)
+                self.calib_func(self.model, self.forward_loop, **self.calib_kwargs)
                 return
 
             # Inline import breaks conversion -> utils -> layerwise_calib import cycle.
@@ -193,7 +184,7 @@ def forward_loop(model):
                             layer, [{"quantizer_name": "*", "enable": False}]
                         )
                     )
-                self.calib_func(self.model, forward_loop, **self.calib_kwargs)
+                self.calib_func(self.model, self.forward_loop, **self.calib_kwargs)
 
 
 class LayerActivationCollector:
diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py
index 664071496b5..07e14d50f2a 100644
--- a/tests/unit/torch/quantization/test_layerwise_calibrate.py
+++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py
@@ -413,10 +413,7 @@ def test_layerwise_offload_warning_gating(monkeypatch, offloaded, with_tail, war
         "modelopt.torch.quantization.utils.layerwise_calib.warn_rank_0", warnings.append
     )
 
-    def calib_func(target, target_forward_loop):
-        target_forward_loop(target)
-
-    layerwise_calibrate(model, lambda m: m(torch.tensor([2.0])), calib_func)
+    layerwise_calibrate(model, lambda m: m(torch.tensor([2.0])), lambda *_args: None)
 
     assert bool(warnings) is warns
     if warns:
@@ -424,21 +421,13 @@ def calib_func(target, target_forward_loop):
 
 
 @pytest.mark.parametrize(
-    ("skip_forward_without_activation_calib", "expected_forward_calls", "warns"),
-    [(None, 2, True), (False, 2, True), (True, 1, False)],
+    ("skip_forward_without_activation_calib", "expected_forward_calls"),
+    [(None, 2), (False, 2), (True, 1)],
 )
 def test_layerwise_max_outside_calibration_uses_configured_forward_behavior(
-    monkeypatch, skip_forward_without_activation_calib, expected_forward_calls, warns
+    monkeypatch, skip_forward_without_activation_calib, expected_forward_calls
 ):
     _register_test_discoverer(monkeypatch)
-    monkeypatch.setattr(
-        "modelopt.torch.quantization.utils.layerwise_calib.has_accelerate_offload",
-        lambda _model: True,
-    )
-    warnings = []
-    monkeypatch.setattr(
-        "modelopt.torch.quantization.utils.layerwise_calib.warn_rank_0", warnings.append
-    )
     config = copy.deepcopy(mtq.INT8_WEIGHT_ONLY_CFG)
     config["quant_cfg"].append({"quantizer_name": "*lm_head*weight_quantizer", "enable": True})
     algorithm = {"method": "max", "layerwise": {"enable": True}}
@@ -456,7 +445,6 @@ def forward_loop(target):
     mtq.quantize(model, config, forward_loop=forward_loop)
 
     assert forward_calls == expected_forward_calls
-    assert bool(warnings) is warns
 
 
 def test_layerwise_export_rejects_enabled_outside_quantizer(monkeypatch, tmp_path):
@@ -474,10 +462,7 @@ def test_layerwise_export_rejects_enabled_outside_quantizer(monkeypatch, tmp_pat
     assert not (tmp_path / "export").exists()
 
 
-@pytest.mark.parametrize("skip_forward_without_activation_calib", [None, False, True])
-def test_layerwise_export_rejects_weight_only_outside_quantizer(
-    monkeypatch, tmp_path, skip_forward_without_activation_calib
-):
+def test_layerwise_export_rejects_weight_only_outside_quantizer(monkeypatch, tmp_path):
     _register_test_discoverer(monkeypatch)
     config = copy.deepcopy(mtq.INT8_WEIGHT_ONLY_CFG)
     config["quant_cfg"].append({"quantizer_name": "*lm_head*weight_quantizer", "enable": True})
@@ -485,8 +470,6 @@ def test_layerwise_export_rejects_weight_only_outside_quantizer(
         "method": "max",
         "layerwise": {"enable": True, "export_dir": str(tmp_path / "export")},
     }
-    if skip_forward_without_activation_calib is not None:
-        algorithm["skip_forward_without_activation_calib"] = skip_forward_without_activation_calib
     config["algorithm"] = algorithm
 
     with pytest.raises(ValueError, match="outside transformer layers"):

From 963fa56a6d2f4734a2e6324997fdb616516cecc1 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Tue, 15 Sep 2026 20:33:18 +0000
Subject: [PATCH 10/11] Document layerwise support and skip broken LoRA test

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 CHANGELOG.rst                          | 2 ++
 tests/examples/llm_qat/test_llm_qat.py | 1 +
 2 files changed, 3 insertions(+)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index f0e1e67dddc..dddfec8e071 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -9,6 +9,7 @@ Changelog
 *Quantization*
 
 - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts.
+- Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration.
 - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/onnx_ptq/bevformer>`_ for details.
 
 *Misc*
@@ -17,6 +18,7 @@ Changelog
 
 **Backward Breaking Changes**
 
+- Layerwise calibration now raises ``ValueError`` when ``layerwise.export_dir`` is combined with an enabled quantizer outside the transformer layers. Calibrate without ``layerwise.export_dir``, then export the completed model separately.
 - Unified HuggingFace export now fails with ``NotImplementedError`` when it meets an MoE block whose expert projection names it does not know, instead of assuming Mixtral's ``w1``/``w2``/``w3``. If you hit this, register a ``ModelSpec`` for the model under ``modelopt/torch/models/``. Every MoE architecture ModelOpt exported correctly before this change is registered, so no supported model regresses.
 
 **Deprecations**
diff --git a/tests/examples/llm_qat/test_llm_qat.py b/tests/examples/llm_qat/test_llm_qat.py
index 4cd67d7f905..a7b610a807c 100644
--- a/tests/examples/llm_qat/test_llm_qat.py
+++ b/tests/examples/llm_qat/test_llm_qat.py
@@ -151,6 +151,7 @@ def test_qwen3_qat_nvfp4(tiny_qwen3_path, tmp_path, backend):
         cache_dir=cache_dir,
     )
 
+@pytest.mark.skip(reason="FSDP2 LoRA checkpoint save omits adapter_model.safetensors")
 def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path):
     ptq_output_dir = tmp_path / "ptq"
     cache_dir = str(tmp_path / "dataset_cache")

From e8a7c006bbfbe27446bb2a2cc553cbc76f381df2 Mon Sep 17 00:00:00 2001
From: realAsma <akuriparambi@nvidia.com>
Date: Tue, 15 Sep 2026 20:44:34 +0000
Subject: [PATCH 11/11] Remove layerwise export restriction changelog entry

Signed-off-by: realAsma <akuriparambi@nvidia.com>
---
 CHANGELOG.rst | 1 -
 1 file changed, 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index dddfec8e071..f8251217365 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -18,7 +18,6 @@ Changelog
 
 **Backward Breaking Changes**
 
-- Layerwise calibration now raises ``ValueError`` when ``layerwise.export_dir`` is combined with an enabled quantizer outside the transformer layers. Calibrate without ``layerwise.export_dir``, then export the completed model separately.
 - Unified HuggingFace export now fails with ``NotImplementedError`` when it meets an MoE block whose expert projection names it does not know, instead of assuming Mixtral's ``w1``/``w2``/``w3``. If you hit this, register a ``ModelSpec`` for the model under ``modelopt/torch/models/``. Every MoE architecture ModelOpt exported correctly before this change is registered, so no supported model regresses.
 
 **Deprecations**