[Fix] Calibrate non-decoder modules during layerwise quantization

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*

modelopt/torch/quantization/model_calib.py

@@ -35,6 +35,7 @@
 from modelopt.torch.quantization.utils.layerwise_calib import (
     LayerActivationCollector,
     _CheckpointState,
+    _OutsideQuantizerCalibrator,
     _reconcile_export_with_resume,
 )
 from modelopt.torch.utils import print_rank_0, warn_rank_0
@@ -2086,6 +2087,20 @@ def layerwise_calibrate(
             "Layerwise calibration requires a model with identifiable transformer layers."
         )
 
+    outside_calibrator = _OutsideQuantizerCalibrator(
+        model,
+        transformer_layers,
+        forward_loop,
+        calib_func,
+        calib_kwargs,
+        qdq_from_prev,
+    )
+    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."
+        )
+
     num_layers = len(transformer_layers)
     print_rank_0(f"Layerwise calibration: Found {num_layers} transformer layers")
 
@@ -2217,6 +2232,8 @@ def _layer_forward_loop(m, _inputs=layer_inputs):
     if ckpt:
         ckpt.full_restore(transformer_layers, model)
 
+    outside_calibrator.calibrate()
+
     if exporter is not None:
         warn_rank_0(
             f"Layerwise export: wrote every layer shard to {exporter.export_dir}.{finalize_hint}"

modelopt/torch/quantization/utils/layerwise_calib.py

@@ -27,22 +27,25 @@
 import os
 import shutil
 from collections import deque
+from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
 import torch
 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
+    from collections.abc import Callable, Sequence
 
     from modelopt.torch.opt.searcher import ForwardLoop
 
@@ -67,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"})
@@ -96,11 +100,93 @@ 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
         )
 
 
+@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): _SkipLayer(child, forward_original=True) 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 _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,
+    ):
+        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
+
+        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):
+        """Calibrate outside quantizers while decoder state is hidden from caller traversal."""
+        if not self.enabled:
+            return
+
+        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, 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.
 

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

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