fix(quantization): disable the KV cache for layerwise calibration

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)

modelopt/torch/quantization/utils/layerwise_calib.py

@@ -51,6 +51,11 @@ class _EarlyStopForwardError(Exception):
     """Raised to halt the forward pass after capturing layer inputs."""
 
 
+def _is_kv_cache(obj: Any) -> bool:
+    """Duck-typed ``transformers.Cache``, so this module need not import transformers."""
+    return hasattr(obj, "update") and hasattr(obj, "get_seq_length")
+
+
 @dataclass
 class _LayerCalibState:
     """Mutable per-layer state used during layerwise calibration.
@@ -241,6 +246,13 @@ def _patched_forward(self, *args, **kwargs):
                 return output
 
             if info.mode == "capture":
+                if any(_is_kv_cache(v) for v in (*args, *kwargs.values())):
+                    raise RuntimeError(
+                        f"Layerwise calibration captured a KV cache on layer {info.name!r}. "
+                        "Replaying it would attend over the layer's own earlier writes. "
+                        "Check that forward_loop does not pass use_cache=True or an "
+                        "explicit past_key_values."
+                    )
                 info.collected_inputs.append((args, kwargs))
                 raise _EarlyStopForwardError()
 

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,48 @@ 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_and_kv_caching(factory):
+    """No cache may reach a decoder layer, and one that does must be refused.
+
+    Layerwise replays each layer's captured inputs, so a cache on them lets a layer
+    attend over its own earlier replay's writes. Prevented by not building one, which
+    rests on the model honouring ``config.use_cache`` -- hence the sweep over attention
+    styles -- and backstopped by the capture check, since a caller may pass ``use_cache``
+    or ``past_key_values`` itself.
+    """
+    cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
+    cfg["algorithm"] = {"method": "max", "layerwise": {"enable": True}}
+    tokens = torch.randint(0, 20, (1, 8))
+
+    model = factory().eval()
+    assert model.config.use_cache, "fixture must start with caching on to be meaningful"
+    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 LayerActivationCollector.get_decoder_layers(model)
+    ]
+    try:
+        mtq.quantize(model, cfg, lambda m: m(tokens))
+    finally:
+        for h in handles:
+            h.remove()
+    assert not seen, f"{len(seen)} KV cache(s) reached a decoder layer"
+    assert model.config.use_cache, "config.use_cache was not restored"
+
+    with pytest.raises(RuntimeError, match="captured a KV cache"):
+        mtq.quantize(factory().eval(), cfg, lambda m: m(tokens, use_cache=True))
+
+
 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.