From 7baf1d095fb9857bf52212a323c3d27d646e5688 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 18:43:41 +0000 Subject: [PATCH 01/13] fix(quantization): stop a layerwise replay attending over its own KV cache Layerwise calibration replays each decoder layer on its captured inputs. The captured kwargs carry the model's ``past_key_values``, which the preceding capture pass has already written to, so the replay had to start from an empty cache. It called ``Cache.reset()`` for that -- but ``reset()`` "resets the cache values while preserving the objects": it zeroes the key/value tensors and leaves them at full length. The replay therefore attended over a same-length, all-zero cache instead of no cache at all. The result was silently wrong activation scales. ``self_attn.o_proj``'s input amax collapsed to exactly ``0.0`` on every layer but the last -- the one with no preceding capture pass, so its cache was never initialized and ``reset()`` was a no-op -- while ``down_proj`` picked up a plausible but wrong value from the residual alone. Models that pass an explicit sliding-window mask raised a shape mismatch instead, since the replay's concatenated cache is twice the mask width. Drop the cache instead; the layer recomputes keys and values from the captured inputs. Activation amaxes now match the non-layerwise path exactly on every architecture checked (llama, mixtral, nemotron, nemotron_h, and gpt_oss, which previously raised). Not a transformers regression: ``reset()`` has these semantics across the whole supported range, verified on 4.57.6 (tf_min) and 5.12.1. Weight-only recipes are unaffected -- weight amaxes never depended on the cache. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- CHANGELOG.rst | 1 + modelopt/torch/quantization/model_calib.py | 34 +++++++------- .../quantization/test_layerwise_calibrate.py | 47 +++++++++++++++++++ 3 files changed, 66 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6687ebd31ea..5238df6fba9 100755 --- a/CHANGELOG.rst +++ b/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 producing wrong **activation** amaxes: a replayed decoder layer attended over the KV cache its own earlier run had written, so ``self_attn.o_proj``'s input amax collapsed to ``0.0`` on every layer but the last and other intra-layer activation scales were silently off. Models with an explicit sliding-window mask raised a shape mismatch instead. Weight-only recipes are unaffected; re-run calibration for any layerwise recipe that quantizes activations. 0.46 (2026-08-17) ^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index c85e97a104d..24b0bec8cfa 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2047,6 +2047,23 @@ def postprocess(module, name): max_calibrate(model, forward_loop) +def _with_empty_kv_cache(kwargs_input: dict) -> dict: + """Return *kwargs_input* with any attention cache dropped. + + A layer replayed during calibration must not see the keys and values its own + earlier run wrote, or it attends over stale state and its activation amaxes + describe the wrong tensor. Dropping the cache is what makes each replay + independent; the layer recomputes keys and values from the captured inputs. + + ``Cache.reset()`` is not usable here: it zeroes the key/value tensors while + keeping them at full length, so the replay attends over a same-length, + all-zero cache rather than no cache at all. + """ + if kwargs_input.get("past_key_values") is None: + return kwargs_input + return {**kwargs_input, "past_key_values": None} + + @torch.no_grad() def layerwise_calibrate( model: nn.Module, @@ -2125,22 +2142,7 @@ 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) + m(*args, **_with_empty_kv_cache(kwargs_input)) is_last = layer_idx + 1 >= num_layers diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py index bda8c6029b1..6795e4c8ca8 100644 --- a/tests/unit/torch/quantization/test_layerwise_calibrate.py +++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py @@ -830,6 +830,53 @@ def fwd(m): _assert_amax_close(_collect_amax(model_lw), seq_amax, "layerwise vs sequential") +def test_layerwise_replay_does_not_attend_over_its_own_kv_cache(): + """A layer replayed during calibration must not see the keys and values its + own earlier run wrote. + + The equivalence test above uses a model with no attention cache, so it cannot + reach this: only the captured ``past_key_values`` makes a replay stateful. + With the cache left in place, ``o_proj`` saw an all-zero attention output on + every layer but the last (the one with no preceding capture pass) and its + input amax collapsed to exactly 0.0, while ``down_proj`` picked up a + plausible but wrong value from the residual alone. + """ + from _test_utils.torch.transformers_models import get_tiny_llama + + calib_data = [torch.randint(0, 32, (2, 8)) for _ in range(2)] + + def fwd(m): + for batch in calib_data: + m(batch) + + def calibrate(algorithm): + torch.manual_seed(0) + model = get_tiny_llama(num_hidden_layers=4).eval() + cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + cfg["algorithm"] = algorithm + mtq.quantize(model, cfg, forward_loop=fwd) + return model + + sequential = calibrate({"method": "max"}) + layerwise = calibrate( + {"method": "max", "layerwise": {"enable": True, "calib_mutates_weights": False}} + ) + + expected = _collect_amax(sequential) + assert expected, "sequential calibration populated no amax values" + layerwise_amax = _collect_amax(layerwise) + _assert_amax_close(layerwise_amax, expected, "layerwise vs sequential (KV cache)") + + # Pinned separately from the comparison: a future regression that made both + # paths collapse to zero would still satisfy the equality above. + collapsed = [ + name + for name, amax in layerwise_amax.items() + if name.endswith("input_quantizer") and not torch.count_nonzero(amax) + ] + assert not collapsed, f"activation amax collapsed to zero: {collapsed}" + + def test_layerwise_no_qdq_captures_inputs_before_calib_func_mutates_weights(monkeypatch): """A destructive ``calib_func`` (zeros weights) must not affect what is captured for downstream layers under ``qdq_from_prev=False`` — otherwise From 20df3d3535bad56b5704b04b3df8e949d846f6df Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 18:56:58 +0000 Subject: [PATCH 02/13] docs: compact the KV-cache fix docstring and changelog entry Trims the helper docstring to the one non-obvious point a reader needs -- that Cache.reset() zeroes in place rather than clearing -- and cuts the changelog entry to what an external user must act on. Also names the shipped recipes affected: the experts-only layerwise recipes enable NVFP4 input quantizers on the MoE experts, which sit after attention and so inherit the zeroed attention output, so "weight-only recipes are unaffected" understated the blast radius. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- CHANGELOG.rst | 2 +- modelopt/torch/quantization/model_calib.py | 12 +++--------- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5238df6fba9..05b029b23fa 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,7 +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 producing wrong **activation** amaxes: a replayed decoder layer attended over the KV cache its own earlier run had written, so ``self_attn.o_proj``'s input amax collapsed to ``0.0`` on every layer but the last and other intra-layer activation scales were silently off. Models with an explicit sliding-window mask raised a shape mismatch instead. Weight-only recipes are unaffected; re-run calibration for any layerwise recipe that quantizes activations. +- Fix layerwise calibration producing wrong activation amaxes: a replayed decoder layer attended over its own stale KV cache, so activation scales downstream of attention were calibrated against a zeroed attention output on every layer but the last (sliding-window models raised a shape mismatch instead). Re-run calibration for any layerwise recipe that quantizes activations, including the shipped ``nvfp4_experts_only-kv_fp8_layerwise`` recipes; weight-only recipes are unchanged. 0.46 (2026-08-17) ^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 24b0bec8cfa..7ba1bc2974b 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2048,16 +2048,10 @@ def postprocess(module, name): def _with_empty_kv_cache(kwargs_input: dict) -> dict: - """Return *kwargs_input* with any attention cache dropped. + """Drop any attention cache, so a replay does not attend over its own earlier writes. - A layer replayed during calibration must not see the keys and values its own - earlier run wrote, or it attends over stale state and its activation amaxes - describe the wrong tensor. Dropping the cache is what makes each replay - independent; the layer recomputes keys and values from the captured inputs. - - ``Cache.reset()`` is not usable here: it zeroes the key/value tensors while - keeping them at full length, so the replay attends over a same-length, - all-zero cache rather than no cache at all. + Not ``Cache.reset()``: that zeroes the key/value tensors but keeps them at full + length, leaving the replay attending over an all-zero cache instead of none. """ if kwargs_input.get("past_key_values") is None: return kwargs_input From 6013c55ca9d6412e02786cd7c7930798e1805776 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 20:16:49 +0000 Subject: [PATCH 03/13] fix(quantization): clear the KV cache on both replay paths, widen the changelog Addresses review on #2248. The run-mode replay in LayerActivationCollector consumes the same captured kwargs as calib_func, so both consumers of that shared cache object now clear it. It was correct before only because calib_func no longer wrote to the cache -- an implicit ordering dependency that reordering the two blocks would silently break. Verified a no-op today: the recipe end-to-end and the cache-independence oracle return identical results either way. The helper moves to layerwise_calib.py, which owns both replay paths (model_calib.py already imports from it, so the reverse would be circular). The changelog said "weight-only recipes are unchanged", which is wrong for the multi-pass calibrators: gptq replays for its Hessian pass and awq_lite for its search pass, so the corrupted activations fed the weight updates. Measured on a W4A16 weight-only config with zero activation quantizers, the fix changes 15 exported .weight tensors under awq_lite and 25 under gptq. Entry now names the affected shipped recipes and drops the root-cause detail that belongs in the PR. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- CHANGELOG.rst | 2 +- modelopt/torch/quantization/model_calib.py | 12 +----------- .../torch/quantization/utils/layerwise_calib.py | 17 ++++++++++++++++- .../quantization/test_layerwise_calibrate.py | 3 +-- 4 files changed, 19 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 05b029b23fa..248b37051a4 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,7 +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 producing wrong activation amaxes: a replayed decoder layer attended over its own stale KV cache, so activation scales downstream of attention were calibrated against a zeroed attention output on every layer but the last (sliding-window models raised a shape mismatch instead). Re-run calibration for any layerwise recipe that quantizes activations, including the shipped ``nvfp4_experts_only-kv_fp8_layerwise`` recipes; weight-only recipes are unchanged. +- 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 for any recipe with ``layerwise.enable: true``, including the shipped ``nvfp4_experts_only-kv_fp8_layerwise*`` and ``nvfp4_default-kv_none-gptq`` recipes — under ``gptq`` and ``awq_lite`` the exported weights change too, not just activation scales. 0.46 (2026-08-17) ^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 7ba1bc2974b..d48b49d8ad3 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -34,6 +34,7 @@ from modelopt.torch.quantization.utils.layerwise_calib import ( LayerActivationCollector, _CheckpointState, + _with_empty_kv_cache, ) from modelopt.torch.utils import print_rank_0, warn_rank_0 from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState, is_master @@ -2047,17 +2048,6 @@ def postprocess(module, name): max_calibrate(model, forward_loop) -def _with_empty_kv_cache(kwargs_input: dict) -> dict: - """Drop any attention cache, so a replay does not attend over its own earlier writes. - - Not ``Cache.reset()``: that zeroes the key/value tensors but keeps them at full - length, leaving the replay attending over an all-zero cache instead of none. - """ - if kwargs_input.get("past_key_values") is None: - return kwargs_input - return {**kwargs_input, "past_key_values": None} - - @torch.no_grad() def layerwise_calibrate( model: nn.Module, diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py index 070ee521cd5..30b4be1ed9f 100644 --- a/modelopt/torch/quantization/utils/layerwise_calib.py +++ b/modelopt/torch/quantization/utils/layerwise_calib.py @@ -51,6 +51,21 @@ class _EarlyStopForwardError(Exception): """Raised to halt the forward pass after capturing layer inputs.""" +def _with_empty_kv_cache(kwargs_input: dict) -> dict: + """Drop any attention cache, so a replay does not attend over its own earlier writes. + + Both replay paths -- the patched ``run`` forward below and ``calib_func``'s replay in + ``layerwise_calibrate`` -- consume the same captured kwargs, so each clears the cache + itself rather than relying on the other not having written to it. + + Not ``Cache.reset()``: that zeroes the key/value tensors but keeps them at full + length, leaving the replay attending over an all-zero cache instead of none. + """ + if kwargs_input.get("past_key_values") is None: + return kwargs_input + return {**kwargs_input, "past_key_values": None} + + @dataclass class _LayerCalibState: """Mutable per-layer state used during layerwise calibration. @@ -236,7 +251,7 @@ def _patched_forward(self, *args, **kwargs): f"Layer {info.name} is in 'run' mode but has no cached inputs to replay." ) real_args, real_kwargs = info.cached_inputs.popleft() - output = self._original_forward(*real_args, **real_kwargs) + output = self._original_forward(*real_args, **_with_empty_kv_cache(real_kwargs)) info.output_meta = LayerActivationCollector._extract_output_meta(output) return output diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py index 6795e4c8ca8..b99aea7a572 100644 --- a/tests/unit/torch/quantization/test_layerwise_calibrate.py +++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py @@ -22,6 +22,7 @@ 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 @@ -841,8 +842,6 @@ def test_layerwise_replay_does_not_attend_over_its_own_kv_cache(): input amax collapsed to exactly 0.0, while ``down_proj`` picked up a plausible but wrong value from the residual alone. """ - from _test_utils.torch.transformers_models import get_tiny_llama - calib_data = [torch.randint(0, 32, (2, 8)) for _ in range(2)] def fwd(m): From 0f120535ce2b41f360ec5bcec05da839bfbe7d4f Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 20:18:23 +0000 Subject: [PATCH 04/13] test: drop the inert calib_mutates_weights override from the KV-cache test Addresses review on #2248. With no checkpoint_dir the flag never reaches the weights.pt / quantizer_buffers.pt branch it controls, and writeback=False is a no-op for a resident CPU model -- so it only made the test diverge from the default that the shipped nvfp4_experts_only-kv_fp8_layerwise recipe runs. Still fails without the fix on o_proj's input amax. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- tests/unit/torch/quantization/test_layerwise_calibrate.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py index b99aea7a572..759640d16af 100644 --- a/tests/unit/torch/quantization/test_layerwise_calibrate.py +++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py @@ -857,9 +857,7 @@ def calibrate(algorithm): return model sequential = calibrate({"method": "max"}) - layerwise = calibrate( - {"method": "max", "layerwise": {"enable": True, "calib_mutates_weights": False}} - ) + layerwise = calibrate({"method": "max", "layerwise": {"enable": True}}) expected = _collect_amax(sequential) assert expected, "sequential calibration populated no amax values" From 3383c7a0e30d645071e75ffbb8eec44551ae0f49 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 21:11:46 +0000 Subject: [PATCH 05/13] fix(quantization): clear the KV cache at capture, generalize the changelog Addresses the re-review on #2248. Sanitizing at capture means the cache is never stored in collected_inputs, so it is not pinned for the whole layer loop and not pickled into next_inputs.pt -- _move_to_device recurses into tensors, dicts and lists only, so a transformers Cache was written out as-is and unpickled with weights_only=False on resume. The two replay-site calls stay: they are free on the None fast path and still cover inputs restored from an older checkpoint. Verified a no-op on the current results. The changelog enumerated gptq and awq_lite, which was non-exhaustive in the same direction as the earlier "weight-only recipes are unchanged": local_hessian also runs a second forward_loop to accumulate its Hessian, and awq_clip and smoothquant are activation-driven too. Generalized to the property rather than a list, and points at a fresh checkpoint_dir, since re-running a recipe on a partial directory resumes rather than recalibrates. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- CHANGELOG.rst | 2 +- modelopt/torch/quantization/utils/layerwise_calib.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 248b37051a4..2fe51aace1b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,7 +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 for any recipe with ``layerwise.enable: true``, including the shipped ``nvfp4_experts_only-kv_fp8_layerwise*`` and ``nvfp4_default-kv_none-gptq`` recipes — under ``gptq`` and ``awq_lite`` the exported weights change too, not just activation scales. +- 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) ^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py index 30b4be1ed9f..0e5a1615b65 100644 --- a/modelopt/torch/quantization/utils/layerwise_calib.py +++ b/modelopt/torch/quantization/utils/layerwise_calib.py @@ -54,9 +54,9 @@ class _EarlyStopForwardError(Exception): def _with_empty_kv_cache(kwargs_input: dict) -> dict: """Drop any attention cache, so a replay does not attend over its own earlier writes. - Both replay paths -- the patched ``run`` forward below and ``calib_func``'s replay in - ``layerwise_calibrate`` -- consume the same captured kwargs, so each clears the cache - itself rather than relying on the other not having written to it. + Applied at capture, so the cache is never stored, checkpointed or held across the + layer loop, and again at both replay sites, which costs nothing on the ``None`` fast + path and covers inputs restored from a checkpoint. Not ``Cache.reset()``: that zeroes the key/value tensors but keeps them at full length, leaving the replay attending over an all-zero cache instead of none. @@ -256,7 +256,7 @@ def _patched_forward(self, *args, **kwargs): return output if info.mode == "capture": - info.collected_inputs.append((args, kwargs)) + info.collected_inputs.append((args, _with_empty_kv_cache(kwargs))) raise _EarlyStopForwardError() return self._original_forward(*args, **kwargs) From e03c53b8ad8b7b642c9146faed76abca2a1efb00 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 21:18:19 +0000 Subject: [PATCH 06/13] refactor(quantization): clear the captured KV cache once, at capture Sanitizing at capture makes the two replay-site clears unreachable: measured over a calibration run, capture sees a live cache on 12 of 12 invocations and both replay sites on 0 of 21. The only input that could still carry one is a next_inputs.pt written before this fix, which layerwise -- an experimental feature -- does not promise to resume. So the fix is now one call at the boundary where model-produced kwargs enter stored state, and everything downstream trusts the invariant. model_calib.py no longer needs the helper at all; its diff is just the removal of the broken block. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/quantization/model_calib.py | 3 +-- modelopt/torch/quantization/utils/layerwise_calib.py | 11 ++++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index d48b49d8ad3..4f797b450d7 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -34,7 +34,6 @@ from modelopt.torch.quantization.utils.layerwise_calib import ( LayerActivationCollector, _CheckpointState, - _with_empty_kv_cache, ) from modelopt.torch.utils import print_rank_0, warn_rank_0 from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState, is_master @@ -2126,7 +2125,7 @@ def _set_layer_status(status: str): def _layer_forward_loop(m, _inputs=layer_inputs): for args, kwargs_input in _inputs: - m(*args, **_with_empty_kv_cache(kwargs_input)) + m(*args, **kwargs_input) is_last = layer_idx + 1 >= num_layers diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py index 0e5a1615b65..64833aa0eb1 100644 --- a/modelopt/torch/quantization/utils/layerwise_calib.py +++ b/modelopt/torch/quantization/utils/layerwise_calib.py @@ -52,11 +52,12 @@ class _EarlyStopForwardError(Exception): def _with_empty_kv_cache(kwargs_input: dict) -> dict: - """Drop any attention cache, so a replay does not attend over its own earlier writes. + """Drop the attention cache from captured layer inputs. - Applied at capture, so the cache is never stored, checkpointed or held across the - layer loop, and again at both replay sites, which costs nothing on the ``None`` fast - path and covers inputs restored from a checkpoint. + Captured kwargs are replayed many times -- by ``calib_func``, once per calibration + pass, and by the ``run`` branch below -- so a retained cache would let a layer attend + over the keys and values its own earlier replay wrote. Clearing once at capture keeps + every consumer independent, and keeps the cache out of ``next_inputs.pt``. Not ``Cache.reset()``: that zeroes the key/value tensors but keeps them at full length, leaving the replay attending over an all-zero cache instead of none. @@ -251,7 +252,7 @@ def _patched_forward(self, *args, **kwargs): f"Layer {info.name} is in 'run' mode but has no cached inputs to replay." ) real_args, real_kwargs = info.cached_inputs.popleft() - output = self._original_forward(*real_args, **_with_empty_kv_cache(real_kwargs)) + output = self._original_forward(*real_args, **real_kwargs) info.output_meta = LayerActivationCollector._extract_output_meta(output) return output From 3310a48607bf1590b946826f18a30cc17714330e Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 21:40:17 +0000 Subject: [PATCH 07/13] fix(quantization): match the KV cache by shape, and clear it on the resume path Addresses review on #2248. Name-matching only past_key_values missed two shapes that occur in practice. Remote-code models written against older transformers pass the cache as past_key_value -- ModelOpt itself patches that for Kimi-K2 (speculative/utils.py:546), which is exactly the class of model layerwise calibration exists for -- and a custom parent may pass it positionally. In both cases the clear silently no-opped and layerwise produced wrong-but-plausible amaxes. Duck-typed detection over args and kwargs covers all three, without importing transformers into a core util. Capture is also not the only entry point into stored inputs: get_first_layer_inputs seeds them straight from next_inputs.pt on resume, and _move_to_device passes a Cache through untouched, so a checkpoint written before this fix still carried a live one. A file is a boundary like the model forward is, so it gets the same clear. Adds a test pinning the clear at its call site. The equivalence test cannot cover it: an unsanitized cache only accumulates across replays, which max-calibration's max reduction absorbs on a small model, so amaxes still matched while the invariant was broken -- that test passed with capture unsanitized. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- .../quantization/utils/layerwise_calib.py | 33 +++++++--- .../quantization/test_layerwise_calibrate.py | 61 ++++++++++++++++++- 2 files changed, 83 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py index 64833aa0eb1..4bf8c936356 100644 --- a/modelopt/torch/quantization/utils/layerwise_calib.py +++ b/modelopt/torch/quantization/utils/layerwise_calib.py @@ -51,20 +51,30 @@ class _EarlyStopForwardError(Exception): """Raised to halt the forward pass after capturing layer inputs.""" -def _with_empty_kv_cache(kwargs_input: dict) -> dict: - """Drop the attention cache from captured layer inputs. +def _is_kv_cache(obj: Any) -> bool: + """Duck-typed ``transformers.Cache``, so a core util need not import transformers.""" + return hasattr(obj, "update") and hasattr(obj, "get_seq_length") - Captured kwargs are replayed many times -- by ``calib_func``, once per calibration - pass, and by the ``run`` branch below -- so a retained cache would let a layer attend - over the keys and values its own earlier replay wrote. Clearing once at capture keeps - every consumer independent, and keeps the cache out of ``next_inputs.pt``. + +def _with_empty_kv_cache(args: tuple, kwargs_input: dict) -> tuple[tuple, dict]: + """Drop the attention cache from a captured layer input. + + Layer inputs are replayed many times -- by ``calib_func``, once per calibration pass, + and by the ``run`` branch below -- so a retained cache would let a layer attend over + the keys and values its own earlier replay wrote. Cleared where inputs enter stored + state (capture, and restore from a checkpoint), so every consumer can trust it. + + Matched by shape rather than by name: the cache reaches a layer as ``past_key_values`` + on HF-native models but as ``past_key_value`` on remote-code models written against + older transformers (Kimi-K2, for one), and a custom parent may pass it positionally. Not ``Cache.reset()``: that zeroes the key/value tensors but keeps them at full length, leaving the replay attending over an all-zero cache instead of none. """ - if kwargs_input.get("past_key_values") is None: - return kwargs_input - return {**kwargs_input, "past_key_values": None} + return ( + tuple(None if _is_kv_cache(a) else a for a in args), + {k: (None if _is_kv_cache(v) else v) for k, v in kwargs_input.items()}, + ) @dataclass @@ -257,7 +267,7 @@ def _patched_forward(self, *args, **kwargs): return output if info.mode == "capture": - info.collected_inputs.append((args, _with_empty_kv_cache(kwargs))) + info.collected_inputs.append(_with_empty_kv_cache(args, kwargs)) raise _EarlyStopForwardError() return self._original_forward(*args, **kwargs) @@ -451,6 +461,9 @@ def get_first_layer_inputs( for i in range(start_layer): self._swap_to_dummy(i) layer = self._decoder_layers[start_layer] + # Second entry point into stored inputs, and the only one not fed by capture: + # a checkpoint written before the cache was cleared still carries a live one. + resumed_inputs = [_with_empty_kv_cache(a, kw) for a, kw in resumed_inputs] layer._layerwise_calib.collected_inputs = resumed_inputs layer._layerwise_calib.mode = "original" return resumed_inputs diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py index 759640d16af..5f2dcc96d39 100644 --- a/tests/unit/torch/quantization/test_layerwise_calibrate.py +++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py @@ -23,11 +23,17 @@ import torch import torch.nn as nn from _test_utils.torch.transformers_models import get_tiny_llama +from transformers.cache_utils import DynamicCache import modelopt.torch.quantization as mtq from modelopt.torch.quantization.model_calib import layerwise_calibrate from modelopt.torch.quantization.nn import TensorQuantizer -from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector, _SkipLayer +from modelopt.torch.quantization.utils.layerwise_calib import ( + LayerActivationCollector, + _is_kv_cache, + _SkipLayer, + _with_empty_kv_cache, +) class _DecoderBlock(nn.Module): @@ -1080,6 +1086,59 @@ def crashing_torch_save(obj, path, *args, **kwargs): assert manifest["last_completed_layer"] == 1, f"manifest leaked mid-window state: {manifest}" +def test_capture_stores_no_kv_cache(): + """Captured layer inputs must carry no cache, so every replay of them is independent. + + Pins the clear at its call site. The equivalence test above cannot: an unsanitized + cache merely *accumulates* across replays, which max-calibration's max reduction + absorbs on a small model, so amaxes can still match while the invariant is broken. + """ + model = get_tiny_llama(num_hidden_layers=3).eval() + collector = LayerActivationCollector(model) + collector._patch_all_layers(decoder_layers=model.model.layers) + try: + captured = collector.get_input_activations( + model.model.layers[0], lambda m: m(torch.randint(0, 32, (2, 8))) + ) + finally: + collector._unpatch_all_layers() + + assert captured, "nothing was captured" + live = [ + k + for args, kwargs in captured + for k, v in [*enumerate(args), *kwargs.items()] + if _is_kv_cache(v) + ] + assert not live, f"captured inputs still hold a KV cache at {live}" + + +def test_with_empty_kv_cache_matches_by_shape_not_by_name(): + """The cache must be cleared however it reaches the layer. + + HF-native layers take it as ``past_key_values``, but remote-code models written + against older transformers use ``past_key_value`` (Kimi-K2 -- see + ``modelopt/torch/speculative/utils.py``), and a custom parent may pass it + positionally. Name-matching only the plural form would silently no-op on those. + """ + cache = DynamicCache() + hidden = torch.randn(1, 4) + + for args, kwargs in ( + ((), {"past_key_values": cache}), + ((), {"past_key_value": cache}), + ((hidden, cache), {}), + ): + out_args, out_kwargs = _with_empty_kv_cache(args, kwargs) + assert not any(_is_kv_cache(a) for a in out_args) + assert not any(_is_kv_cache(v) for v in out_kwargs.values()) + + # Non-cache values pass through untouched. + out_args, out_kwargs = _with_empty_kv_cache((hidden,), {"attention_mask": None}) + assert out_args[0] is hidden + assert out_kwargs == {"attention_mask": None} + + 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. From 519a78963bf2e6689938f1252a4937c742665887 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 21:44:52 +0000 Subject: [PATCH 08/13] docs: compact the KV-cache comments and docstrings Cuts the helper docstring from fourteen lines of prose to two facts a reader cannot infer from the code: why a replay must not keep the cache, and why Cache.reset() is not the way to drop it. The name-vs-shape rationale moves onto _is_kv_cache, which is what does the matching. Same for the three test docstrings and the resume-site comment. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- .../quantization/utils/layerwise_calib.py | 24 ++++++--------- .../quantization/test_layerwise_calibrate.py | 29 ++++++++----------- 2 files changed, 21 insertions(+), 32 deletions(-) diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py index 4bf8c936356..26bc277eb09 100644 --- a/modelopt/torch/quantization/utils/layerwise_calib.py +++ b/modelopt/torch/quantization/utils/layerwise_calib.py @@ -52,24 +52,19 @@ class _EarlyStopForwardError(Exception): def _is_kv_cache(obj: Any) -> bool: - """Duck-typed ``transformers.Cache``, so a core util need not import transformers.""" + """Duck-typed ``transformers.Cache``, avoiding a transformers import here. + + Matched by shape, not name: the keyword is ``past_key_values`` on HF-native layers + but ``past_key_value`` on older remote-code ones (Kimi-K2), and may be positional. + """ return hasattr(obj, "update") and hasattr(obj, "get_seq_length") def _with_empty_kv_cache(args: tuple, kwargs_input: dict) -> tuple[tuple, dict]: - """Drop the attention cache from a captured layer input. - - Layer inputs are replayed many times -- by ``calib_func``, once per calibration pass, - and by the ``run`` branch below -- so a retained cache would let a layer attend over - the keys and values its own earlier replay wrote. Cleared where inputs enter stored - state (capture, and restore from a checkpoint), so every consumer can trust it. - - Matched by shape rather than by name: the cache reaches a layer as ``past_key_values`` - on HF-native models but as ``past_key_value`` on remote-code models written against - older transformers (Kimi-K2, for one), and a custom parent may pass it positionally. + """Strip the attention cache, so a replay cannot attend over its own earlier writes. - Not ``Cache.reset()``: that zeroes the key/value tensors but keeps them at full - length, leaving the replay attending over an all-zero cache instead of none. + Not ``Cache.reset()``: it zeroes the keys and values but keeps them at full length, + so the replay attends over an all-zero cache instead of none. """ return ( tuple(None if _is_kv_cache(a) else a for a in args), @@ -461,8 +456,7 @@ def get_first_layer_inputs( for i in range(start_layer): self._swap_to_dummy(i) layer = self._decoder_layers[start_layer] - # Second entry point into stored inputs, and the only one not fed by capture: - # a checkpoint written before the cache was cleared still carries a live one. + # Not fed by capture: an older checkpoint can still carry a live cache. resumed_inputs = [_with_empty_kv_cache(a, kw) for a, kw in resumed_inputs] layer._layerwise_calib.collected_inputs = resumed_inputs layer._layerwise_calib.mode = "original" diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py index 5f2dcc96d39..60e3b559f76 100644 --- a/tests/unit/torch/quantization/test_layerwise_calibrate.py +++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py @@ -838,15 +838,11 @@ def fwd(m): def test_layerwise_replay_does_not_attend_over_its_own_kv_cache(): - """A layer replayed during calibration must not see the keys and values its - own earlier run wrote. - - The equivalence test above uses a model with no attention cache, so it cannot - reach this: only the captured ``past_key_values`` makes a replay stateful. - With the cache left in place, ``o_proj`` saw an all-zero attention output on - every layer but the last (the one with no preceding capture pass) and its - input amax collapsed to exactly 0.0, while ``down_proj`` picked up a - plausible but wrong value from the residual alone. + """Same equivalence as above, on a model that actually has a KV cache. + + The toy model used above has none, so it cannot go stale -- which is why a + replay attending over its own writes shipped, collapsing ``o_proj``'s input + amax to 0.0 on every layer but the last. """ calib_data = [torch.randint(0, 32, (2, 8)) for _ in range(2)] @@ -1087,11 +1083,11 @@ def crashing_torch_save(obj, path, *args, **kwargs): def test_capture_stores_no_kv_cache(): - """Captured layer inputs must carry no cache, so every replay of them is independent. + """Captured inputs must carry no cache. - Pins the clear at its call site. The equivalence test above cannot: an unsanitized - cache merely *accumulates* across replays, which max-calibration's max reduction - absorbs on a small model, so amaxes can still match while the invariant is broken. + Asserted structurally because the amax comparison cannot see it: a retained cache + only accumulates across replays, and max calibration's max absorbs that on a small + model, so amaxes match while the invariant is broken. """ model = get_tiny_llama(num_hidden_layers=3).eval() collector = LayerActivationCollector(model) @@ -1116,10 +1112,9 @@ def test_capture_stores_no_kv_cache(): def test_with_empty_kv_cache_matches_by_shape_not_by_name(): """The cache must be cleared however it reaches the layer. - HF-native layers take it as ``past_key_values``, but remote-code models written - against older transformers use ``past_key_value`` (Kimi-K2 -- see - ``modelopt/torch/speculative/utils.py``), and a custom parent may pass it - positionally. Name-matching only the plural form would silently no-op on those. + Kimi-K2's remote code passes it as ``past_key_value`` (see + ``modelopt/torch/speculative/utils.py``), so matching only the plural keyword + would silently no-op on a model layerwise calibration exists for. """ cache = DynamicCache() hidden = torch.randn(1, 4) From 555d6bc06811f18597b420e656c755f199287b79 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 22:15:47 +0000 Subject: [PATCH 09/13] test: keep only the KV-cache tests that can fail Drops the layerwise-vs-non-layerwise equivalence test added earlier. It duplicated test_layerwise_no_qdq_matches_sequential_amax's property, and it could not fail for the reason it was written: without the clear the cache merely accumulates across replays, which max calibration's max absorbs, so the amaxes still matched. The structural test covers a reset()-style regression too, since a zeroed cache is still a cache in the captured inputs. What remains is three tests, one per call site, each verified to fail when that site's clear is removed: capture stores no cache, resume strips one left in a checkpoint, and the helper matches a cache by shape rather than by keyword. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- .../quantization/test_layerwise_calibrate.py | 61 +++++++------------ 1 file changed, 22 insertions(+), 39 deletions(-) diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py index 60e3b559f76..1080e79a4db 100644 --- a/tests/unit/torch/quantization/test_layerwise_calibrate.py +++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py @@ -837,45 +837,6 @@ def fwd(m): _assert_amax_close(_collect_amax(model_lw), seq_amax, "layerwise vs sequential") -def test_layerwise_replay_does_not_attend_over_its_own_kv_cache(): - """Same equivalence as above, on a model that actually has a KV cache. - - The toy model used above has none, so it cannot go stale -- which is why a - replay attending over its own writes shipped, collapsing ``o_proj``'s input - amax to 0.0 on every layer but the last. - """ - calib_data = [torch.randint(0, 32, (2, 8)) for _ in range(2)] - - def fwd(m): - for batch in calib_data: - m(batch) - - def calibrate(algorithm): - torch.manual_seed(0) - model = get_tiny_llama(num_hidden_layers=4).eval() - cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) - cfg["algorithm"] = algorithm - mtq.quantize(model, cfg, forward_loop=fwd) - return model - - sequential = calibrate({"method": "max"}) - layerwise = calibrate({"method": "max", "layerwise": {"enable": True}}) - - expected = _collect_amax(sequential) - assert expected, "sequential calibration populated no amax values" - layerwise_amax = _collect_amax(layerwise) - _assert_amax_close(layerwise_amax, expected, "layerwise vs sequential (KV cache)") - - # Pinned separately from the comparison: a future regression that made both - # paths collapse to zero would still satisfy the equality above. - collapsed = [ - name - for name, amax in layerwise_amax.items() - if name.endswith("input_quantizer") and not torch.count_nonzero(amax) - ] - assert not collapsed, f"activation amax collapsed to zero: {collapsed}" - - def test_layerwise_no_qdq_captures_inputs_before_calib_func_mutates_weights(monkeypatch): """A destructive ``calib_func`` (zeros weights) must not affect what is captured for downstream layers under ``qdq_from_prev=False`` — otherwise @@ -1109,6 +1070,28 @@ def test_capture_stores_no_kv_cache(): assert not live, f"captured inputs still hold a KV cache at {live}" +def test_resume_clears_a_kv_cache_left_in_the_checkpoint(): + """Inputs restored from ``next_inputs.pt`` must be cleared too. + + ``_move_to_device`` passes a ``Cache`` through untouched, so a checkpoint written + before the clear still holds a live one, and resume is the one path into stored + inputs that capture does not feed. + """ + model = get_tiny_llama(num_hidden_layers=3).eval() + collector = LayerActivationCollector(model) + collector._patch_all_layers(decoder_layers=model.model.layers) + try: + restored = [((torch.randn(1, 4, 32),), {"past_key_values": DynamicCache()})] + returned = collector.get_first_layer_inputs(0, restored, forward_loop=None) + stored = model.model.layers[0]._layerwise_calib.collected_inputs + finally: + collector._unpatch_all_layers() + + for label, entries in (("returned", returned), ("stored", stored)): + live = [v for _, kwargs in entries for v in kwargs.values() if _is_kv_cache(v)] + assert not live, f"{label} resume inputs still hold a KV cache" + + def test_with_empty_kv_cache_matches_by_shape_not_by_name(): """The cache must be cleared however it reaches the layer. From 1e721a01b74da278ba3ca61e0060a104c0ba3b05 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 22:26:45 +0000 Subject: [PATCH 10/13] test: fold the two structural KV-cache tests into one parametrized case capture and resume assert the same invariant -- stored layer inputs hold no cache -- at the two entry points into stored state, so they are one test over an entry_point parameter. Both parameters still fail independently when their own call site's clear is removed. Also asserts against transformers' Cache type rather than _is_kv_cache. Using the predicate under test made the assertions self-referential: breaking _is_kv_cache to return False left the shape-matching test passing, because the check it makes went false too. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- .../quantization/test_layerwise_calibrate.py | 63 +++++++------------ 1 file changed, 22 insertions(+), 41 deletions(-) diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py index 1080e79a4db..f59e1571091 100644 --- a/tests/unit/torch/quantization/test_layerwise_calibrate.py +++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py @@ -23,14 +23,13 @@ import torch import torch.nn as nn from _test_utils.torch.transformers_models import get_tiny_llama -from transformers.cache_utils import DynamicCache +from transformers.cache_utils import Cache, DynamicCache import modelopt.torch.quantization as mtq from modelopt.torch.quantization.model_calib import layerwise_calibrate from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.quantization.utils.layerwise_calib import ( LayerActivationCollector, - _is_kv_cache, _SkipLayer, _with_empty_kv_cache, ) @@ -1043,53 +1042,35 @@ def crashing_torch_save(obj, path, *args, **kwargs): assert manifest["last_completed_layer"] == 1, f"manifest leaked mid-window state: {manifest}" -def test_capture_stores_no_kv_cache(): - """Captured inputs must carry no cache. +@pytest.mark.parametrize("entry_point", ["capture", "resume"]) +def test_stored_layer_inputs_never_hold_a_kv_cache(entry_point): + """Neither way inputs enter stored state may leave a cache on them. - Asserted structurally because the amax comparison cannot see it: a retained cache - only accumulates across replays, and max calibration's max absorbs that on a small - model, so amaxes match while the invariant is broken. - """ - model = get_tiny_llama(num_hidden_layers=3).eval() - collector = LayerActivationCollector(model) - collector._patch_all_layers(decoder_layers=model.model.layers) - try: - captured = collector.get_input_activations( - model.model.layers[0], lambda m: m(torch.randint(0, 32, (2, 8))) - ) - finally: - collector._unpatch_all_layers() - - assert captured, "nothing was captured" - live = [ - k - for args, kwargs in captured - for k, v in [*enumerate(args), *kwargs.items()] - if _is_kv_cache(v) - ] - assert not live, f"captured inputs still hold a KV cache at {live}" + Stored inputs are replayed many times, so a retained cache lets a layer attend over + its own earlier writes. Asserted structurally because an amax comparison cannot see + it: a retained cache only accumulates across replays, and max calibration's max + absorbs that on a small model, so amaxes match while the invariant is broken. - -def test_resume_clears_a_kv_cache_left_in_the_checkpoint(): - """Inputs restored from ``next_inputs.pt`` must be cleared too. - - ``_move_to_device`` passes a ``Cache`` through untouched, so a checkpoint written - before the clear still holds a live one, and resume is the one path into stored - inputs that capture does not feed. + ``resume`` is the entry point capture does not feed -- ``_move_to_device`` passes a + ``Cache`` through untouched, so a checkpoint predating the clear still holds one. """ model = get_tiny_llama(num_hidden_layers=3).eval() collector = LayerActivationCollector(model) collector._patch_all_layers(decoder_layers=model.model.layers) try: - restored = [((torch.randn(1, 4, 32),), {"past_key_values": DynamicCache()})] - returned = collector.get_first_layer_inputs(0, restored, forward_loop=None) - stored = model.model.layers[0]._layerwise_calib.collected_inputs + if entry_point == "capture": + stored = collector.get_input_activations( + model.model.layers[0], lambda m: m(torch.randint(0, 32, (2, 8))) + ) + else: + restored = [((torch.randn(1, 4, 32),), {"past_key_values": DynamicCache()})] + stored = collector.get_first_layer_inputs(0, restored, forward_loop=None) finally: collector._unpatch_all_layers() - for label, entries in (("returned", returned), ("stored", stored)): - live = [v for _, kwargs in entries for v in kwargs.values() if _is_kv_cache(v)] - assert not live, f"{label} resume inputs still hold a KV cache" + assert stored, f"{entry_point} produced no inputs" + live = [v for args, kwargs in stored for v in (*args, *kwargs.values()) if isinstance(v, Cache)] + assert not live, f"{entry_point} inputs still hold a KV cache" def test_with_empty_kv_cache_matches_by_shape_not_by_name(): @@ -1108,8 +1089,8 @@ def test_with_empty_kv_cache_matches_by_shape_not_by_name(): ((hidden, cache), {}), ): out_args, out_kwargs = _with_empty_kv_cache(args, kwargs) - assert not any(_is_kv_cache(a) for a in out_args) - assert not any(_is_kv_cache(v) for v in out_kwargs.values()) + assert not any(isinstance(a, Cache) for a in out_args) + assert not any(isinstance(v, Cache) for v in out_kwargs.values()) # Non-cache values pass through untouched. out_args, out_kwargs = _with_empty_kv_cache((hidden,), {"attention_mask": None}) From c3300d821077d84192d9a5b01741b0b94790e88c Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 25 Aug 2026 23:08:57 +0000 Subject: [PATCH 11/13] fix(quantization): disable the KV cache for layerwise calibration Layerwise replays each layer's captured inputs several times, and those inputs carried the model's past_key_values, so a layer attended over the keys and values its own earlier replay wrote. Everything downstream of attention in the layer was then calibrated against a zeroed attention output on all but the last layer. The code tried to prevent exactly this with Cache.reset(), but that "resets the cache values while preserving the objects" -- it zeroes the key/value tensors and leaves them at full length, so the replay attended over a same-length all-zero cache instead of no cache. Sliding-window models raised a shape mismatch instead, since the next update then doubled kv_len. Calibration never reads a cache, so rather than clearing it per input, don't build one: wrap the layerwise loop in the existing _disable_use_cache. That helper already handles nested multimodal configs and configs that never assign the attribute, and its docstring already names this failure class for hybrid Mamba/attention models. It also drops peak calibration memory, since the cache was allocated and never read. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/quantization/model_calib.py | 106 +++++++++--------- .../quantization/utils/layerwise_calib.py | 25 +---- .../quantization/test_layerwise_calibrate.py | 97 +++++++--------- 3 files changed, 99 insertions(+), 129 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 4f797b450d7..91dc6af9e60 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/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,61 +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 - - # 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: - m(*args, **kwargs_input) + # 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 + ) - is_last = layer_idx + 1 >= num_layers + # Bootstrap: get first layer's inputs (or use resumed inputs). + layer_inputs = input_getter.get_first_layer_inputs( + start_layer, resumed_inputs, forward_loop + ) - 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) diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py index 26bc277eb09..070ee521cd5 100644 --- a/modelopt/torch/quantization/utils/layerwise_calib.py +++ b/modelopt/torch/quantization/utils/layerwise_calib.py @@ -51,27 +51,6 @@ class _EarlyStopForwardError(Exception): """Raised to halt the forward pass after capturing layer inputs.""" -def _is_kv_cache(obj: Any) -> bool: - """Duck-typed ``transformers.Cache``, avoiding a transformers import here. - - Matched by shape, not name: the keyword is ``past_key_values`` on HF-native layers - but ``past_key_value`` on older remote-code ones (Kimi-K2), and may be positional. - """ - return hasattr(obj, "update") and hasattr(obj, "get_seq_length") - - -def _with_empty_kv_cache(args: tuple, kwargs_input: dict) -> tuple[tuple, dict]: - """Strip the attention cache, so a replay cannot attend over its own earlier writes. - - Not ``Cache.reset()``: it zeroes the keys and values but keeps them at full length, - so the replay attends over an all-zero cache instead of none. - """ - return ( - tuple(None if _is_kv_cache(a) else a for a in args), - {k: (None if _is_kv_cache(v) else v) for k, v in kwargs_input.items()}, - ) - - @dataclass class _LayerCalibState: """Mutable per-layer state used during layerwise calibration. @@ -262,7 +241,7 @@ def _patched_forward(self, *args, **kwargs): return output if info.mode == "capture": - info.collected_inputs.append(_with_empty_kv_cache(args, kwargs)) + info.collected_inputs.append((args, kwargs)) raise _EarlyStopForwardError() return self._original_forward(*args, **kwargs) @@ -456,8 +435,6 @@ def get_first_layer_inputs( for i in range(start_layer): self._swap_to_dummy(i) layer = self._decoder_layers[start_layer] - # Not fed by capture: an older checkpoint can still carry a live cache. - resumed_inputs = [_with_empty_kv_cache(a, kw) for a, kw in resumed_inputs] layer._layerwise_calib.collected_inputs = resumed_inputs layer._layerwise_calib.mode = "original" return resumed_inputs diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py index f59e1571091..e58e18cad3c 100644 --- a/tests/unit/torch/quantization/test_layerwise_calibrate.py +++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py @@ -22,17 +22,17 @@ import pytest import torch import torch.nn as nn -from _test_utils.torch.transformers_models import get_tiny_llama -from transformers.cache_utils import Cache, DynamicCache +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 from modelopt.torch.quantization.nn import TensorQuantizer -from modelopt.torch.quantization.utils.layerwise_calib import ( - LayerActivationCollector, - _SkipLayer, - _with_empty_kv_cache, -) +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector, _SkipLayer class _DecoderBlock(nn.Module): @@ -1042,60 +1042,47 @@ def crashing_torch_save(obj, path, *args, **kwargs): assert manifest["last_completed_layer"] == 1, f"manifest leaked mid-window state: {manifest}" -@pytest.mark.parametrize("entry_point", ["capture", "resume"]) -def test_stored_layer_inputs_never_hold_a_kv_cache(entry_point): - """Neither way inputs enter stored state may leave a cache on them. +@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. - Stored inputs are replayed many times, so a retained cache lets a layer attend over - its own earlier writes. Asserted structurally because an amax comparison cannot see - it: a retained cache only accumulates across replays, and max calibration's max - absorbs that on a small model, so amaxes match while the invariant is broken. + 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. - ``resume`` is the entry point capture does not feed -- ``_move_to_device`` passes a - ``Cache`` through untouched, so a checkpoint predating the clear still holds one. + 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 = get_tiny_llama(num_hidden_layers=3).eval() - collector = LayerActivationCollector(model) - collector._patch_all_layers(decoder_layers=model.model.layers) + 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: - if entry_point == "capture": - stored = collector.get_input_activations( - model.model.layers[0], lambda m: m(torch.randint(0, 32, (2, 8))) - ) - else: - restored = [((torch.randn(1, 4, 32),), {"past_key_values": DynamicCache()})] - stored = collector.get_first_layer_inputs(0, restored, forward_loop=None) + 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: - collector._unpatch_all_layers() - - assert stored, f"{entry_point} produced no inputs" - live = [v for args, kwargs in stored for v in (*args, *kwargs.values()) if isinstance(v, Cache)] - assert not live, f"{entry_point} inputs still hold a KV cache" + for h in handles: + h.remove() - -def test_with_empty_kv_cache_matches_by_shape_not_by_name(): - """The cache must be cleared however it reaches the layer. - - Kimi-K2's remote code passes it as ``past_key_value`` (see - ``modelopt/torch/speculative/utils.py``), so matching only the plural keyword - would silently no-op on a model layerwise calibration exists for. - """ - cache = DynamicCache() - hidden = torch.randn(1, 4) - - for args, kwargs in ( - ((), {"past_key_values": cache}), - ((), {"past_key_value": cache}), - ((hidden, cache), {}), - ): - out_args, out_kwargs = _with_empty_kv_cache(args, kwargs) - assert not any(isinstance(a, Cache) for a in out_args) - assert not any(isinstance(v, Cache) for v in out_kwargs.values()) - - # Non-cache values pass through untouched. - out_args, out_kwargs = _with_empty_kv_cache((hidden,), {"attention_mask": None}) - assert out_args[0] is hidden - assert out_kwargs == {"attention_mask": None} + 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): From dd3e40cb4f7487a26719896e66b06e8a57088692 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Wed, 26 Aug 2026 00:02:15 +0000 Subject: [PATCH 12/13] fix(quantization): refuse a KV cache at capture, and scope the changelog Disabling use_cache stops the model building a cache, but a caller can still hand one in: use_cache resolves to the caller's value when passed, and past_key_values can be passed directly -- examples/alpamayo/quantize.py does both. Layerwise replays captured inputs, so either route silently reproduces the bug. Capture now refuses a cache instead. The changelog also overstated the blast radius. create_forward_loop already wraps its body in _disable_use_cache, so the shipped hf_ptq path never built a cache to begin with: measured on Qwen2.5-1.5B with nvfp4_default-kv_none-gptq and the fix simulated away, the shipped loop puts 0 caches on a layer and 0 in the checkpoint, while a bare user loop puts 222 and 54. Telling users to re-run the shipped recipes was wrong; the exposure is custom forward loops via mtq.quantize. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- CHANGELOG.rst | 2 +- .../quantization/utils/layerwise_calib.py | 12 ++++++ .../quantization/test_layerwise_calibrate.py | 37 +++++++++---------- 3 files changed, 31 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2fe51aace1b..df78cd6dd58 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,7 +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. +- Fix layerwise calibration miscalibrating every module downstream of attention within a decoder layer when the calibration forward loop left KV caching enabled; models using an explicit sliding-window mask raised a shape mismatch instead. ``create_forward_loop`` already disabled caching, so ``examples/hf_ptq`` and the shipped recipes were unaffected — re-run calibration from a fresh ``layerwise.checkpoint_dir`` only if you passed your own ``forward_loop`` to ``mtq.quantize``, where ``gptq``, ``awq_lite``, ``awq_clip``, ``local_hessian`` and ``smoothquant`` change exported weights too, not just activation scales. 0.46 (2026-08-17) ^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py index 070ee521cd5..087f4e24fb1 100644 --- a/modelopt/torch/quantization/utils/layerwise_calib.py +++ b/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() diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py index e58e18cad3c..2659c3042a9 100644 --- a/tests/unit/torch/quantization/test_layerwise_calibrate.py +++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py @@ -1047,22 +1047,21 @@ def crashing_torch_save(obj, path, *args, **kwargs): [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. +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" - layers = LayerActivationCollector.get_decoder_layers(model) - seen = [] handles = [ layer.register_forward_pre_hook( @@ -1071,19 +1070,19 @@ def test_layerwise_calibration_builds_no_kv_cache(factory): ), with_kwargs=True, ) - for layer in layers + for layer in LayerActivationCollector.get_decoder_layers(model) ] 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)))) + 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 during calibration" + 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 From 0e351201fec8d2c261a0986b2267460dab89598a Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Wed, 26 Aug 2026 18:01:29 +0000 Subject: [PATCH 13/13] docs(changelog): drop the layerwise KV-cache entry No shipped recipe was affected. create_forward_loop already disables caching, so the bug needed a calibration loop that does not -- verified by running hf_ptq.py itself on main with full NVFP4 W4A4 including o_proj, which exported 24 o_proj input_scale values with none collapsed. The remaining exposure is a caller-supplied forward_loop, which no released checkpoint went through, so there is nothing for a user to act on. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- CHANGELOG.rst | 1 - 1 file changed, 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index df78cd6dd58..6687ebd31ea 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,7 +43,6 @@ 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 when the calibration forward loop left KV caching enabled; models using an explicit sliding-window mask raised a shape mismatch instead. ``create_forward_loop`` already disabled caching, so ``examples/hf_ptq`` and the shipped recipes were unaffected — re-run calibration from a fresh ``layerwise.checkpoint_dir`` only if you passed your own ``forward_loop`` to ``mtq.quantize``, where ``gptq``, ``awq_lite``, ``awq_clip``, ``local_hessian`` and ``smoothquant`` change exported weights too, not just activation scales. 0.46 (2026-08-17) ^^^^^^^^^^^^^^^^^