[OMNIML-5899] Add IQ post-training quantization recipes

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.

modelopt/torch/quantization/ggml/backend.py

@@ -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":

modelopt/torch/quantization/ggml/common.py

@@ -27,14 +27,36 @@
 
 @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
     packed_weights: torch.Tensor
     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()
 

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,
     )

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,
     )

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

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

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

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

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

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