mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
[OMNIML-5899] Add IQ post-training quantization recipes (#2449)
## Summary - add numerics configs for IQ1_S and IQ2_XS - add model presets and general PTQ recipes - document the supported weight-shape requirement - reuse the packed IQ weight across forwards instead of re-encoding it every time - add end-to-end `hf_ptq` coverage for both formats - add the release-note entry ## PR split This work is split into four focused PRs. Each PR targets `main` and owns a disjoint file set: 1. **Kernel** — [#2448: Add CUDA kernels for IQ packing](https://github.com/NVIDIA/Model-Optimizer/pull/2448) 2. **Quantization** — [#2446: Add IQ quantization codecs and backend](https://github.com/NVIDIA/Model-Optimizer/pull/2446) 3. **Export** — [#2447: Export IQ checkpoints from HF and Megatron](https://github.com/NVIDIA/Model-Optimizer/pull/2447) 4. **Recipes** — [#2449: Add IQ post-training quantization recipes](https://github.com/NVIDIA/Model-Optimizer/pull/2449) The required merge order is #2448, #2446, #2447, then #2449. ## Scope Sixteen files. The PR started as eight recipe config, documentation and changelog files; the end-to-end test added for them surfaced two performance bugs in the already-merged codec, and fixing those pulled in the codec files and their tests. Beyond the original recipe set it now touches four files owned by #2446 — `ggml/common.py`, `ggml/iq1_s.py`, `ggml/iq2_xs.py` and `ggml/backend.py` — plus three test files. It still contains no kernel or export files. Keeping those fixes here rather than moving them to #2446 is deliberate and confirmed with the stack owner: #2446 is already merged, and both bugs are only observable through the end-to-end test this PR adds, so splitting them would separate each fix from the test that demonstrates it. ## Packed-weight cache fix Adding the end-to-end test made the cost visible: a TinyLlama IQ1_S `hf_ptq` run spent **499 of its 537 seconds inside the IQ1_S encoder**, packing the same 154 weights 15400 times — 100 times each. The repacking is not calibration. These recipes set `algorithm: null` and `hf_ptq` logs `Dynamic quantization. Calibration skipped.`. The 100 passes are the sample `generate()` calls `hf_ptq` makes before and after quantization: one decode step re-runs weight fake-quant on every linear, and the packed payload was thrown away each time. `_PackedWeightCache` was already there to prevent exactly this, and it never hit. It keyed on the identity of the tensor the backend was handed, but `TensorQuantizer` passes a fresh *view* of the weight on every forward, so the identity check never matched twice. The fix anchors the entry to `inputs._base` — the parameter the view is taken from — held as a weakref. The parameter is stable across forwards, so the cache hits; the reference stays weak, so the payload is released with the weight and offloaded/meta-device flows are unaffected. (A strong reference does make the cache hit, but pins full-precision storage for the life of the quantizer, which is the opposite of what those flows need.) Measured on TinyLlama IQ1_S `hf_ptq`, 2×H100, same command before and after: | | packer calls | packing time | wall clock | |---|---|---|---| | before | 15400 | 499.0 s | 8m57s | | after | 154 (one per weight) | 6.2 s | 3m21s | `test_ggml_weight_is_packed_once_across_forwards` pins this: it counts encoder calls across five forwards under `torch.inference_mode()` (what `generate()` runs under) and asserts exactly one. ## Decode chunk fix With packing cached, the end-to-end cost moved entirely into the decode, and IQ2_XS was still 4x slower than IQ1_S (1081s vs 260s per case). Instrumenting both showed packing was no longer the cost at all — IQ2_XS packs *faster*: | | pack calls | packing time | wall clock | |---|---|---|---| | IQ1_S | 154 | 6.2 s | 3m21s | | IQ2_XS | 154 | 1.6 s | 17m59s | The cause was one constant serving two loops with opposite characteristics. `_DEFAULT_BLOCK_CHUNK_SIZE` bounds the torch encode fallback, which holds the large codebook-search temporaries and runs once per weight; IQ2_XS sets it to 256 rather than IQ1_S's 1024 because its search sweeps sixteen local scales per grid tile. But the *decode* shared it — and the decode has tiny temporaries, runs on every forward, and is never cached, so a small chunk only multiplies kernel launches. Decoding a 2048x5632 weight: | chunk | IQ2_XS decode | transient peak | |---|---|---| | 256 (was) | 91.4 ms | +24 MiB | | 1024 | 22.9 ms | +31 MiB | | 4096 (now) | 5.8 ms | +56 MiB | The decode now takes its own `_DEFAULT_DECODE_CHUNK_SIZE`, threaded through `fake_quantize_with_cache`. The encode bounds are untouched, so the memory ceiling stays where it was aimed. The seven-iteration sign-parity loop in `dequantize_iq2_xs` is also folded into three XOR steps, off the same per-forward path. Per case in `tests/examples/hf_ptq` on 2xH100: | | before | after | |---|---|---| | IQ1_S | 259.65 s | 94.64 s | | IQ2_XS | 1081.05 s | 102.77 s | Both now fit the 300s `tests/examples` default, so the cases carry no explicit timeout. ## Integration contracts - `block_sizes: {-1: 256}` records the native packed-block contract; it does not drive the GGML fake-quant scale search. The export path reads `TensorQuantizer.block_sizes[-1]` through `get_weight_block_size`, validates it against the format block size, and records `group_size: 256` in checkpoint metadata. - The model presets intentionally expose `--qformat iq1_s` and `--qformat iq2_xs` in `hf_ptq`. ## Testing - `tests/unit/torch/quantization/ -k 'ggml or iq1 or iq2'` — 52 passed, 1 skipped - `tests/unit/recipe/test_presets.py` — both shipped IQ recipes load with the expected backend and no unused search option - `tests/examples/hf_ptq/test_llm_ptq.py -k 'iq1_s or iq2_xs'` — 2 passed on 2×H100 (TinyLlama, both formats end to end through export), 3m17s for the pair - pre-commit hooks pass on all changed files - larger-model sanity check outside CI: Qwen3.8-27B (2256 quantizers) quantizes and exports with `general/ptq/iq1_s` on 2xH100 in 66m46s, export itself 192s --------- Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
hychiang
co-authored by
Chenjie Luo
Claude Opus 5
Chenjie Luo
Claude Opus 5
parent
5bb7343592
commit
1b4e7dfb14
@@ -26,7 +26,7 @@ def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor:
|
||||
"""Dispatch an IQ quantizer to its format-specific implementation."""
|
||||
num_bits = getattr(quantizer, "num_bits", None)
|
||||
extra_args = getattr(quantizer, "backend_extra_args", None) or {}
|
||||
unknown_args = set(extra_args) - {"block_chunk_size"}
|
||||
unknown_args = set(extra_args) - {"block_chunk_size", "decode_chunk_size"}
|
||||
if unknown_args:
|
||||
raise ValueError(f"Unsupported ggml backend_extra_args: {sorted(unknown_args)}")
|
||||
if num_bits == "iq1_s":
|
||||
|
||||
@@ -27,7 +27,19 @@ GGML_BLOCK_SIZE = 256
|
||||
|
||||
@dataclass
|
||||
class _PackedWeightCache:
|
||||
input_ref: weakref.ReferenceType
|
||||
"""One weight's packed payload, reused across forwards.
|
||||
|
||||
``base_ref`` points at the parameter, not at the tensor the backend was handed.
|
||||
TensorQuantizer passes a fresh view of the weight on every forward, so a weakref to that
|
||||
view dies as soon as the forward returns and an identity check against it never matches
|
||||
again -- which is what kept this cache from ever hitting.
|
||||
|
||||
Keeping it a weakref matters: a strong reference would pin full-precision storage alive and
|
||||
defeat offloaded or meta-device flows. Tying the entry to the parameter's lifetime means the
|
||||
payload stops being reused exactly when the weight it came from is released.
|
||||
"""
|
||||
|
||||
base_ref: weakref.ReferenceType
|
||||
input_key: tuple[object, ...]
|
||||
format_name: str
|
||||
block_chunk_size: int
|
||||
@@ -35,6 +47,16 @@ class _PackedWeightCache:
|
||||
weight_shape: torch.Tensor
|
||||
|
||||
|
||||
def _cache_base(inputs: torch.Tensor) -> torch.Tensor:
|
||||
"""The tensor whose lifetime the cached payload should follow.
|
||||
|
||||
``inputs`` is a per-forward view; ``inputs._base`` is the parameter behind it, which lives
|
||||
as long as the module does.
|
||||
"""
|
||||
base = inputs._base
|
||||
return inputs if base is None else base
|
||||
|
||||
|
||||
def _input_cache_key(inputs: torch.Tensor) -> tuple[object, ...] | None:
|
||||
try:
|
||||
version = inputs._version
|
||||
@@ -57,16 +79,18 @@ def fake_quantize_with_cache(
|
||||
*,
|
||||
format_name: str,
|
||||
block_chunk_size: int,
|
||||
decode_chunk_size: int,
|
||||
quantize: Callable[..., tuple[torch.Tensor, torch.Tensor]],
|
||||
dequantize: Callable[..., torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
"""Fake-quantize a weight while caching its compact packed representation."""
|
||||
input_key = _input_cache_key(inputs)
|
||||
cache_base = _cache_base(inputs)
|
||||
cache = getattr(quantizer, "_quantizer_cache", None)
|
||||
if (
|
||||
isinstance(cache, _PackedWeightCache)
|
||||
and input_key is not None
|
||||
and cache.input_ref() is inputs
|
||||
and cache.base_ref() is cache_base
|
||||
and cache.input_key == input_key
|
||||
and cache.format_name == format_name
|
||||
and cache.block_chunk_size == block_chunk_size
|
||||
@@ -76,7 +100,7 @@ def fake_quantize_with_cache(
|
||||
packed_weights, weight_shape = quantize(inputs, block_chunk_size=block_chunk_size)
|
||||
if input_key is not None:
|
||||
quantizer._quantizer_cache = _PackedWeightCache(
|
||||
input_ref=weakref.ref(inputs),
|
||||
base_ref=weakref.ref(cache_base),
|
||||
input_key=input_key,
|
||||
format_name=format_name,
|
||||
block_chunk_size=block_chunk_size,
|
||||
@@ -86,11 +110,13 @@ def fake_quantize_with_cache(
|
||||
else:
|
||||
quantizer._quantizer_cache = None
|
||||
|
||||
# Sized separately from the encode chunk: packing happens once per weight and is bounded
|
||||
# by its search temporaries, while this runs on every forward and is bounded by launches.
|
||||
reconstructed = dequantize(
|
||||
packed_weights,
|
||||
weight_shape,
|
||||
dtype=inputs.dtype,
|
||||
block_chunk_size=block_chunk_size,
|
||||
block_chunk_size=decode_chunk_size,
|
||||
)
|
||||
return inputs + (reconstructed - inputs).detach()
|
||||
|
||||
|
||||
@@ -61,8 +61,14 @@ IQ1_S_EFFECTIVE_BITS = IQ1_S_BLOCK_BYTES * 8 / IQ1_S_BLOCK_SIZE
|
||||
_IQ1_S_DELTA = 0.125
|
||||
_IQ1_S_NATIVE_MAX = 16.875
|
||||
_IQ1_S_SCALE_ANCHOR = 0.61
|
||||
# At 1024 blocks, each largest IQ1_S search temporary is about 16 MiB in FP32.
|
||||
# Bounds the torch encode fallback, whose codebook search holds the large temporaries: at
|
||||
# 1024 blocks each is about 16 MiB in FP32. The CUDA encoder ignores this entirely.
|
||||
_DEFAULT_BLOCK_CHUNK_SIZE = 1024
|
||||
# The decode's temporaries are far smaller, so it is launch-bound rather than memory-bound
|
||||
# and wants a bigger chunk -- and unlike packing it is not cached, so it runs on every
|
||||
# forward. Measured decoding a 2048x5632 weight: 66.5 ms at 256 blocks, 4.2 ms at 4096,
|
||||
# where the transient peak is +42 MiB.
|
||||
_DEFAULT_DECODE_CHUNK_SIZE = 4096
|
||||
|
||||
|
||||
_GRID_CACHE: dict[torch.device, torch.Tensor] = {}
|
||||
@@ -198,7 +204,7 @@ def dequantize_iq1_s(
|
||||
weight_shape: torch.Tensor,
|
||||
*,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE,
|
||||
block_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE,
|
||||
) -> torch.Tensor:
|
||||
"""Decode GGML-compatible IQ1_S payload bytes."""
|
||||
shape = validate_packed_weights(
|
||||
@@ -234,6 +240,7 @@ def iq1_s_fake_quant(
|
||||
quantizer,
|
||||
*,
|
||||
block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE,
|
||||
decode_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE,
|
||||
) -> torch.Tensor:
|
||||
"""IQ1_S weight backend for TensorQuantizer, with pass-through backward."""
|
||||
if getattr(quantizer, "num_bits", None) != "iq1_s":
|
||||
@@ -243,6 +250,7 @@ def iq1_s_fake_quant(
|
||||
quantizer,
|
||||
format_name="iq1_s",
|
||||
block_chunk_size=block_chunk_size,
|
||||
decode_chunk_size=decode_chunk_size,
|
||||
quantize=quantize_iq1_s,
|
||||
dequantize=dequantize_iq1_s,
|
||||
)
|
||||
|
||||
@@ -60,8 +60,17 @@ _IQ2_XS_NATIVE_MAX = 43 * 31 / 8
|
||||
_IQ2_XS_SCALE_ANCHOR_MIN = 0.65
|
||||
_IQ2_XS_SCALE_ANCHOR_MAX = 0.92
|
||||
_IQ2_XS_PEAK_TO_RMS_TAPER = 0.035
|
||||
# At 256 blocks, the largest IQ2_XS search temporary is about 16 MiB in FP32.
|
||||
# Bounds the torch encode fallback, whose codebook search holds the large temporaries: at
|
||||
# 256 blocks the largest is about 16 MiB in FP32. The CUDA encoder ignores this entirely.
|
||||
# This is four times smaller than the IQ1_S bound because the IQ2_XS search sweeps sixteen
|
||||
# local scales per grid tile.
|
||||
_DEFAULT_BLOCK_CHUNK_SIZE = 256
|
||||
# The decode's temporaries are far smaller, so it is launch-bound rather than memory-bound
|
||||
# and wants a bigger chunk -- and unlike packing it is not cached, so it runs on every
|
||||
# forward. Sharing the encode bound above is what made IQ2_XS four times slower end to end
|
||||
# than IQ1_S. Measured decoding a 2048x5632 weight: 91.4 ms at 256 blocks, 22.9 ms at 1024,
|
||||
# 5.8 ms at 4096, where the transient peak is +56 MiB.
|
||||
_DEFAULT_DECODE_CHUNK_SIZE = 4096
|
||||
_SCALE_BLOCK_CHUNK_SIZE = 4096
|
||||
|
||||
|
||||
@@ -204,7 +213,7 @@ def dequantize_iq2_xs(
|
||||
weight_shape: torch.Tensor,
|
||||
*,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE,
|
||||
block_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE,
|
||||
) -> torch.Tensor:
|
||||
"""Decode GGML-compatible IQ2_XS payload bytes."""
|
||||
shape = validate_packed_weights(
|
||||
@@ -225,10 +234,13 @@ def dequantize_iq2_xs(
|
||||
)
|
||||
entries = codes & 0x1FF
|
||||
sign_index = codes >> 9
|
||||
parity = torch.zeros_like(sign_index)
|
||||
for bit in range(7):
|
||||
parity ^= (sign_index >> bit) & 1
|
||||
sign_mask = sign_index | (parity << 7)
|
||||
# XOR-fold the seven payload bits down to bit 0 to recover the eighth sign bit.
|
||||
# The loop this replaces cost seven elementwise passes per chunk, and unlike packing
|
||||
# the decode is not cached -- it runs again on every forward.
|
||||
folded = sign_index ^ (sign_index >> 4)
|
||||
folded ^= folded >> 2
|
||||
folded ^= folded >> 1
|
||||
sign_mask = sign_index | ((folded & 1) << 7)
|
||||
signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float()
|
||||
|
||||
scale_bytes = block_chunk[:, 66:].to(torch.int64)
|
||||
@@ -248,6 +260,7 @@ def iq2_xs_fake_quant(
|
||||
quantizer,
|
||||
*,
|
||||
block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE,
|
||||
decode_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE,
|
||||
) -> torch.Tensor:
|
||||
"""IQ2_XS weight backend for TensorQuantizer, with pass-through backward."""
|
||||
if getattr(quantizer, "num_bits", None) != "iq2_xs":
|
||||
@@ -257,6 +270,7 @@ def iq2_xs_fake_quant(
|
||||
quantizer,
|
||||
format_name="iq2_xs",
|
||||
block_chunk_size=block_chunk_size,
|
||||
decode_chunk_size=decode_chunk_size,
|
||||
quantize=quantize_iq2_xs,
|
||||
dequantize=dequantize_iq2_xs,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user