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

modelopt/torch/quantization/mode.py

@@ -60,6 +60,7 @@
     update_quantize_metadata,
 )
 from .model_calib import (
+    _warn_on_kv_cache_during_calibration,
     awq,
     gptq,
     layerwise_calibrate,
@@ -247,6 +248,7 @@ def wrapped_calib_func(
                 module._moe_calib_experts_ratio = moe_calib_experts_ratio
 
     if func is not None:
+        forward_loop = _warn_on_kv_cache_during_calibration(forward_loop)
         if layerwise:
             # All currently implemented PTQ algorithms support layerwise calibration;
             # future algorithms that need full-model context must add a guard here.

modelopt/torch/quantization/model_calib.py

@@ -2047,6 +2047,63 @@ def postprocess(module, name):
     max_calibrate(model, forward_loop)
 
 
+def _is_kv_cache(obj) -> bool:
+    """Duck-typed ``transformers.Cache``, so this stays framework-agnostic.
+
+    Tensors are by far the common argument, so they are rejected before the attribute
+    probes, which are comparatively slow.
+    """
+    if obj is None or isinstance(obj, torch.Tensor):
+        return False
+    return hasattr(obj, "update") and hasattr(obj, "get_seq_length")
+
+
+_KV_CACHE_WARNING = (
+    "Calibration ran with KV caching enabled. Calibration only gathers activation "
+    "statistics and never reads a cache, so it is wasted memory and compute; under "
+    "layerwise calibration it is also incorrect, because each layer's captured inputs "
+    "are replayed and a cache among them makes the layer attend over the keys and "
+    "values its own earlier replay wrote, corrupting everything downstream of "
+    "attention. Disable it in the calibration forward loop, for example "
+    "`model(**batch, use_cache=False)`, or set `model.config.use_cache = False` around "
+    "it. `modelopt.torch.utils.dataset_utils.create_forward_loop` already does this."
+)
+
+
+def _warn_on_kv_cache_during_calibration(forward_loop):
+    """Wrap *forward_loop* to warn once if a KV cache is live during calibration.
+
+    Checked on what modules receive rather than on what the model returns, because a
+    layerwise forward stops early and returns nothing.
+    """
+    if forward_loop is None:
+        return None
+
+    warned = False
+
+    def _check(module, args, kwargs):
+        nonlocal warned
+        if not warned and any(_is_kv_cache(v) for v in (*args, *kwargs.values())):
+            warned = True
+            warn_rank_0(_KV_CACHE_WARNING)
+
+    def checked_forward_loop(m):
+        # A cache is handed to composite blocks (the decoder layer, its attention), never
+        # to a leaf such as a Linear -- which is most of the module tree.
+        handles = [
+            mod.register_forward_pre_hook(_check, with_kwargs=True)
+            for mod in m.modules()
+            if next(mod.children(), None) is not None
+        ]
+        try:
+            return forward_loop(m)
+        finally:
+            for h in handles:
+                h.remove()
+
+    return checked_forward_loop
+
+
 @torch.no_grad()
 def layerwise_calibrate(
     model: nn.Module,
@@ -2125,21 +2182,6 @@ def _set_layer_status(status: str):
 
             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

tests/unit/torch/quantization/plugins/test_attention_quant.py

@@ -99,7 +99,9 @@ def test_kv_quant_hf(model_getter, attn_cls):
                 setattr(parent, attention_module, attn_cls())
 
     model_test(input_ids, **kwargs)
-    mtq.quantize(model_test, kv_cache_config, lambda model: model(input_ids, **kwargs))
+    mtq.quantize(
+        model_test, kv_cache_config, lambda model: model(input_ids, use_cache=False, **kwargs)
+    )
 
     for name, module in model_test.named_modules():
         if name.endswith(attention_module):
@@ -128,7 +130,7 @@ def test_kv_quant_bert():
     mtq.quantize(
         model_test,
         kv_cache_config,
-        lambda model: model(input_ids, attention_mask=attention_mask),
+        lambda model: model(input_ids, attention_mask=attention_mask, use_cache=False),
     )
 
     # BERT attention modules are at encoder.layer.X.attention.self

tests/unit/torch/quantization/plugins/test_huggingface.py

@@ -250,7 +250,8 @@ def test_autoquantize_huggingface(model_provider, method):
     input_ids = model.dummy_inputs["input_ids"]
 
     def forward_step(model, batch):
-        return model(**batch) if method == "gradient" else model(**batch).logits
+        out = model(**batch, use_cache=False)
+        return out if method == "gradient" else out.logits
 
     warnings.filterwarnings(
         "error", message="AutoQuantize: Error enabling gradient checkpointing for huggingface model"
@@ -301,7 +302,9 @@ def test_quantized_transformers_save_restore(tmp_path, model_cls, quant_config):
         raise ValueError(f"Unsupported quant_config: {quant_config}")
 
     model_ref = model_cls.from_pretrained(tiny_llama_dir)
-    mtq.quantize(model_ref, quant_config, lambda model: model(**model.dummy_inputs))
+    mtq.quantize(
+        model_ref, quant_config, lambda model: model(**model.dummy_inputs, use_cache=False)
+    )
     mtq.compress(model_ref)
     model_ref.save_pretrained(tiny_llama_dir / "modelopt_model")
     assert os.path.exists(tiny_llama_dir / "modelopt_model/modelopt_state.pth")

tests/unit/torch/quantization/plugins/test_peft.py

@@ -81,7 +81,7 @@ def test_peft_flow(tmp_path):
     input_ids = torch.randint(0, model_original.config.vocab_size, (1, 4))
 
     def forward_loop(model):
-        return model(input_ids)
+        return model(input_ids, use_cache=False)
 
     mtq.quantize(peft_model, mtq.INT8_DEFAULT_CFG, forward_loop)
     mtq.quantize(model_full, mtq.INT8_DEFAULT_CFG, forward_loop)

tests/unit/torch/quantization/test_layerwise_calibrate.py

@@ -17,11 +17,13 @@
 
 import copy
 import json
+import warnings
 from collections import deque
 
 import pytest
 import torch
 import torch.nn as nn
+from _test_utils.torch.transformers_models import get_tiny_llama
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.quantization.model_calib import layerwise_calibrate
@@ -1036,6 +1038,35 @@ def crashing_torch_save(obj, path, *args, **kwargs):
     assert manifest["last_completed_layer"] == 1, f"manifest leaked mid-window state: {manifest}"
 
 
+@pytest.mark.parametrize("layerwise", [False, True], ids=["non_layerwise", "layerwise"])
+def test_calibration_warns_when_a_kv_cache_is_live(layerwise):
+    """Calibration never reads a KV cache, and layerwise is corrupted by one.
+
+    Layerwise replays each layer's captured inputs, so a cache among them makes the
+    layer attend over the keys and values its own earlier replay wrote. Detected on what
+    modules receive, since a layerwise forward stops early and returns nothing.
+    """
+    tokens = torch.randint(0, 32, (1, 8))
+    cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
+    cfg["algorithm"] = (
+        {"method": "max", "layerwise": {"enable": True}} if layerwise else {"method": "max"}
+    )
+
+    with pytest.warns(UserWarning, match="KV caching enabled"):
+        mtq.quantize(get_tiny_llama(num_hidden_layers=3).eval(), cfg, lambda m: m(tokens))
+
+    # A loop that disables caching calibrates silently -- as do models that never build
+    # one at all, which is why non-HF paths (e.g. Megatron) are unaffected.
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        mtq.quantize(
+            get_tiny_llama(num_hidden_layers=3).eval(),
+            cfg,
+            lambda m: m(tokens, use_cache=False),
+        )
+    assert not [w for w in caught if "KV caching enabled" in str(w.message)]
+
+
 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.