diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py index c096aaeb00e..f247d10efa3 100644 --- a/modelopt/torch/quantization/mode.py +++ b/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. diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index c85e97a104d..171e2e2e472 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/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 diff --git a/tests/unit/torch/quantization/plugins/test_attention_quant.py b/tests/unit/torch/quantization/plugins/test_attention_quant.py index 702cf3ad1db..fdf2078da9b 100644 --- a/tests/unit/torch/quantization/plugins/test_attention_quant.py +++ b/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 diff --git a/tests/unit/torch/quantization/plugins/test_huggingface.py b/tests/unit/torch/quantization/plugins/test_huggingface.py index 0da6cfbb1fc..7ddee5c1eaa 100644 --- a/tests/unit/torch/quantization/plugins/test_huggingface.py +++ b/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") diff --git a/tests/unit/torch/quantization/plugins/test_peft.py b/tests/unit/torch/quantization/plugins/test_peft.py index a3ef95cf6c1..dd29a970a74 100644 --- a/tests/unit/torch/quantization/plugins/test_peft.py +++ b/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) diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py index bda8c6029b1..46c7d913de1 100644 --- a/tests/unit/torch/quantization/test_layerwise_calibrate.py +++ b/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.