diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1fc54679a9a..2bc8282b2e9 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -45,6 +45,8 @@ Changelog - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root. - Fix GPT-OSS draft training with pre-normalization teacher hidden states by enabling the target model's final normalization. - 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 GGML fake-quant silently reusing stale codes after optimizer steps that bypass the PyTorch + version counter, such as bitsandbytes 8-bit optimizers. - 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 GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU. - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index 59a24cbdb14..a18de2b972f 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/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,26 @@ def _input_cache_key(inputs: torch.Tensor) -> _CacheKey | None: ) +def _input_fingerprint(inputs: torch.Tensor) -> torch.Tensor: + """Return a position-sensitive device-resident fingerprint of every input element.""" + bits = inputs.contiguous().view(_SIGNED_INT_OF_SIZE[inputs.element_size()]).reshape(-1) + summary = torch.zeros(4, dtype=torch.int64, device=bits.device) + for start in range(0, bits.numel(), _ROW_KEY_CHUNK): + chunk = bits[start : start + _ROW_KEY_CHUNK].to(torch.int64) + positions = torch.arange( + start + 1, start + chunk.numel() + 1, dtype=torch.int64, device=bits.device + ) + summary += torch.stack( + ( + chunk.sum(), + (chunk * chunk).sum(), + (chunk * positions).sum(), + (chunk * chunk * positions).sum(), + ) + ) + return summary + + def _is_contiguous(shape: tuple[int, ...], stride: tuple[int, ...]) -> bool: expected = 1 for size, step in zip(reversed(shape), reversed(stride)): @@ -109,16 +130,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,9 +178,20 @@ 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. + same_layout = _same_input_layout(cache.input_key, input_key) or ( + any_contiguous_view and _same_contiguous_layout(cache.input_key, input_key) + ) + if not same_layout: + return None + same_version = cache.input_key.version == input_key.version + if not same_version: + return None + trainable_weight = _cache_base(inputs).requires_grad + if not trainable_weight or torch.is_inference_mode_enabled(): + return cache + if torch.equal(cache.input_fingerprint, _input_fingerprint(inputs)): return cache return None @@ -209,6 +269,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, diff --git a/tests/unit/torch/quantization/test_ggml_backend.py b/tests/unit/torch/quantization/test_ggml_backend.py index 136ce92f6ed..277519b48cb 100644 --- a/tests/unit/torch/quantization/test_ggml_backend.py +++ b/tests/unit/torch/quantization/test_ggml_backend.py @@ -154,6 +154,89 @@ 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) + + +@pytest.mark.parametrize("bypass_version", [False, True]) +def test_ggml_cache_detects_swapped_weight_elements(bypass_version): + block_size = GGML_FORMAT_REGISTRY["q8_0"].block_size + weight = torch.randn(1, block_size * 16, requires_grad=True) + quantizer = SimpleNamespace(num_bits="q8_0", _quantizer_cache=None) + fake_quant = GGML_FORMAT_REGISTRY["q8_0"].fake_quant + + fake_quant(weight, quantizer) + before = quantizer._quantizer_cache.packed_weights.clone() + with torch.no_grad(): + if bypass_version: + weight.data[0, 1], weight.data[0, 2] = weight.data[0, 2], weight.data[0, 1] + else: + weight[0, 1], weight[0, 2] = weight[0, 2], weight[0, 1] + + fake_quant(weight, quantizer) + assert not torch.equal(quantizer._quantizer_cache.packed_weights, before) + + +def test_ggml_cache_uses_version_fast_path_in_inference_mode(): + block_size = GGML_FORMAT_REGISTRY["q8_0"].block_size + weight = torch.randn(1, block_size, requires_grad=True) + quantizer = SimpleNamespace(num_bits="q8_0", _quantizer_cache=None) + fake_quant = GGML_FORMAT_REGISTRY["q8_0"].fake_quant + + with torch.inference_mode(): + fake_quant(weight, quantizer) + cached = quantizer._quantizer_cache.packed_weights + fake_quant(weight, quantizer) + + assert quantizer._quantizer_cache.packed_weights is cached + + 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