Make folding exact for GPTQ'd GGML weights, so the vLLM fake-quant export evaluates them

docs/source/deployment/3_unified_hf.rst

@@ -89,6 +89,13 @@ The generated configuration records ``quant_method: modelopt``, ``packing: ggml`
 block size, and the payload byte count. GGML payloads are not represented as compressed-tensors
 integer ``weights`` groups because all scales and indices are embedded in each packed block.
 
+To evaluate a GGML-quantized model in a runtime without GGML kernels, such as vLLM, export it with
+``--vllm_fakequant_export`` in ``examples/hf_ptq/hf_ptq.py`` instead. Each GGML weight is folded into
+the values its blocks decode to, including the codes GPTQ chose, and the checkpoint carries no
+quantization configuration, so a stock vLLM image serves it as an ordinary model. The extra
+``vllm_fq_modelopt_state.pth`` is read only by the fake-quant launcher in ``examples/vllm_serve``, so
+activation or KV-cache quantization elsewhere in the model applies only when served through it.
+
 Each 74-byte IQ2_XS block represents 256 logical weights:
 
 * bytes 0--1 are the little-endian FP16 super-block scale ``d``;

examples/hf_ptq/README.md

@@ -671,7 +671,7 @@ with torch.inference_mode():
 python hf_ptq.py --pyt_ckpt_path <huggingface_model_card> --qformat fp8 --export_path <quantized_ckpt_path> --trust_remote_code
 ```
 
-> *For exporting fake-quantized models for vLLM serving (e.g., for research or kernels not yet supported in real-quant), use the `--vllm_fakequant_export` flag. See [vllm_serve/README.md](../vllm_serve/README.md) for details.*
+> *For exporting fake-quantized models for vLLM serving (e.g., for research or kernels not yet supported in real-quant), use the `--vllm_fakequant_export` flag. See [vllm_serve/README.md](../vllm_serve/README.md) for details. When every quantizer is a weight quantizer, as in the GGML recipes (IQ, Q8_0), the export is a plain BF16 checkpoint that a stock vLLM image serves directly.*
 
 ### Hugging Face framework [Script](./scripts/huggingface_example.sh)
 

modelopt/torch/export/quant_utils.py

@@ -1003,6 +1003,8 @@ def from_quantized_weight(
     "_bias_value",
     "input_quantizer._pre_quant_scale",
     "weight_shape",
+    # The GGML payload GPTQ pins to a weight quantizer; the exported weight already holds it.
+    "_ggml_pinned_",
 )
 
 

modelopt/torch/quantization/ggml/common.py

@@ -26,6 +26,8 @@
 
 GGML_BLOCK_SIZE = 256
 _SIGNED_INT_OF_SIZE = {2: torch.int16, 4: torch.int32, 8: torch.int64}
+# Columns hashed at a time, bounding _row_keys' int64 temporaries to this slice of the weight.
+_ROW_KEY_CHUNK = 256
 
 
 class _CacheKey(NamedTuple):
@@ -168,13 +170,18 @@ def _row_keys(x: torch.Tensor) -> torch.Tensor:
     """A 64-bit hash of each row's bits; equal rows always get equal keys.
 
     Integer sums wrap and are associative, so unlike a floating-point projection the key does not
-    depend on how the reduction is ordered, which differs between differently shaped products.
+    depend on how the reduction is ordered, which differs between differently shaped products. That
+    is also what lets the columns be summed a chunk at a time.
     """
     multipliers = torch.randint(
         -(2**62), 2**62, (x.shape[-1],), generator=torch.Generator().manual_seed(0)
-    )
-    bits = x.contiguous().view(_SIGNED_INT_OF_SIZE[x.element_size()]).to(torch.int64)
-    return (bits * multipliers.to(x.device)).sum(-1)
+    ).to(x.device)
+    bits = x.contiguous().view(_SIGNED_INT_OF_SIZE[x.element_size()])
+    keys = torch.zeros(x.shape[:-1], dtype=torch.int64, device=x.device)
+    for start in range(0, x.shape[-1], _ROW_KEY_CHUNK):
+        columns = slice(start, start + _ROW_KEY_CHUNK)
+        keys += (bits[..., columns].to(torch.int64) * multipliers[columns]).sum(-1)
+    return keys
 
 
 def fake_quantize_with_cache(
@@ -305,8 +312,10 @@ def _pinned_rows(self, quantizer, weight: torch.Tensor) -> torch.Tensor | None:
         """The pinned payload of ``weight``, a view of the pinned weight or rows taken from it.
 
         Fake quant hands over the pinned weight itself in 256-value rows; export may hand over row
-        slices of it, such as the gate and up halves of a fused projection. The comparison syncs,
-        but it only runs for weights the identity cache has not seen.
+        slices of it, such as the gate and up halves of a fused projection. Folding hands over a
+        float32 copy of a narrower weight, which holds the decoding rounded to that narrower dtype,
+        so a float32 weight also matches the decoding rounded through BF16 or FP16. The comparison
+        syncs, but it only runs for weights the identity cache has not seen.
         """
         pinned = pinned_packed_weight(quantizer, self.name)
         width = weight.shape[-1] // self.block_size
@@ -317,17 +326,28 @@ def _pinned_rows(self, quantizer, weight: torch.Tensor) -> torch.Tensor | None:
         ):
             return None
         pinned = pinned.to(weight.device).reshape(-1, width, self.block_bytes)
-        table = self.dequantize(
-            pinned, torch.tensor((pinned.shape[0], weight.shape[-1])), dtype=torch.float32
-        ).to(weight.dtype)
+        shape = torch.tensor((pinned.shape[0], weight.shape[-1]))
         rows = weight.reshape(-1, weight.shape[-1])
-        if rows.shape == table.shape and torch.equal(rows, table):
-            index = slice(None)
-        else:
-            index = _matching_rows(table, rows)
-            if index is None:
-                return None
-        return pinned[index].reshape(*weight.shape[:-1], width, self.block_bytes)
+        storage_dtypes = [weight.dtype]
+        if weight.dtype == torch.float32:
+            storage_dtypes += [torch.bfloat16, torch.float16]
+
+        def decoding(storage_dtype):
+            # Decoding straight into the storage dtype rounds once, as GPTQ did writing the weight,
+            # and holds one table at a time rather than a float32 copy beside each rounding.
+            return self.dequantize(pinned, shape, dtype=storage_dtype).to(weight.dtype)
+
+        # Every whole-weight comparison first: matching rows builds int64 keys several times the
+        # weight's size, and only row slices need it.
+        if rows.shape[0] == pinned.shape[0]:
+            for storage_dtype in storage_dtypes:
+                if torch.equal(rows, decoding(storage_dtype)):
+                    return pinned.reshape(*weight.shape[:-1], width, self.block_bytes)
+        for storage_dtype in storage_dtypes:
+            index = _matching_rows(decoding(storage_dtype), rows)
+            if index is not None:
+                return pinned[index].reshape(*weight.shape[:-1], width, self.block_bytes)
+        return None
 
 
 def narrow_to_float32(blocks: torch.Tensor) -> torch.Tensor:

tests/unit/torch/export/test_export_weight.py

@@ -29,6 +29,7 @@
 )
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY
+from modelopt.torch.quantization.ggml.common import pin_packed_weight
 from modelopt.torch.quantization.nn import TensorQuantizer
 from modelopt.torch.quantization.utils import quantizer_attr_names
 
@@ -187,6 +188,19 @@ def test_export_repacks_a_weight_changed_since_fake_quant(num_bits):
     assert torch.equal(linear.weight, expected)
 
 
+@pytest.mark.parametrize("num_bits", sorted(GGML_FORMAT_REGISTRY))
+def test_export_drops_the_payload_gptq_pinned(num_bits):
+    """GPTQ pins its payload to the quantizer as a buffer; only the packed weight is exported."""
+    linear = _iq_linear(num_bits)
+    packed, _ = GGML_FORMAT_REGISTRY[num_bits].quantize(linear.weight.detach())
+    pin_packed_weight(linear.weight_quantizer, num_bits, packed)
+
+    _export_quantized_weight(linear, torch.bfloat16)
+    state_dict = postprocess_state_dict(linear.state_dict(), maxbound=448, quantization=None)
+
+    assert list(state_dict) == ["weight"]
+
+
 class QuantMoELinear(nn.Module):
     def __init__(self):
         super().__init__()

tests/unit/torch/quantization/test_ggml_gptq.py

@@ -215,6 +215,26 @@ def test_pin_outlives_the_weight_tensor():
     assert not torch.equal(IQ1_S.quantize(model.weight.detach())[0], _pin(model))
 
 
+def test_folding_a_bf16_weight_decodes_the_pinned_payload():
+    """Folding fake-quantizes a float32 copy of the weight; it must still decode GPTQ's codes.
+
+    A BF16 weight holds the pinned payload's decoding rounded to BF16, so the float32 copy that
+    ``mtq.fold_weight`` and the vLLM fake-quant export hand to the quantizer is not that decoding
+    in float32. Re-encoding it instead would fold in different codes than GPTQ chose.
+    """
+    torch.manual_seed(0)
+    model = torch.nn.Linear(512, 8, bias=False, dtype=torch.bfloat16)
+    inputs = _inputs().to(torch.bfloat16)
+    mtq.quantize(model, _iq_config(GPTQ), forward_loop=lambda m: m(inputs))
+    decoded = _decoded(_pin(model), model.weight).to(torch.bfloat16)
+    assert torch.equal(model.weight, decoded)
+
+    # The vLLM fake-quant export folds this way.
+    assert torch.equal(model.weight_quantizer(model.weight.float()).to(torch.bfloat16), decoded)
+    mtq.fold_weight(model)
+    assert torch.equal(model.weight, decoded)
+
+
 def test_pin_survives_save_and_restore():
     model, inputs = _gptq_model()
     buffer = io.BytesIO()