fix(quantization): warn once when calibration runs with a KV cache

CHANGELOG.rst

@@ -43,6 +43,7 @@ Changelog
 - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7).
 - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``.
 - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs.
+- Fix layerwise calibration miscalibrating every module downstream of attention within a decoder layer; models using an explicit sliding-window mask raised a shape mismatch instead. Re-run calibration from a fresh ``layerwise.checkpoint_dir`` for any recipe with ``layerwise.enable: true``, including the shipped ``nvfp4_experts_only-kv_fp8_layerwise*`` and ``nvfp4_default-kv_none-gptq`` recipes; algorithms that derive weight scales or weight updates from activations (``gptq``, ``awq_lite``, ``awq_clip``, ``local_hessian``, ``smoothquant``) change their exported weights too, not just activation scales.
 
 0.46 (2026-08-17)
 ^^^^^^^^^^^^^^^^^

modelopt/torch/quantization/model_calib.py

@@ -36,6 +36,7 @@
     _CheckpointState,
 )
 from modelopt.torch.utils import print_rank_0, warn_rank_0
+from modelopt.torch.utils.dataset_utils import _disable_use_cache
 from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState, is_master
 from modelopt.torch.utils.distributed import is_initialized as dist_is_initialized
 from modelopt.torch.utils.distributed import size as dist_size
@@ -2111,76 +2112,66 @@ def _set_layer_status(status: str):
 
     input_getter = LayerActivationCollector(model, status_callback=_set_layer_status)
 
-    try:
-        input_getter._patch_all_layers(decoder_layers=transformer_layers)
-        resumed_inputs = ckpt.setup_resume(transformer_layers) if ckpt and start_layer > 0 else None
+    # Calibration never reads a KV cache, and a layer replayed with one would attend
+    # over the keys and values its own earlier replay wrote.
+    with _disable_use_cache(model):
+        try:
+            input_getter._patch_all_layers(decoder_layers=transformer_layers)
+            resumed_inputs = (
+                ckpt.setup_resume(transformer_layers) if ckpt and start_layer > 0 else None
+            )
 
-        # Bootstrap: get first layer's inputs (or use resumed inputs).
-        layer_inputs = input_getter.get_first_layer_inputs(
-            start_layer, resumed_inputs, forward_loop
-        )
+            # Bootstrap: get first layer's inputs (or use resumed inputs).
+            layer_inputs = input_getter.get_first_layer_inputs(
+                start_layer, resumed_inputs, forward_loop
+            )
 
-        for layer_idx in range(start_layer, num_layers):
-            layer = transformer_layers[layer_idx]
-
-            def _layer_forward_loop(m, _inputs=layer_inputs):
-                for args, kwargs_input in _inputs:
-                    # Reset past_key_values to prevent the KV cache from
-                    # accumulating across multiple forward replays (e.g.
-                    # max_calibrate then Hessian collection in GPTQ).
-                    # The layer doesn't need stale KV data — each replay
-                    # should start with a fresh cache.
-                    if (
-                        "past_key_values" in kwargs_input
-                        and kwargs_input["past_key_values"] is not None
-                    ):
-                        kwargs_input = dict(kwargs_input)
-                        cache = kwargs_input["past_key_values"]
-                        if hasattr(cache, "reset"):
-                            cache.reset()
-                        else:
-                            kwargs_input["past_key_values"] = None
-                    m(*args, **kwargs_input)
-
-            is_last = layer_idx + 1 >= num_layers
-
-            with persistent_materialization(layer, writeback=calib_mutates_weights):
-                # qdq_from_prev=False: capture before calib_func so the forward
-                # replay uses the original FP weights. Disable quantizers too in
-                # case any pre-calibration observer behavior would perturb the
-                # captured activations.
-                if not is_last and not qdq_from_prev:
-                    with set_quantizer_by_cfg_context(
-                        layer, [{"quantizer_name": "*", "enable": False}]
-                    ):
+            for layer_idx in range(start_layer, num_layers):
+                layer = transformer_layers[layer_idx]
+
+                def _layer_forward_loop(m, _inputs=layer_inputs):
+                    for args, kwargs_input in _inputs:
+                        m(*args, **kwargs_input)
+
+                is_last = layer_idx + 1 >= num_layers
+
+                with persistent_materialization(layer, writeback=calib_mutates_weights):
+                    # qdq_from_prev=False: capture before calib_func so the forward
+                    # replay uses the original FP weights. Disable quantizers too in
+                    # case any pre-calibration observer behavior would perturb the
+                    # captured activations.
+                    if not is_last and not qdq_from_prev:
+                        with set_quantizer_by_cfg_context(
+                            layer, [{"quantizer_name": "*", "enable": False}]
+                        ):
+                            next_inputs = input_getter.cache_outputs_for_next_layer_calib(
+                                layer, forward_loop
+                            )
+                        # cache_outputs left this layer in "run" mode with an empty
+                        # deque; reset so calib_func's replay hits the real forward.
+                        layer._layerwise_calib.mode = "original"
+
+                    calib_func(layer, _layer_forward_loop, **calib_kwargs)
+
+                    # qdq_from_prev=True: capture after calib_func so the next layer
+                    # sees QDQ error and any in-place weight updates from this layer.
+                    if not is_last and qdq_from_prev:
                         next_inputs = input_getter.cache_outputs_for_next_layer_calib(
                             layer, forward_loop
                         )
-                    # cache_outputs left this layer in "run" mode with an empty
-                    # deque; reset so calib_func's replay hits the real forward.
-                    layer._layerwise_calib.mode = "original"
-
-                calib_func(layer, _layer_forward_loop, **calib_kwargs)
-
-                # qdq_from_prev=True: capture after calib_func so the next layer
-                # sees QDQ error and any in-place weight updates from this layer.
-                if not is_last and qdq_from_prev:
-                    next_inputs = input_getter.cache_outputs_for_next_layer_calib(
-                        layer, forward_loop
-                    )
-                elif is_last:
-                    next_inputs = None
-
-                if ckpt:
-                    ckpt.save(layer_idx, model, transformer_layers, next_inputs)
-
-            layer_pbar.update(1)
-            del layer_inputs
-            torch.cuda.empty_cache()
-            layer_inputs = next_inputs  # noqa: F841 (used in next iteration's closure)
-    finally:
-        input_getter._unpatch_all_layers()
-        layer_pbar.close()
+                    elif is_last:
+                        next_inputs = None
+
+                    if ckpt:
+                        ckpt.save(layer_idx, model, transformer_layers, next_inputs)
+
+                layer_pbar.update(1)
+                del layer_inputs
+                torch.cuda.empty_cache()
+                layer_inputs = next_inputs  # noqa: F841 (used in next iteration's closure)
+        finally:
+            input_getter._unpatch_all_layers()
+            layer_pbar.close()
 
     if ckpt:
         ckpt.full_restore(transformer_layers, model)

tests/unit/torch/quantization/test_layerwise_calibrate.py

@@ -22,6 +22,12 @@
 import pytest
 import torch
 import torch.nn as nn
+from _test_utils.torch.transformers_models import (
+    get_tiny_gpt_oss,
+    get_tiny_llama,
+    get_tiny_nemotron_h,
+)
+from transformers.cache_utils import Cache
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.quantization.model_calib import layerwise_calibrate
@@ -1036,6 +1042,49 @@ def crashing_torch_save(obj, path, *args, **kwargs):
     assert manifest["last_completed_layer"] == 1, f"manifest leaked mid-window state: {manifest}"
 
 
+@pytest.mark.parametrize(
+    "factory",
+    [get_tiny_llama, get_tiny_nemotron_h, get_tiny_gpt_oss],
+    ids=["llama", "nemotron_h_hybrid", "gpt_oss_sliding_window"],
+)
+def test_layerwise_calibration_builds_no_kv_cache(factory):
+    """No cache may reach a decoder layer while layerwise calibration runs.
+
+    Layerwise replays each layer's captured inputs several times, so a cache on them
+    lets a layer attend over the keys and values its own earlier replay wrote. This is
+    prevented upstream, by not building one -- which rests on the model honouring
+    ``config.use_cache``, hence the sweep over attention styles.
+
+    Asserted structurally: a retained cache only accumulates across replays, and max
+    calibration's max absorbs that on a small model, so amaxes can match while the
+    invariant is broken.
+    """
+    model = factory().eval()
+    assert model.config.use_cache, "fixture must start with caching on to be meaningful"
+    layers = LayerActivationCollector.get_decoder_layers(model)
+
+    seen = []
+    handles = [
+        layer.register_forward_pre_hook(
+            lambda mod, args, kwargs: seen.extend(
+                v for v in (*args, *kwargs.values()) if isinstance(v, Cache)
+            ),
+            with_kwargs=True,
+        )
+        for layer in layers
+    ]
+    try:
+        cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
+        cfg["algorithm"] = {"method": "max", "layerwise": {"enable": True}}
+        mtq.quantize(model, cfg, lambda m: m(torch.randint(0, 20, (1, 8))))
+    finally:
+        for h in handles:
+            h.remove()
+
+    assert not seen, f"{len(seen)} KV cache(s) reached a decoder layer during calibration"
+    assert model.config.use_cache, "config.use_cache was not restored"
+
+
 def test_layerwise_checkpoint_mismatch_save_every_raises(monkeypatch, tmp_path):
     """Resuming with a different ``save_every`` than the checkpoint was produced
     with must raise — the on-disk window layout assumes a fixed value.