Fix stale GGML fake-quant cache after version-bypassing optimizer steps (#2857)

CHANGELOG.rst

@@ -32,6 +32,8 @@ Changelog
 
 **Bug Fixes**
 
+- Fix GGML fake-quant silently reusing stale codes after optimizer steps that bypass the PyTorch
+  version counter, such as bitsandbytes 8-bit optimizers.
 - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``.
 - Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source.
 - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11.

modelopt/torch/quantization/ggml/common.py

@@ -57,6 +57,7 @@ class _PackedWeightCache:
 
     base_ref: weakref.ReferenceType
     input_key: _CacheKey
+    input_fingerprint: torch.Tensor
     format_name: str
     block_chunk_size: int
     packed_weights: torch.Tensor
@@ -89,6 +90,16 @@ def _input_cache_key(inputs: torch.Tensor) -> _CacheKey | None:
     )
 
 
+def _input_fingerprint(inputs: torch.Tensor) -> torch.Tensor:
+    """Return a small device-resident fingerprint of the input's bit representation."""
+    bits = inputs.contiguous().view(_SIGNED_INT_OF_SIZE[inputs.element_size()]).reshape(-1)
+    sample = bits[:: max(1, bits.numel() // 1024)]
+    summary = torch.stack(
+        (bits.sum(), (bits * bits).sum(), sample.sum(), (sample * sample).sum())
+    )
+    return torch.cat((summary, sample))
+
+
 def _is_contiguous(shape: tuple[int, ...], stride: tuple[int, ...]) -> bool:
     expected = 1
     for size, step in zip(reversed(shape), reversed(stride)):
@@ -109,16 +120,44 @@ def _same_elements(key: _CacheKey, other: _CacheKey) -> bool:
     )
 
 
+def _same_input_layout(key: _CacheKey, other: _CacheKey) -> bool:
+    return (
+        key.data_ptr,
+        key.dtype,
+        key.device,
+        key.shape,
+        key.stride,
+    ) == (
+        other.data_ptr,
+        other.dtype,
+        other.device,
+        other.shape,
+        other.stride,
+    )
+
+
+def _same_contiguous_layout(key: _CacheKey, other: _CacheKey) -> bool:
+    return (
+        key.dtype == other.dtype
+        and key.device == other.device
+        and math.prod(key.shape) == math.prod(other.shape)
+        and _is_contiguous(key.shape, key.stride)
+        and _is_contiguous(other.shape, other.stride)
+    )
+
+
 def _matching_cache(
     quantizer, inputs: torch.Tensor, format_name: str, *, any_contiguous_view: bool = False
 ) -> _PackedWeightCache | None:
     """The quantizer's cached payload if it was packed from ``inputs``, else None.
 
     It matches only while the weight is unchanged since it was packed: the key carries the
-    tensor's version counter, so an in-place update misses. With ``any_contiguous_view`` it also
-    matches another contiguous view of the same values: TensorQuantizer hands fake quant the
-    weight reshaped into 256-value blocks, so export, which holds the weight itself, sees a
-    different shape over the same elements in the same order, and so the same GGML blocks.
+    tensor's version counter, so an in-place update misses. For trainable weights, it additionally
+    compares a device-resident content fingerprint because optimizers may update them without
+    incrementing that counter. With ``any_contiguous_view`` it also matches another
+    contiguous view of the same values: TensorQuantizer hands fake quant the weight reshaped into
+    256-value blocks, so export, which holds the weight itself, sees a different shape over the
+    same elements in the same order, and so the same GGML blocks.
     """
     cache = getattr(quantizer, "_quantizer_cache", None)
     input_key = _input_cache_key(inputs)
@@ -129,8 +168,28 @@ def _matching_cache(
         and cache.format_name == format_name
     ):
         return None
-    if cache.input_key == input_key or (
-        any_contiguous_view and _same_elements(cache.input_key, input_key)
+    # ``inputs`` is often a view created in the forward; under no_grad its requires_grad is False
+    # even for a trainable parameter, so read it from the parameter itself.
+    trainable_weight = _cache_base(inputs).requires_grad
+    if (
+        (
+            trainable_weight
+            and (
+                _same_input_layout(cache.input_key, input_key)
+                or (
+                    any_contiguous_view
+                    and _same_contiguous_layout(cache.input_key, input_key)
+                )
+            )
+            and torch.equal(cache.input_fingerprint, _input_fingerprint(inputs))
+        )
+        or (
+            not trainable_weight
+            and (
+                cache.input_key == input_key
+                or (any_contiguous_view and _same_elements(cache.input_key, input_key))
+            )
+        )
     ):
         return cache
     return None
@@ -209,6 +268,7 @@ def fake_quantize_with_cache(
             quantizer._quantizer_cache = _PackedWeightCache(
                 base_ref=weakref.ref(_cache_base(inputs)),
                 input_key=input_key,
+                input_fingerprint=_input_fingerprint(inputs),
                 format_name=format_name,
                 block_chunk_size=block_chunk_size,
                 packed_weights=packed_weights,

tests/unit/torch/quantization/test_ggml_backend.py

@@ -154,6 +154,56 @@ def counted_quantize(*args, **kwargs):
     assert call_count == 3
 
 
+@pytest.mark.parametrize("num_bits", ["iq1_s", "q8_0"])
+def test_ggml_cache_invalidates_after_version_bypassing_weight_update(num_bits):
+    torch.manual_seed(1234)
+    model = torch.nn.Linear(256, 64, bias=False)
+    inputs = torch.randn(2, 256)
+    config = {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*weight_quantizer",
+                "cfg": {
+                    "num_bits": num_bits,
+                    "block_sizes": {-1: 256},
+                    "backend": "ggml",
+                },
+                "enable": True,
+            },
+        ],
+        "algorithm": "max",
+    }
+
+    mtq.quantize(model, config, forward_loop=lambda module: module(inputs))
+    quantizer = model.weight_quantizer
+    model(inputs)
+    before = quantizer._quantizer_cache.packed_weights.clone()
+    version = model.weight._version
+    stale_output = model(inputs).detach()
+
+    model.weight.data.add_(0.01)
+    assert model.weight._version == version
+
+    with torch.no_grad():
+        model(inputs)
+    after = quantizer._quantizer_cache.packed_weights.clone()
+    assert not torch.equal(after, before)
+
+    quantizer._quantizer_cache = None
+    fresh_output = model(inputs)
+    fresh = quantizer._quantizer_cache.packed_weights
+    assert torch.equal(after, fresh)
+    assert not torch.equal(stale_output, fresh_output)
+    assert torch.equal(model(inputs), fresh_output)
+
+    with torch.no_grad():
+        model.weight.add_(0.01)
+    assert model.weight._version != version
+    model(inputs)
+    assert not torch.equal(quantizer._quantizer_cache.packed_weights, fresh)
+
+
 def test_narrow_to_float32_matches_the_cuda_load_float_policy():
     """Non-finite elements become zero; finite out-of-range elements saturate."""
     largest = torch.finfo(torch.float32).max