diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2a5a4f73dbb..0a6afcb0017 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,6 +13,7 @@ Changelog *Quantization* +- Add IQ1_S and IQ2_XS weight-only quantization with GGML-compatible 256-value block encoders, built-in ``iq1_s`` / ``iq2_xs`` PTQ recipes, and unified HF and Megatron export of the packed blocks. Quantized weights must have a final dimension divisible by 256, and Megatron export requires tensor and pipeline parallel sizes of 1. - A recipe can now **delegate its whole body to another recipe** with a top-level ``$import``; any top-level key given alongside it overrides the imported one. ``metadata.recipe_type`` became optional along with it: a recipe states its kind with a ``# modelopt-schema:`` comment, with ``metadata.recipe_type``, or by delegating to a recipe that does, and only a recipe that another file imports has to carry the schema comment. Whatever a recipe does state must be true: a schema comment and a ``recipe_type`` must agree, and so must a recipe and the recipe it delegates to. ``modelopt_recipes/models/`` uses this for checkpoint entries that a portable recipe already reproduces: the entry aliases that recipe instead of copying it. - Backfill the recipes behind NVIDIA's already-published checkpoints under ``modelopt_recipes/models/``, so a released checkpoint's quantization scheme is reachable from its own model-hub path rather than only from the general tier. For example, ``moonshotai/Kimi-K2.6`` (published as ``nvidia/Kimi-K2.6-NVFP4``) and ``Qwen/Qwen3.5-397B-A17B`` (published as ``nvidia/Qwen3.5-397B-A17B-NVFP4-V2``) each alias a portable recipe wholesale -- the general expert-only NVFP4 recipe and the ``qwen3_5_moe`` architecture recipe respectively -- rather than copying its body; other checkpoints follow in separate changes. - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index 80e8cf2e2e2..574db9bcac8 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -27,7 +27,19 @@ @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() diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py index 804d27af058..a43e1699db1 100644 --- a/modelopt/torch/quantization/ggml/iq1_s.py +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -61,8 +61,14 @@ _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, ) diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py index 2f43e71d479..e2adc2fa44f 100644 --- a/modelopt/torch/quantization/ggml/iq2_xs.py +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -60,8 +60,17 @@ _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, ) diff --git a/modelopt_recipes/configs/numerics/iq1_s.yaml b/modelopt_recipes/configs/numerics/iq1_s.yaml new file mode 100644 index 00000000000..78e309fd433 --- /dev/null +++ b/modelopt_recipes/configs/numerics/iq1_s.yaml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# IQ1_S weight quantizer using the built-in fixed-scale codebook search. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig +num_bits: iq1_s +# Cost metadata for AutoQuantize's compression estimate only; it drives no packing or +# numerics. num_bits is the string "iq1_s", so the generic estimator cannot derive the +# storage cost: 50 packed bytes * 8 / 256 weights. Keep in sync with IQ1_S_BLOCK_BYTES. +effective_bits: 1.5625 +block_sizes: + -1: 256 +backend: ggml diff --git a/modelopt_recipes/configs/numerics/iq2_xs.yaml b/modelopt_recipes/configs/numerics/iq2_xs.yaml new file mode 100644 index 00000000000..e30102ada82 --- /dev/null +++ b/modelopt_recipes/configs/numerics/iq2_xs.yaml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# IQ2_XS weight quantizer using the built-in fixed-scale codebook search. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig +num_bits: iq2_xs +# Cost metadata for AutoQuantize's compression estimate only; it drives no packing or +# numerics. num_bits is the string "iq2_xs", so the generic estimator cannot derive the +# storage cost: 74 packed bytes * 8 / 256 weights. Keep in sync with IQ2_XS_BLOCK_BYTES. +effective_bits: 2.3125 +block_sizes: + -1: 256 +backend: ggml diff --git a/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml b/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml new file mode 100644 index 00000000000..57bddbe7180 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# QuantizeConfig preset for IQ1_S weight-only quantization. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + iq1_s: configs/numerics/iq1_s + +algorithm: +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: iq1_s + - quantizer_name: '*input_quantizer' + enable: false + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml new file mode 100644 index 00000000000..4fcf1c29213 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# QuantizeConfig preset for IQ2_XS weight-only quantization. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + iq2_xs: configs/numerics/iq2_xs + +algorithm: +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: iq2_xs + - quantizer_name: '*input_quantizer' + enable: false + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/general/ptq/iq1_s.yaml b/modelopt_recipes/general/ptq/iq1_s.yaml new file mode 100644 index 00000000000..4bc31221b79 --- /dev/null +++ b/modelopt_recipes/general/ptq/iq1_s.yaml @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# IQ1_S weight-only PTQ. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + preset: configs/ptq/presets/model/iq1_s + +metadata: + description: >- + Applies uniform GGML-compatible IQ1_S weight-only quantization to eligible linear layers. + This is not a mixed per-tensor precision preset. No calibration data is required. +quantize: + $import: preset diff --git a/modelopt_recipes/general/ptq/iq2_xs.yaml b/modelopt_recipes/general/ptq/iq2_xs.yaml new file mode 100644 index 00000000000..f690c55cf09 --- /dev/null +++ b/modelopt_recipes/general/ptq/iq2_xs.yaml @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# IQ2_XS weight-only PTQ. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + preset: configs/ptq/presets/model/iq2_xs + +metadata: + description: >- + Applies uniform GGML-compatible IQ2_XS weight-only quantization to eligible linear layers. + This is not a mixed per-tensor precision preset. No calibration data is required. +quantize: + $import: preset diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 3bb131ac18d..fa493a5c945 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 26 general/ptq/ recipes (click to expand) +All 28 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -58,6 +58,8 @@ supported combinations. | `int4_blockwise_weight_only` | INT4 W4A16, block 128, weights only | none | max | | `nvfp4_mlp_weight_only` | NVFP4 W4A16 (block 32), MLP + MoE weights only | none | max | | `mxfp4_mlp_weight_only` | MXFP4 W4A16, MLP + MoE weights only | none | none (no calibration) | +| `iq1_s` | IQ1_S W1A16, eligible linears | none | none (no calibration) | +| `iq2_xs` | IQ2_XS W2A16, eligible linears | none | none (no calibration) |
@@ -113,7 +115,7 @@ activations are quantized too** (W4A4/W8A8 vs weight-only W4A16). > keeps that sensitive path at a safer precision than NVFP4 while still halving those > weights vs. BF16. -#### Weight-only schemes (W4A16 — activations stay BF16) +#### Weight-only schemes (activations stay BF16) Quantize weights only; activations run in BF16. This shrinks the model (memory-bound decode win) with much lower accuracy risk than W4A4, and **needs no @@ -136,6 +138,13 @@ activations and tensor-core math are what deliver the throughput. - **`mxfp4_mlp_weight_only`** — MXFP4 weights on MLP/MoE layers only, BF16 activations. Needs no calibration forward pass; the QAT starting point for the GPT-OSS family (see `examples/gpt-oss`). +- **`iq1_s` / `iq2_xs`** — GGML-compatible IQ1_S or IQ2_XS weights on the eligible + linear layers, with BF16 activations; `lm_head`, MoE routers, `conv1d` and the + vision branch stay in BF16 like every other preset. No calibration data is + required. Quantized weights must have a final dimension divisible by 256. + Unified HF export writes the packed GGML blocks; Megatron export additionally + requires tensor and pipeline parallel sizes of 1, and does not support + fused-MoE experts. --- diff --git a/tests/examples/hf_ptq/test_llm_ptq.py b/tests/examples/hf_ptq/test_llm_ptq.py index d307868c372..6ae697b3c09 100644 --- a/tests/examples/hf_ptq/test_llm_ptq.py +++ b/tests/examples/hf_ptq/test_llm_ptq.py @@ -76,6 +76,13 @@ def test_ptq_whisper(command): PTQCommand(quant="int8_weight_only", kv_cache_quant="none"), PTQCommand(quant="int4_awq", kv_cache_quant="none"), PTQCommand(quant="w4a8_awq_beta", kv_cache_quant="none"), + # GGML IQ weight-only, recipe-driven. These encoders require every weight's input + # dimension to be a multiple of 256; TinyLlama's 2048 and 5632 both are. Neither + # recipe calibrates -- both set algorithm: null -- so the only IQ-specific cost is + # packing each weight once and decoding it on each forward. 95s and 103s on 2xH100, + # inside the 300s tests/examples default. + PTQCommand(recipe="general/ptq/iq1_s", kv_cache_quant="none"), + PTQCommand(recipe="general/ptq/iq2_xs", kv_cache_quant="none"), PTQCommand(quant="nvfp4"), PTQCommand(quant="nvfp4_awq_lite"), # autoquant (recipe-driven) diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index df90fcfd7d3..2bf899cf650 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -33,6 +33,12 @@ from modelopt.recipe.presets import RecipeSupersededAction from modelopt.torch.opt.config_loader import BUILTIN_CONFIG_ROOT from modelopt.torch.quantization.config import LocalHessianCalibConfig, QuantizeConfig +from modelopt.torch.quantization.ggml import ( + IQ1_S_BLOCK_SIZE, + IQ1_S_EFFECTIVE_BITS, + IQ2_XS_BLOCK_SIZE, + IQ2_XS_EFFECTIVE_BITS, +) def _yaml_basenames(subdir: str) -> set[str]: @@ -125,6 +131,27 @@ def test_mlp_weight_only_recipe_matches_its_mtq_cfg(recipe_name, cfg_name): assert recipe_cfg == mtq_cfg +@pytest.mark.parametrize( + ("qformat", "block_size", "effective_bits"), + [ + ("iq1_s", IQ1_S_BLOCK_SIZE, IQ1_S_EFFECTIVE_BITS), + ("iq2_xs", IQ2_XS_BLOCK_SIZE, IQ2_XS_EFFECTIVE_BITS), + ], +) +def test_iq_recipe_matches_packing_contract(qformat, block_size, effective_bits): + recipe = load_recipe(f"general/ptq/{qformat}") + quant_cfg = recipe.quantize.model_dump(exclude_unset=True)["quant_cfg"] + weight_cfg = next( + entry["cfg"] for entry in quant_cfg if entry.get("quantizer_name") == "*weight_quantizer" + ) + + assert qformat in presets.QUANT_CFG_CHOICES + assert weight_cfg["backend"] == "ggml" + assert weight_cfg["num_bits"] == qformat + assert weight_cfg["block_sizes"][-1] == block_size + assert weight_cfg["effective_bits"] == effective_bits + + # --- RecipeSupersededAction: the flags --recipe replaces ---------------------------------------- diff --git a/tests/unit/torch/quantization/test_ggml_backend.py b/tests/unit/torch/quantization/test_ggml_backend.py index ee78bab5713..1370fdd312b 100644 --- a/tests/unit/torch/quantization/test_ggml_backend.py +++ b/tests/unit/torch/quantization/test_ggml_backend.py @@ -22,8 +22,10 @@ import modelopt.torch.quantization.ggml.backend as backend_module import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module +from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.ggml.backend import ggml_fake_quant from modelopt.torch.quantization.ggml.common import narrow_to_float32 +from modelopt.torch.quantization.nn import TensorQuantizer @pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"]) @@ -135,3 +137,79 @@ def test_narrow_to_float32_matches_the_cuda_load_float_policy(): assert torch.equal( narrowed, torch.tensor([0.0, 0.0, 0.0, largest, -largest, 1.5], dtype=torch.float32) ) + + +@pytest.mark.parametrize( + ("num_bits", "module"), [("iq1_s", iq1_s_module), ("iq2_xs", iq2_xs_module)] +) +def test_ggml_weight_is_packed_once_across_forwards(monkeypatch, num_bits, module): + """The packed weight is reused across forwards rather than re-encoded each time. + + TensorQuantizer hands the backend a fresh view of the weight on every forward, so a cache + that checked tensor identity never hit: the codebook search reran on every forward, roughly + 100x during a generate loop and over 90% of a PTQ run's wall clock. + """ + packer = f"quantize_{num_bits}" + original = getattr(module, packer) + calls = [] + + def counting(weight, **kwargs): + calls.append(tuple(weight.shape)) + return original(weight, **kwargs) + + monkeypatch.setattr(module, packer, counting) + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=num_bits, block_sizes={-1: 256}, backend="ggml") + ) + weight = torch.randn(4, 256) + + with torch.inference_mode(): # what generate() runs under + for _ in range(5): + quantizer(weight) + + assert calls == [(4, 256)], f"expected one pack, got {len(calls)}" + + +@pytest.mark.parametrize( + ("num_bits", "module"), [("iq1_s", iq1_s_module), ("iq2_xs", iq2_xs_module)] +) +def test_ggml_decode_chunk_is_sized_independently_of_the_encode_chunk( + monkeypatch, num_bits, module +): + """The decode runs every forward; the encode runs once and holds the big temporaries. + + Sharing one constant between them is what made IQ2_XS four times slower end to end than + IQ1_S, so pin that the decode gets its own, larger chunk. + """ + seen = {} + original = getattr(module, f"dequantize_{num_bits}") + + def recording(packed_weights, weight_shape, **kwargs): + seen["block_chunk_size"] = kwargs["block_chunk_size"] + return original(packed_weights, weight_shape, **kwargs) + + monkeypatch.setattr(module, f"dequantize_{num_bits}", recording) + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=num_bits, block_sizes={-1: 256}, backend="ggml") + ) + quantizer(torch.randn(4, 256)) + + assert seen["block_chunk_size"] == module._DEFAULT_DECODE_CHUNK_SIZE + assert module._DEFAULT_DECODE_CHUNK_SIZE > module._DEFAULT_BLOCK_CHUNK_SIZE + + +@pytest.mark.parametrize( + ("num_bits", "module"), [("iq1_s", iq1_s_module), ("iq2_xs", iq2_xs_module)] +) +def test_ggml_decode_is_invariant_to_chunk_size(num_bits, module): + """Chunking the decode is a memory bound, not a numerical choice.""" + torch.manual_seed(0) + weight = torch.randn(3, 1024, dtype=torch.bfloat16) + packed, shape = getattr(module, f"quantize_{num_bits}")(weight) + dequantize = getattr(module, f"dequantize_{num_bits}") + + reference = dequantize(packed, shape, dtype=weight.dtype, block_chunk_size=1) + for chunk in (2, 7, 4096): + assert torch.equal( + dequantize(packed, shape, dtype=weight.dtype, block_chunk_size=chunk), reference + ) diff --git a/tests/unit/torch/quantization/test_iq2_xs.py b/tests/unit/torch/quantization/test_iq2_xs.py index 73a1b352c8a..3dde5a8a902 100644 --- a/tests/unit/torch/quantization/test_iq2_xs.py +++ b/tests/unit/torch/quantization/test_iq2_xs.py @@ -102,6 +102,32 @@ def test_iq2_xs_dequantizes_pinned_scale_factor(): assert torch.equal(decoded, torch.full((1, 256), 43 * 31 / 8, dtype=torch.float32)) +@pytest.mark.parametrize("sign_index", [0b0000001, 0b0000011, 0b0000111, 0b1010101, 0b1111111]) +def test_iq2_xs_dequantizes_the_implied_eighth_sign_bit(sign_index): + """The eighth sign is not stored; it is the parity of the seven that are. + + The pinned-scale test above only covers sign_index 0, where every value is positive + whether or not the implied bit is derived correctly. These indices vary which payload + bits are set so the parity actually has to be computed. + """ + codes = 511 | (sign_index << 9) # entry 511 holds eight 43s + packed = torch.zeros((1, 1, 74), dtype=torch.uint8) + packed[0, 0, :2] = torch.tensor([1.0], dtype=torch.float16).view(torch.uint8) + packed[0, 0, 2:66:2] = codes & 0xFF + packed[0, 0, 3:66:2] = codes >> 8 + packed[0, 0, 66:] = 0xFF # local code 15 in both nibbles + + decoded = dequantize_iq2_xs(packed, torch.tensor([1, 256]), dtype=torch.float32) + + payload_bits = [(sign_index >> bit) & 1 for bit in range(7)] + magnitude = 43 * 31 / 8 # entry value 43, local code 15 -> (2 * 15 + 1) / 8 + expected_group = torch.tensor( + [-magnitude if bit else magnitude for bit in (*payload_bits, sum(payload_bits) % 2)], + dtype=torch.float32, + ) + assert torch.equal(decoded.reshape(32, 8), expected_group.expand(32, 8)) + + def test_iq2_xs_requires_complete_last_dimension_blocks(): with pytest.raises(ValueError, match="last weight dimension"): quantize_iq2_xs(torch.ones(2, 257))