From a3052f93c211ea7976e3faa883711f864807b515 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 22 Sep 2026 22:04:44 +0000 Subject: [PATCH 1/4] Add the IQ2_XXS weight-only quantization format llama.cpp defines five GGML IQ formats at one and two bits; we ship two. This adds IQ2_XXS, at 2.0625 bits per weight between IQ1_S and IQ2_XS, and is the first of three changes that close the gap. On a real mixed-precision checkpoint (unsloth/Qwen3.8-27B-GGUF) IQ2_XXS alone covers 59 tensors and 2.84B parameters, 10.6% of the file, which a reader limited to IQ1_S and IQ2_XS cannot consume. The encoder follows the existing single-pass grid search at a fixed anchored super-block scale, and the CUDA kernel follows the existing per-block structure. IQ2_XXS reuses IQ2_XS's even-parity sign rule but packs a 4-bit sub-block scale into the same 32-bit word as four 7-bit sign indices, and its 256-entry grid needs no high index bits. Two pieces of groundwork come with it, both of which the next two formats reuse. Export spelled the IQ family as a two-element tuple at nine sites; those become an IQ_FORMATS frozenset plus per-format packer and block-geometry tables, so a format is a row rather than a sweep through the exporters. And the per-format test files, which had drifted apart, become one parametrized module per layer, so every format is held to the same contract. The decoder is validated against llama.cpp's own output over 11,100,160 blocks from that checkpoint, all bit-identical to dequantize_row_iq2_xxs, and the new codebook matches the ggml-common.h table entry for entry. Blocks lifted from the checkpoint ship as conformance vectors so CI keeps checking bytes we did not produce. Measured on a 5632x2048 weight, the CUDA encoder runs at 1047.9 M elem/s against the torch search's 10.7. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 1 + modelopt/torch/export/quant_format.py | 16 +- modelopt/torch/export/quant_utils.py | 29 +- modelopt/torch/export/unified_export_hf.py | 18 +- .../torch/export/unified_export_megatron.py | 32 +- .../kernels/quantization/ggml/common.cuh | 1 + .../torch/kernels/quantization/ggml/ggml.cpp | 26 ++ .../kernels/quantization/ggml/iq2_xxs.cu | 240 +++++++++++++++ modelopt/torch/quantization/extensions.py | 1 + modelopt/torch/quantization/ggml/__init__.py | 8 +- modelopt/torch/quantization/ggml/backend.py | 20 +- modelopt/torch/quantization/ggml/codebooks.py | 19 ++ modelopt/torch/quantization/ggml/iq2_xxs.py | 286 ++++++++++++++++++ .../configs/numerics/iq2_xxs.yaml | 26 ++ .../configs/ptq/presets/model/iq2_xxs.yaml | 32 ++ modelopt_recipes/general/ptq/iq2_xxs.yaml | 27 ++ modelopt_recipes/ptq.md | 12 +- .../quantization/iq_llama_cpp_vectors.py | 162 ++++++++++ tests/examples/hf_ptq/test_llm_ptq.py | 12 +- .../quantization/test_iq_formats_cuda.py | 179 +++++++++++ tests/unit/recipe/test_presets.py | 3 + .../torch/quantization/test_ggml_backend.py | 3 +- .../torch/quantization/test_iq_formats.py | 274 +++++++++++++++++ 23 files changed, 1379 insertions(+), 48 deletions(-) create mode 100644 modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu create mode 100644 modelopt/torch/quantization/ggml/iq2_xxs.py create mode 100644 modelopt_recipes/configs/numerics/iq2_xxs.yaml create mode 100644 modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml create mode 100644 modelopt_recipes/general/ptq/iq2_xxs.yaml create mode 100644 tests/_test_utils/torch/quantization/iq_llama_cpp_vectors.py create mode 100644 tests/gpu/torch/quantization/test_iq_formats_cuda.py create mode 100644 tests/unit/torch/quantization/test_iq_formats.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0a6afcb0017..1f2ad30b05e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -14,6 +14,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. +- Add ``iq1_m``, ``iq2_xxs`` and ``iq2_s`` weight-only quantization with CUDA encoders and ``general/ptq`` recipes, completing the GGML IQ formats at one and two bits. The same 256-value block constraint applies as for ``iq1_s`` and ``iq2_xs``. - 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/export/quant_format.py b/modelopt/torch/export/quant_format.py index d1bb213b13d..dea3148404d 100644 --- a/modelopt/torch/export/quant_format.py +++ b/modelopt/torch/export/quant_format.py @@ -37,16 +37,26 @@ QUANTIZATION_FP8_PB_WO = "fp8_pb_wo" QUANTIZATION_FP8_PC_PT = "fp8_pc_pt" QUANTIZATION_IQ1_S = "iq1_s" +QUANTIZATION_IQ2_XXS = "iq2_xxs" QUANTIZATION_IQ2_XS = "iq2_xs" +# Every GGML IQ format. They share the weight-only, 256-value-block, per-module-scale +# shape, so export treats them as one family; adding a format means adding it here +# rather than extending a tuple at each use site. +IQ_FORMATS = frozenset( + { + QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_XXS, + QUANTIZATION_IQ2_XS, + } +) + # Formats whose scales are purely per-module, so export never merges them across the q/k/v # and gate/up groups that share an input. Every other format unifies input_amax (and, for # NVFP4, weight_scale_2) across such a group, which only a whole-model forward can discover. -FUSION_FREE_FORMATS = frozenset( +FUSION_FREE_FORMATS = IQ_FORMATS | frozenset( { QUANTIZATION_FP8, - QUANTIZATION_IQ1_S, - QUANTIZATION_IQ2_XS, QUANTIZATION_NONE, QUANTIZATION_FP8_PB_REAL, } diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index c36369bab19..8caa4ce2c19 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -34,6 +34,9 @@ IQ2_XS_BLOCK_BYTES, IQ2_XS_BLOCK_SIZE, IQ2_XS_EFFECTIVE_BITS, + IQ2_XXS_BLOCK_BYTES, + IQ2_XXS_BLOCK_SIZE, + IQ2_XXS_EFFECTIVE_BITS, ) from modelopt.torch.quantization.model_calib import ( enable_stats_collection, @@ -59,6 +62,7 @@ from ..quantization.nn import NVFP4StaticQuantizer, SequentialQuantizer, TensorQuantizer from .model_utils import TiedWeightMap, get_language_model_from_vl from .quant_format import ( + IQ_FORMATS, KV_CACHE_FP8, KV_CACHE_FP8_K_NVFP4_V, KV_CACHE_INT8, @@ -73,6 +77,7 @@ QUANTIZATION_INT8_WO, QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS, + QUANTIZATION_IQ2_XXS, QUANTIZATION_MXFP4, QUANTIZATION_MXFP8, QUANTIZATION_NONE, @@ -88,6 +93,14 @@ logger = logging.getLogger(__name__) +# Block geometry per GGML IQ format, used to describe the packed payload in the checkpoint. +_IQ_BLOCK_METADATA = { + QUANTIZATION_IQ1_S: (IQ1_S_BLOCK_SIZE, IQ1_S_BLOCK_BYTES, IQ1_S_EFFECTIVE_BITS), + QUANTIZATION_IQ2_XXS: (IQ2_XXS_BLOCK_SIZE, IQ2_XXS_BLOCK_BYTES, IQ2_XXS_EFFECTIVE_BITS), + QUANTIZATION_IQ2_XS: (IQ2_XS_BLOCK_SIZE, IQ2_XS_BLOCK_BYTES, IQ2_XS_EFFECTIVE_BITS), +} + + def _has_large_fp8_scale(value: torch.Tensor) -> bool: """Return whether an FP8 scale exceeds the recommended threshold.""" return bool(torch.any(value > 0.5)) @@ -474,8 +487,7 @@ def uses_iq_quantization(module) -> bool: if ( weight_quantizer is not None and weight_quantizer.is_enabled - and getattr(weight_quantizer, "num_bits", None) - in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) + and getattr(weight_quantizer, "num_bits", None) in IQ_FORMATS ): return True return any(uses_iq_quantization(child) for _, child in module.named_children()) @@ -515,7 +527,7 @@ def _get_quantization_from_layer(layer, quantizer_attr_names: QuantizerAttrNames return QUANTIZATION_W4A8_AWQ # Handle individual num_bits cases - if weight_quantizer.num_bits in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if weight_quantizer.num_bits in IQ_FORMATS: if weight_quantizer.backend != "ggml": raise ValueError("IQ formats require the built-in 'ggml' quantization backend") # Both exporters return before collecting input_scale and before the pre_quant_scale @@ -781,15 +793,8 @@ def process_layer_quant_config(layer_config_dict): "quant_algo": "MXFP8", "group_size": block_size_value, } - elif v in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): - if v == QUANTIZATION_IQ1_S: - block_size = IQ1_S_BLOCK_SIZE - payload_bytes = IQ1_S_BLOCK_BYTES - effective_bits = IQ1_S_EFFECTIVE_BITS - else: - block_size = IQ2_XS_BLOCK_SIZE - payload_bytes = IQ2_XS_BLOCK_BYTES - effective_bits = IQ2_XS_EFFECTIVE_BITS + elif v in IQ_FORMATS: + block_size, payload_bytes, effective_bits = _IQ_BLOCK_METADATA[v] if block_size_value != block_size: raise ValueError( f"{v.upper()} requires block size {block_size}, got {block_size_value}" diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 28edade5d7b..bced8866c71 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -67,7 +67,7 @@ from modelopt.torch.opt.conversion import ModeloptStateManager, modelopt_state from modelopt.torch.opt.plugins.huggingface import _MODELOPT_STATE_SAVE_NAME from modelopt.torch.quantization import set_quantizer_by_cfg_context -from modelopt.torch.quantization.ggml import quantize_iq1_s, quantize_iq2_xs +from modelopt.torch.quantization.ggml import quantize_iq1_s, quantize_iq2_xs, quantize_iq2_xxs from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper @@ -101,11 +101,13 @@ ) from .quant_format import ( FUSION_FREE_FORMATS, + IQ_FORMATS, QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, QUANTIZATION_FP8_PC_PT, QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS, + QUANTIZATION_IQ2_XXS, QUANTIZATION_MXFP8, QUANTIZATION_NONE, QUANTIZATION_NVFP4, @@ -134,6 +136,14 @@ ) from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry +# One packer per GGML IQ format; the checkpoint stores the packed blocks directly. +_IQ_PACKERS = { + QUANTIZATION_IQ1_S: quantize_iq1_s, + QUANTIZATION_IQ2_XXS: quantize_iq2_xxs, + QUANTIZATION_IQ2_XS: quantize_iq2_xs, +} + + __all__ = ["export_hf_checkpoint", "export_speculative_decoding"] @@ -630,15 +640,13 @@ def _export_quantized_weight( "which dispatches to the streaming writer that materialises weights layer-by-layer." ) - if quantization_format in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if quantization_format in IQ_FORMATS: if weight_name != "weight": raise NotImplementedError( "IQ unified export currently supports modules with a standard 'weight' " f"attribute, got {weight_name!r} on {type(sub_module).__name__}" ) - quantize_iq = ( - quantize_iq1_s if quantization_format == QUANTIZATION_IQ1_S else quantize_iq2_xs - ) + quantize_iq = _IQ_PACKERS[quantization_format] packed_weight, _ = quantize_iq(weight.to(dtype)) setattr(sub_module, weight_name, nn.Parameter(packed_weight, requires_grad=False)) maybe_clear_cuda_cache() diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 215a62ff85e..39be2118651 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -35,7 +35,7 @@ from safetensors.torch import save_file from modelopt import __version__ -from modelopt.torch.quantization.ggml import quantize_iq1_s, quantize_iq2_xs +from modelopt.torch.quantization.ggml import quantize_iq1_s, quantize_iq2_xs, quantize_iq2_xxs from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer from modelopt.torch.utils import import_plugin, warn_rank_0 from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -57,6 +57,7 @@ ) from .plugins.megatron_importer import GPTModelImporter, _get_mamba_conv1d from .quant_format import ( + IQ_FORMATS, KV_CACHE_FP8, KV_CACHE_NVFP4, QUANTIZATION_FP8, @@ -64,6 +65,7 @@ QUANTIZATION_FP8_PB_WO, QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS, + QUANTIZATION_IQ2_XXS, QUANTIZATION_NONE, QUANTIZATION_NVFP4, QUANTIZATION_W4A16_NVFP4, @@ -85,6 +87,14 @@ import transformers from transformers import AutoProcessor +# One packer per GGML IQ format, mirroring the HF exporter's table. +_IQ_PACKERS = { + QUANTIZATION_IQ1_S: quantize_iq1_s, + QUANTIZATION_IQ2_XXS: quantize_iq2_xxs, + QUANTIZATION_IQ2_XS: quantize_iq2_xs, +} + + has_mcore = False with import_plugin("megatron"): from megatron.core.models.gpt import GPTModel @@ -351,7 +361,7 @@ def save_pretrained( quantization = "NVFP4" elif quantization_format == QUANTIZATION_W4A16_NVFP4: quantization = "W4A16_NVFP4" - elif quantization_format in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + elif quantization_format in IQ_FORMATS: quantization = quantization_format.upper() if is_last_stage_main_rank: @@ -1115,7 +1125,7 @@ def _get_quantized_state( self._record_excluded_module(prefix) block_size = get_weight_block_size(module) - is_iq = qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) + is_iq = qformat in IQ_FORMATS name_to_value = self._get_weight_bias( module, dtype, name_to_value, keep_weight_device=is_iq ) @@ -1185,7 +1195,7 @@ def _get_weight_scales(self, quantized_state: dict[str, Any], qformat: str): @staticmethod def _pack_iq_weight(weight: torch.Tensor, qformat: str) -> torch.Tensor: """Pack one ``[out, in]`` weight and return its CPU payload.""" - quantize_iq = quantize_iq1_s if qformat == QUANTIZATION_IQ1_S else quantize_iq2_xs + quantize_iq = _IQ_PACKERS[qformat] packed_weight, _ = quantize_iq(weight) return packed_weight.detach().cpu() @@ -1210,7 +1220,7 @@ def _reject_unsupported_fused_iq_export(qformat: str) -> None: The one gap left is a rank holding no local expert at all, which needs expert-parallel size to exceed the expert count. Worth revisiting if that becomes a supported topology. """ - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat in IQ_FORMATS: raise NotImplementedError( "Fused-MoE IQ export requires a deployment loader that supports " "[num_experts, out_features, in_features // 256, payload_bytes]" @@ -1280,7 +1290,7 @@ def _name_remapping( weight = weight + 1.0 weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat in IQ_FORMATS: self._state_dict.update(self._get_iq_weight_state(prefix + "weight", weight, qformat)) elif weight_scale is None: self._state_dict[prefix + "weight"] = weight @@ -1327,7 +1337,7 @@ def _gated_mlp_slicing( gate_proj_weight = weight[:ffn_hidden_size, :] up_proj_weight = weight[ffn_hidden_size:, :] - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat in IQ_FORMATS: self._state_dict.update( self._get_iq_weight_state(gate_proj_prefix + "weight", gate_proj_weight, qformat) ) @@ -1501,7 +1511,7 @@ def _grouped_mlp_slicing( seen_qformat, seen_block_size = qformat, block_size weight = state_dict[weight_key].to(self.dtype) - if qformat not in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat not in IQ_FORMATS: weight = weight.cpu() weight_scale_cpu = ( weight_scale.detach().cpu().clone() if weight_scale is not None else None @@ -1533,7 +1543,7 @@ def _grouped_mlp_slicing( ] for shard_prefix, shard_weight, shard_scale in shards: - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat in IQ_FORMATS: local_expert_state.update( self._get_iq_weight_state( shard_prefix + "weight", shard_weight, qformat @@ -1702,7 +1712,7 @@ def _take(tensor, index, last_dim, with_gate=False): proj_weights = [_take(weight, s, hidden_size, g) for s, g in zip(slices, gated)] proj_keys = [p + "weight" for p in prefixes] - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat in IQ_FORMATS: for key, weight in zip(proj_keys, proj_weights): self._state_dict.update(self._get_iq_weight_state(key, weight, qformat)) elif weight_scale is None: @@ -1820,7 +1830,7 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False): proj_keys = [p + "weight" for p in proj_prefixes] weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if qformat in IQ_FORMATS: for proj_prefix, proj_weight in zip(proj_prefixes, proj_weights): if proj_prefix in keep_bf16: self._state_dict[proj_prefix + "weight"] = proj_weight.cpu() diff --git a/modelopt/torch/kernels/quantization/ggml/common.cuh b/modelopt/torch/kernels/quantization/ggml/common.cuh index 30818f0da3b..cd234d3423b 100644 --- a/modelopt/torch/kernels/quantization/ggml/common.cuh +++ b/modelopt/torch/kernels/quantization/ggml/common.cuh @@ -49,6 +49,7 @@ constexpr int kScaleBytes = 2; // them cannot drift apart. constexpr int kIq1sEntries = 2048; constexpr int kIq2xsEntries = 512; +constexpr int kIq2xxsEntries = 256; // One CUDA block encodes one GGML block. The reductions below fold over exactly this many warps, // and each kernel static_asserts that its codebook divides evenly among the threads. diff --git a/modelopt/torch/kernels/quantization/ggml/ggml.cpp b/modelopt/torch/kernels/quantization/ggml/ggml.cpp index 4af30719f2f..5ee5b0a3fbf 100644 --- a/modelopt/torch/kernels/quantization/ggml/ggml.cpp +++ b/modelopt/torch/kernels/quantization/ggml/ggml.cpp @@ -23,6 +23,7 @@ at::Tensor iq1_s_pack_cuda(at::Tensor input, at::Tensor grid); at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales); +at::Tensor iq2_xxs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales); namespace { @@ -53,6 +54,23 @@ at::Tensor iq2_xs_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) { return iq2_xs_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous()); } +at::Tensor iq2_xxs_pack(at::Tensor input, at::Tensor grid, at::Tensor scales) { + TORCH_CHECK(input.is_cuda(), "IQ2_XXS packing requires a CUDA input"); + TORCH_CHECK(grid.is_cuda(), "IQ2_XXS packing requires a CUDA grid"); + TORCH_CHECK(scales.is_cuda(), "IQ2_XXS packing requires CUDA scales"); + modelopt::ggml::check_pack_inputs("IQ2_XXS", input, grid, modelopt::ggml::kIq2xxsEntries); + const auto num_blocks = input.numel() / modelopt::ggml::kBlockSize; + TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && + scales.numel() == num_blocks, + "scales must be float16 [numel / 256]"); + // Same rule as IQ2_XS: these bits become the block scale verbatim, and a negative or non-finite + // one packs cleanly while decoding to garbage. + TORCH_CHECK((scales.isfinite() & (scales >= 0)).all().item(), + "scales must be finite and non-negative"); + TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device"); + return iq2_xxs_pack_cuda(input.contiguous(), grid.contiguous(), scales.contiguous()); +} + } // namespace PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { @@ -69,4 +87,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { "Returns uint8 [numel / 256, 74] on the input device. Non-finite input elements are " "treated as zero during packing, and finite elements outside the float32 range " "saturate."); + module.def("iq2_xxs_pack", &iq2_xxs_pack, + "Pack a non-empty float32, float64, float16, or bfloat16 CUDA tensor whose innermost " + "dimension is a multiple of 256. The grid must be float32 [256, 8] holding " + "non-negative codebook magnitudes, and scales must be finite non-negative float16 " + "[numel / 256]. " + "Returns uint8 [numel / 256, 66] on the input device. Non-finite input elements are " + "treated as zero during packing, and finite elements outside the float32 range " + "saturate."); } diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu new file mode 100644 index 00000000000..c501273a669 --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq2_xxs.cu @@ -0,0 +1,240 @@ +/* + * 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. + */ + +#include "common.cuh" + +namespace { + +using namespace modelopt::ggml; + +// The IQ2_XXS packed payload layout and format constants below follow the GGML +// definition at: +// https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h +constexpr int kEntries = kIq2xxsEntries; +constexpr int kGroups = 8; // one 4-bit local scale per 32 values +constexpr int kVectorsPerGroup = 4; // four 8-value codebook vectors per group +constexpr int kLocalScales = 16; +constexpr int kRecordBytes = 8; // four index bytes then one little-endian uint32 +constexpr int kCodeOffset = kScaleBytes; +constexpr int kPayloadBytes = kCodeOffset + kGroups * kRecordBytes; +constexpr float kLocalScaleStep = 0.125f; // Encoded scale is d * (2 * ls + 1) / 8. + +static_assert(kEntries % kThreads == 0, "every thread must visit the same number of entries"); +static_assert((kEntries & (kEntries - 1)) == 0, "the codebook index mask assumes a power of two"); +static_assert(kPayloadBytes == 66, "IQ2_XXS blocks are 66 bytes"); + +// Dot product of |x| against one codebook vector, under the format's even-parity sign rule. +// IQ2_XXS stores seven sign bits per vector and recovers the eighth from their parity, exactly as +// IQ2_XS does, so an odd sign pattern must flip the coordinate with the smallest |x| * q penalty. +__device__ __forceinline__ float even_parity_dot(const float *x, const float *q, bool odd_parity) { + float dot = 0.0f; + float weakest = FLT_MAX; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + const float term = fabsf(x[j]) * q[j]; + dot += term; + weakest = fminf(weakest, term); + } + return odd_parity ? dot - 2.0f * weakest : dot; +} + +template +__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, + const __half *scales, uint8_t *output) { + __shared__ float shared_grid[kEntries * kVectorSize]; + __shared__ float grid_norm[kEntries]; + __shared__ float warp_best[kWarps * kLocalScales]; + __shared__ float group_error[kLocalScales]; + __shared__ unsigned long long warp_keys[kWarps]; + __shared__ int selected_local; + __shared__ uint8_t locals[kGroups]; + __shared__ uint8_t entry_bytes[kGroups * kVectorsPerGroup]; + __shared__ uint8_t sign_bits[kGroups * kVectorsPerGroup]; + + const int tid = threadIdx.x; + const int64_t block = blockIdx.x; + if (block >= num_blocks) + return; + + for (int i = tid; i < kEntries * kVectorSize; i += blockDim.x) + shared_grid[i] = grid[i]; + __syncthreads(); + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + float norm = 0.0f; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + const float q = shared_grid[entry * kVectorSize + j]; + norm = fmaf(q, q, norm); + } + grid_norm[entry] = norm; + } + __syncthreads(); + + const scalar_t *source = input + block * kBlockSize; + uint8_t *payload = output + block * kPayloadBytes; + const __half d_half = scales[block]; + const uint16_t d_bits = __half_as_ushort(d_half); + const float d = __half2float(d_half); + if (!store_block_scale(payload, d_bits)) + return; + +#pragma unroll 1 + for (int group = 0; group < kGroups; ++group) { + if (tid < kLocalScales) + group_error[tid] = 0.0f; + __syncthreads(); + +#pragma unroll + for (int vector = 0; vector < kVectorsPerGroup; ++vector) { + float x[kVectorSize]; + float xnorm = 0.0f; + int negative_count = 0; + const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + x[j] = load_float(source + offset + j); + xnorm = fmaf(x[j], x[j], xnorm); + negative_count += x[j] < 0.0f; + } + const bool odd_parity = (negative_count & 1) != 0; + float local_best[kLocalScales]; +#pragma unroll + for (int local = 0; local < kLocalScales; ++local) + local_best[local] = FLT_MAX; + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + const float *q = shared_grid + entry * kVectorSize; + const float dot = even_parity_dot(x, q, odd_parity); +#pragma unroll + for (int local = 0; local < kLocalScales; ++local) { + const float scale = d * (2 * local + 1) * kLocalScaleStep; + local_best[local] = + fminf(local_best[local], clamped_quant_error(xnorm, dot, grid_norm[entry], scale)); + } + } + block_min_accumulate(local_best, warp_best, group_error); + } + + if (tid == 0) { + selected_local = 0; + float best = group_error[0]; +#pragma unroll + for (int local = 1; local < kLocalScales; ++local) { + if (group_error[local] < best) { + best = group_error[local]; + selected_local = local; + } + } + locals[group] = static_cast(selected_local); + } + __syncthreads(); + const float selected_scale = d * (2 * selected_local + 1) * kLocalScaleStep; + +#pragma unroll + for (int vector = 0; vector < kVectorsPerGroup; ++vector) { + float x[kVectorSize]; + float xnorm = 0.0f; + int negative_count = 0; + const int offset = group * (kVectorsPerGroup * kVectorSize) + vector * kVectorSize; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + x[j] = load_float(source + offset + j); + xnorm = fmaf(x[j], x[j], xnorm); + negative_count += x[j] < 0.0f; + } + const bool odd_parity = (negative_count & 1) != 0; + unsigned long long key = ~0ULL; + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + const float error = clamped_quant_error( + xnorm, even_parity_dot(x, shared_grid + entry * kVectorSize, odd_parity), + grid_norm[entry], selected_scale); + const unsigned long long candidate = error_key(error, entry); + key = candidate < key ? candidate : key; + } + key = block_min_key(key, warp_keys); + if (tid == 0) { + const int entry = static_cast(key & (kEntries - 1)); + const float *q = shared_grid + entry * kVectorSize; + int flip_index = 0; + float weakest = fabsf(x[0]) * q[0]; +#pragma unroll + for (int j = 1; j < kVectorSize; ++j) { + const float term = fabsf(x[j]) * q[j]; + if (term < weakest) { + weakest = term; + flip_index = j; + } + } + int sign_mask = 0; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + bool is_negative = x[j] < 0.0f; + if (odd_parity && j == flip_index) + is_negative = !is_negative; + sign_mask |= static_cast(is_negative) << j; + } + const int slot = group * kVectorsPerGroup + vector; + entry_bytes[slot] = static_cast(entry); + sign_bits[slot] = static_cast(sign_mask & 0x7f); + } + __syncthreads(); + } + } + + // One 8-byte record per group: four index bytes, then a uint32 holding four 7-bit sign + // indices in bits 0..27 and the 4-bit local scale in bits 28..31. + if (tid < kGroups) { + uint8_t *record = payload + kCodeOffset + tid * kRecordBytes; + const int base = tid * kVectorsPerGroup; +#pragma unroll + for (int j = 0; j < kVectorsPerGroup; ++j) + record[j] = entry_bytes[base + j]; + const uint32_t aux = static_cast(sign_bits[base]) | + (static_cast(sign_bits[base + 1]) << 7) | + (static_cast(sign_bits[base + 2]) << 14) | + (static_cast(sign_bits[base + 3]) << 21) | + (static_cast(locals[tid]) << 28); +#pragma unroll + for (int j = 0; j < 4; ++j) + record[kVectorsPerGroup + j] = static_cast(aux >> (8 * j)); + } +} + +} // namespace + +at::Tensor iq2_xxs_pack_cuda(at::Tensor input, at::Tensor grid, at::Tensor scales) { + TORCH_CHECK(input.is_contiguous() && grid.is_contiguous() && scales.is_contiguous(), + "inputs must be contiguous"); + check_pack_inputs("IQ2_XXS", input, grid, kEntries); + const int64_t num_blocks = input.numel() / kBlockSize; + TORCH_CHECK(scales.scalar_type() == at::kHalf && scales.dim() == 1 && + scales.numel() == num_blocks, + "scales must be float16 [numel / 256]"); + TORCH_CHECK(input.get_device() == scales.get_device(), "input and scales must share a device"); + c10::cuda::CUDAGuard guard(input.device()); + auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); + const auto stream = c10::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq2_xxs_pack", [&] { + encode<<(num_blocks), kThreads, 0, stream>>>( + input.data_ptr(), num_blocks, grid.data_ptr(), + reinterpret_cast(scales.data_ptr()), + output.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + return output; +} diff --git a/modelopt/torch/quantization/extensions.py b/modelopt/torch/quantization/extensions.py index 08ca0dbb719..0b2238e452b 100644 --- a/modelopt/torch/quantization/extensions.py +++ b/modelopt/torch/quantization/extensions.py @@ -95,6 +95,7 @@ def get_cuda_ext_ggml(raise_if_failed: bool = False): kernels_ggml / "ggml.cpp", kernels_ggml / "iq1_s.cu", kernels_ggml / "iq2_xs.cu", + kernels_ggml / "iq2_xxs.cu", ], cuda_version_specifiers=">=11.8", fail_msg="GGML IQ CUDA packing extension is unavailable.", diff --git a/modelopt/torch/quantization/ggml/__init__.py b/modelopt/torch/quantization/ggml/__init__.py index 0fbb63ee405..66065b19017 100644 --- a/modelopt/torch/quantization/ggml/__init__.py +++ b/modelopt/torch/quantization/ggml/__init__.py @@ -21,5 +21,11 @@ from .iq1_s import __all__ as _iq1_s_all from .iq2_xs import * from .iq2_xs import __all__ as _iq2_xs_all +from .iq2_xxs import * +from .iq2_xxs import __all__ as _iq2_xxs_all -__all__ = [*_iq1_s_all, *_iq2_xs_all] # noqa: PLE0604 +__all__ = [ # noqa: PLE0604 + *_iq1_s_all, + *_iq2_xs_all, + *_iq2_xxs_all, +] diff --git a/modelopt/torch/quantization/ggml/backend.py b/modelopt/torch/quantization/ggml/backend.py index dc8e15ffb04..d1d51e85607 100644 --- a/modelopt/torch/quantization/ggml/backend.py +++ b/modelopt/torch/quantization/ggml/backend.py @@ -20,6 +20,14 @@ from ..nn.modules.tensor_quantizer import register_quant_backend from .iq1_s import iq1_s_fake_quant from .iq2_xs import iq2_xs_fake_quant +from .iq2_xxs import iq2_xxs_fake_quant + +# One entry per GGML IQ format; adding a format is adding a row here. +_FAKE_QUANTS = { + "iq1_s": iq1_s_fake_quant, + "iq2_xs": iq2_xs_fake_quant, + "iq2_xxs": iq2_xxs_fake_quant, +} def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: @@ -29,11 +37,13 @@ def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: 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": - return iq1_s_fake_quant(inputs, quantizer, **extra_args) - if num_bits == "iq2_xs": - return iq2_xs_fake_quant(inputs, quantizer, **extra_args) - raise ValueError("The ggml backend requires num_bits='iq1_s' or 'iq2_xs'") + # num_bits arrives untyped from the quantizer and is a tuple for scalar formats, + # so narrow before the lookup rather than relying on the dict to reject it. + fake_quant = _FAKE_QUANTS.get(num_bits) if isinstance(num_bits, str) else None + if fake_quant is None: + supported = ", ".join(repr(name) for name in sorted(_FAKE_QUANTS)) + raise ValueError(f"The ggml backend requires num_bits in ({supported})") + return fake_quant(inputs, quantizer, **extra_args) register_quant_backend("ggml", ggml_fake_quant) diff --git a/modelopt/torch/quantization/ggml/codebooks.py b/modelopt/torch/quantization/ggml/codebooks.py index 04402f4b1d0..5e015cd3765 100644 --- a/modelopt/torch/quantization/ggml/codebooks.py +++ b/modelopt/torch/quantization/ggml/codebooks.py @@ -166,6 +166,19 @@ ) +# Compact byte representation of the canonical [256, 8] IQ2_XXS grid. Same +# 8/25/43 magnitude alphabet as IQ2_XS, so it compresses well. +_IQ2_XXS_GRID_ZLIB_B64 = ( + "eNqFVVuS5CAM++UKOoPuf78ZLMmYLLXT1SkSYvyQZGct/egVuDfoFQtPA34M6RXLB74HqRW/N9rm5ZBZmWc5hK+vY+ATII55Oz4X" + "c1+OMZ2PANgWtccrIBOYqYzwEdoH7cfP5EyQSZRJmNpHlbLNMZHCBghKgy/kfh/pOLPgxOgrwUkXoRKrCjo7OpsJlGOXAY+LMhvQ" + "fIGM8/OST2CBvwAWw4JlAk5ZIEC/AZZ7Gz0A54Z5AJ4rkTr1hYuAijSIEEdlXMqryIOQ7paoHJEMMaVPr5LNJGRq90isoTxEVbGC" + "6xCHFaiptZQ2CXVH7Jcidiuqkf2HaGMR8XTsdXUBrOaA1wwdgeAjFOca+ZsFwOu65FtxRlUllEtQfAsLlCKiobobAotS+sSnwyXK" + "0elKQHnT7Wo7wXjB2T3XGr/hOQJmhKxuiaqLvuNeOMKh91qGAl9aAOZkESSOW4KnIS9sJL3/NQJ5NYRiUX9pGrNBgEysAb4ahXru" + "STY/O/mcdAJRbDOSXlzd1WuOymbMktGQnRrkNQF7LBgFz106e/rDkjZmRpFm1KMx3WvRuqemoa19esbHuBgw7VKggELYclb7VH91" + "Xf5hzoD21KEk6CpbkT/vBJbU" +) + + @cache def iq1_s_grid_bytes() -> bytes: """Decoded little-endian int8 bytes of the [2048, 8] IQ1_S ternary table.""" @@ -176,3 +189,9 @@ def iq1_s_grid_bytes() -> bytes: def iq2_xs_grid_bytes() -> bytes: """Decoded bytes of the [512, 8] IQ2_XS magnitude table.""" return base64.b64decode(_IQ2_XS_GRID_B64) + + +@cache +def iq2_xxs_grid_bytes() -> bytes: + """Decoded bytes of the [256, 8] IQ2_XXS magnitude table.""" + return zlib.decompress(base64.b64decode(_IQ2_XXS_GRID_ZLIB_B64)) diff --git a/modelopt/torch/quantization/ggml/iq2_xxs.py b/modelopt/torch/quantization/ggml/iq2_xxs.py new file mode 100644 index 00000000000..01d8bce0a76 --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq2_xxs.py @@ -0,0 +1,286 @@ +# 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_XXS fake quantization and GGML-compatible block packing. + +The encoder performs a single-pass squared-error grid search at a fixed, +empirically anchored super-block scale, mirroring :mod:`.iq2_xs`. Every 256 +logical values become one 66-byte block_iq2_xxs payload: + +* bytes 0..1: little-endian FP16 super-block scale d +* bytes 2..65: eight 8-byte sub-block records, each holding four 8-bit grid + indices followed by a little-endian uint32 of four 7-bit sign indices + (bits 0..27) and one 4-bit local scale (bits 28..31) + +IQ2_XXS differs from IQ2_XS in three ways: the grid is 256 entries rather than +512 so an index needs no high bits, one local scale covers a whole 32-value +sub-block rather than 16 values, and the scale shares a word with the signs +instead of living in a trailing array. + +The canonical 256 x 8 magnitude grid lives in :mod:`.codebooks`, carried from +llama.cpp ggml-common.h revision 9b05354ec6fb58b4e665e9a39ebc40285c015638. +The matching dequantization formula is in ggml-quants.c at the same revision: +https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-quants.c#L2489-L2514 +""" + +import torch + +from ..extensions import get_cuda_ext_ggml +from .codebooks import iq2_xxs_grid_bytes +from .common import ( + GGML_BLOCK_SIZE, + fake_quantize_with_cache, + narrow_to_float32, + validate_block_chunk_size, + validate_packed_weights, + validate_weight, +) + +__all__ = [ + "IQ2_XXS_BLOCK_BYTES", + "IQ2_XXS_BLOCK_SIZE", + "IQ2_XXS_EFFECTIVE_BITS", + "dequantize_iq2_xxs", + "iq2_xxs_fake_quant", + "iq2_xxs_grid", + "quantize_iq2_xxs", +] + +IQ2_XXS_BLOCK_SIZE = GGML_BLOCK_SIZE +IQ2_XXS_BLOCK_BYTES = 66 +IQ2_XXS_EFFECTIVE_BITS = IQ2_XXS_BLOCK_BYTES * 8 / IQ2_XXS_BLOCK_SIZE +_IQ2_XXS_GRID_ENTRIES = 256 +_IQ2_XXS_LOCAL_SCALES = 16 +_IQ2_XXS_GROUPS = 32 +_IQ2_XXS_GROUPS_PER_SUBBLOCK = 4 +_IQ2_XXS_SUBBLOCKS = _IQ2_XXS_GROUPS // _IQ2_XXS_GROUPS_PER_SUBBLOCK +# Largest representable magnitude: grid entry 43 at local scale 15 -> (0.5 + 15) * 0.25. +_IQ2_XXS_NATIVE_MAX = 43 * 31 / 8 +_IQ2_XXS_SCALE_ANCHOR_MIN = 0.65 +_IQ2_XXS_SCALE_ANCHOR_MAX = 0.92 +_IQ2_XXS_PEAK_TO_RMS_TAPER = 0.035 +# Bounds the encode search temporaries; see the note in .iq2_xs. The grid is half the size +# of IQ2_XS's, so the same chunk holds half the search tile. +_DEFAULT_BLOCK_CHUNK_SIZE = 512 +# The decode runs on every forward and is launch-bound, so it takes a much larger chunk. +_DEFAULT_DECODE_CHUNK_SIZE = 4096 +_SCALE_BLOCK_CHUNK_SIZE = 4096 + +_GRID_CACHE: dict[torch.device, torch.Tensor] = {} + + +def iq2_xxs_grid(device: torch.device | str | None = None) -> torch.Tensor: + """Return the canonical IQ2_XXS magnitude grid as float32.""" + resolved_device = torch.device(device or "cpu") + if resolved_device.type == "cuda" and resolved_device.index is None: + resolved_device = torch.device("cuda", torch.cuda.current_device()) + if resolved_device not in _GRID_CACHE: + values = torch.tensor(list(iq2_xxs_grid_bytes()), dtype=torch.float32) + _GRID_CACHE[resolved_device] = values.reshape(_IQ2_XXS_GRID_ENTRIES, 8).to( + device=resolved_device + ) + return _GRID_CACHE[resolved_device] + + +def _predict_iq2_xxs_scales(blocks: torch.Tensor) -> torch.Tensor: + """Predict one FP16 super-block scale for each flattened block.""" + x = narrow_to_float32(blocks) + amax = x.abs().amax(dim=1) + rms = x.square().mean(dim=1).sqrt() + peak_to_rms = torch.where(rms > 0, amax / rms, torch.zeros_like(rms)) + anchor_ratio = (1.0 - _IQ2_XXS_PEAK_TO_RMS_TAPER * peak_to_rms).clamp( + _IQ2_XXS_SCALE_ANCHOR_MIN, _IQ2_XXS_SCALE_ANCHOR_MAX + ) + return ((amax / _IQ2_XXS_NATIVE_MAX) * anchor_ratio).clamp(max=65504.0).to(torch.float16) + + +def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: + """Encode a moderate-size batch of flattened 256-value blocks.""" + x = narrow_to_float32(blocks) + block_count = x.shape[0] + vectors = x.reshape(block_count, _IQ2_XXS_GROUPS, 8) + magnitudes = vectors.abs() + negative = vectors < 0 + # Only seven sign bits are stored; the eighth is their parity, so an odd sign pattern + # has to flip one element. Account for that cost while searching, not after. + odd_parity = negative.sum(dim=-1).remainder(2).bool() + + d = _predict_iq2_xxs_scales(x) + d_float = d.float() + + xnorm = vectors.square().sum(dim=-1) + qnorm = grid.square().sum(dim=-1) + shape = (block_count, _IQ2_XXS_GROUPS, _IQ2_XXS_LOCAL_SCALES) + best_error = torch.full(shape, torch.inf, dtype=torch.float32, device=x.device) + best_entry = torch.zeros(shape, dtype=torch.int64, device=x.device) + # Search the codebook in tiles to cap temporary memory. Strict comparison + # preserves the lowest grid index on equal error. + for entry_start in range(0, _IQ2_XXS_GRID_ENTRIES, 64): + grid_tile = grid[entry_start : entry_start + 64] + products = magnitudes.unsqueeze(2) * grid_tile.reshape(1, 1, -1, 8) + dot = products.sum(dim=-1) + dot = torch.where(odd_parity.unsqueeze(-1), dot - 2.0 * products.amin(dim=-1), dot) + tile_qnorm = qnorm[entry_start : entry_start + 64].reshape(1, 1, -1) + + for local in range(_IQ2_XXS_LOCAL_SCALES): + scale = d_float.reshape(-1, 1, 1) * ((2 * local + 1) / 8.0) + error = ( + xnorm.unsqueeze(-1) - 2.0 * scale * dot + scale.square() * tile_qnorm + ).clamp_min_(0) + tile_error, tile_index = error.min(dim=-1) + replace = tile_error < best_error[:, :, local] + best_error[:, :, local] = torch.where(replace, tile_error, best_error[:, :, local]) + best_entry[:, :, local] = torch.where( + replace, tile_index + entry_start, best_entry[:, :, local] + ) + + # One local scale covers four groups here, against two for IQ2_XS. + subblock_error = best_error.reshape( + block_count, _IQ2_XXS_SUBBLOCKS, _IQ2_XXS_GROUPS_PER_SUBBLOCK, _IQ2_XXS_LOCAL_SCALES + ).sum(dim=2) + selected_local = subblock_error.argmin(dim=-1) + group_local = selected_local.repeat_interleave(_IQ2_XXS_GROUPS_PER_SUBBLOCK, dim=1) + selected_entry = best_entry.gather(2, group_local.unsqueeze(-1)).squeeze(-1) + + selected_grid = grid[selected_entry] + weakest_index = (magnitudes * selected_grid).argmin(dim=-1) + flip = torch.nn.functional.one_hot(weakest_index, num_classes=8).bool() + encoded_negative = negative ^ (flip & odd_parity.unsqueeze(-1)) + sign_bits = torch.arange(8, dtype=torch.int64, device=x.device) + sign_mask = (encoded_negative.to(torch.int64) << sign_bits).sum(dim=-1) & 0x7F + + signs = sign_mask.reshape(block_count, _IQ2_XXS_SUBBLOCKS, _IQ2_XXS_GROUPS_PER_SUBBLOCK) + aux = ( + signs[:, :, 0] + | (signs[:, :, 1] << 7) + | (signs[:, :, 2] << 14) + | (signs[:, :, 3] << 21) + | (selected_local << 28) + ) + + body = torch.empty((block_count, _IQ2_XXS_SUBBLOCKS, 8), dtype=torch.uint8, device=x.device) + body[:, :, 0:4] = ( + selected_entry.reshape(block_count, _IQ2_XXS_SUBBLOCKS, _IQ2_XXS_GROUPS_PER_SUBBLOCK) + ).to(torch.uint8) + for byte in range(4): + body[:, :, 4 + byte] = ((aux >> (8 * byte)) & 0xFF).to(torch.uint8) + + packed = torch.empty((block_count, IQ2_XXS_BLOCK_BYTES), dtype=torch.uint8, device=x.device) + packed[:, :2] = d.contiguous().view(torch.uint8).reshape(block_count, 2) + packed[:, 2:] = body.reshape(block_count, -1) + return torch.where((d_float == 0).unsqueeze(1), 0, packed) + + +@torch.no_grad() +def quantize_iq2_xxs( + weight: torch.Tensor, *, block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack a floating-point weight into GGML-compatible IQ2_XXS blocks. + + Returned shapes are ``[*weight.shape[:-1], weight.shape[-1] // 256, 66]`` + and ``[weight.ndim]``. The packed payload remains on the weight's device; + the logical-shape metadata is kept on CPU. Non-finite input elements are + treated as zero during packing. + """ + validate_weight(weight, "IQ2_XXS") + validate_block_chunk_size(block_chunk_size) + + logical_shape = torch.tensor(weight.shape, dtype=torch.int64) + blocks = weight.contiguous().reshape(-1, IQ2_XXS_BLOCK_SIZE) + grid = iq2_xxs_grid(weight.device) + packed_shape = ( + *weight.shape[:-1], + weight.shape[-1] // IQ2_XXS_BLOCK_SIZE, + IQ2_XXS_BLOCK_BYTES, + ) + if weight.is_cuda: + extension = get_cuda_ext_ggml() + if extension is not None: + scale_chunks = [ + _predict_iq2_xxs_scales(blocks[start : start + _SCALE_BLOCK_CHUNK_SIZE]) + for start in range(0, blocks.shape[0], _SCALE_BLOCK_CHUNK_SIZE) + ] + packed = extension.iq2_xxs_pack(blocks, grid, torch.cat(scale_chunks)) + return packed.reshape(packed_shape), logical_shape + + chunks = [ + _encode_blocks(blocks[start : start + block_chunk_size], grid) + for start in range(0, blocks.shape[0], block_chunk_size) + ] + return torch.cat(chunks).reshape(packed_shape), logical_shape + + +@torch.no_grad() +def dequantize_iq2_xxs( + packed_weights: torch.Tensor, + weight_shape: torch.Tensor, + *, + dtype: torch.dtype = torch.bfloat16, + block_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE, +) -> torch.Tensor: + """Decode GGML-compatible IQ2_XXS payload bytes.""" + shape = validate_packed_weights( + packed_weights, weight_shape, block_bytes=IQ2_XXS_BLOCK_BYTES, format_name="IQ2_XXS" + ) + validate_block_chunk_size(block_chunk_size) + + blocks = packed_weights.contiguous().reshape(-1, IQ2_XXS_BLOCK_BYTES) + bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) + sign_shifts = 7 * torch.arange( + _IQ2_XXS_GROUPS_PER_SUBBLOCK, dtype=torch.int64, device=blocks.device + ) + grid = iq2_xxs_grid(blocks.device) + decoded = torch.empty((blocks.shape[0], IQ2_XXS_BLOCK_SIZE), dtype=dtype, device=blocks.device) + for start in range(0, blocks.shape[0], block_chunk_size): + stop = min(start + block_chunk_size, blocks.shape[0]) + block_chunk = blocks[start:stop] + count = block_chunk.shape[0] + d = block_chunk[:, :2].contiguous().view(torch.float16).reshape(-1).float() + body = block_chunk[:, 2:].reshape(count, _IQ2_XXS_SUBBLOCKS, 8).to(torch.int64) + entries = body[:, :, 0:4] + aux = body[:, :, 4] | (body[:, :, 5] << 8) | (body[:, :, 6] << 16) | (body[:, :, 7] << 24) + # Top nibble is the sub-block scale; the low 28 bits are four 7-bit sign indices. + scales = d.unsqueeze(-1) * (0.5 + ((aux >> 28) & 0xF).float()) * 0.25 + sign_index = (aux.unsqueeze(-1) >> sign_shifts) & 0x7F + 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() + values = grid[entries] * signs + chunk_decoded = values * scales.unsqueeze(-1).unsqueeze(-1) + decoded[start:stop] = chunk_decoded.reshape(-1, IQ2_XXS_BLOCK_SIZE) + return decoded.reshape(shape) + + +def iq2_xxs_fake_quant( + inputs: torch.Tensor, + quantizer, + *, + block_chunk_size: int = _DEFAULT_BLOCK_CHUNK_SIZE, + decode_chunk_size: int = _DEFAULT_DECODE_CHUNK_SIZE, +) -> torch.Tensor: + """IQ2_XXS weight backend for TensorQuantizer, with pass-through backward.""" + if getattr(quantizer, "num_bits", None) != "iq2_xxs": + raise ValueError("The ggml IQ2_XXS backend requires num_bits='iq2_xxs'") + return fake_quantize_with_cache( + inputs, + quantizer, + format_name="iq2_xxs", + block_chunk_size=block_chunk_size, + decode_chunk_size=decode_chunk_size, + quantize=quantize_iq2_xxs, + dequantize=dequantize_iq2_xxs, + ) diff --git a/modelopt_recipes/configs/numerics/iq2_xxs.yaml b/modelopt_recipes/configs/numerics/iq2_xxs.yaml new file mode 100644 index 00000000000..2c7a61a08ab --- /dev/null +++ b/modelopt_recipes/configs/numerics/iq2_xxs.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_XXS weight quantizer using the built-in fixed-scale codebook search. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig +num_bits: iq2_xxs +# Cost metadata for AutoQuantize's compression estimate only; it drives no packing or +# numerics. num_bits is the string "iq2_xxs", so the generic estimator cannot derive the +# storage cost: 66 packed bytes * 8 / 256 weights. Keep in sync with IQ2_XXS_BLOCK_BYTES. +effective_bits: 2.0625 +block_sizes: + -1: 256 +backend: ggml diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.yaml new file mode 100644 index 00000000000..b02b17b417b --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xxs.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_XXS 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_xxs: configs/numerics/iq2_xxs + +algorithm: +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: iq2_xxs + - quantizer_name: '*input_quantizer' + enable: false + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/general/ptq/iq2_xxs.yaml b/modelopt_recipes/general/ptq/iq2_xxs.yaml new file mode 100644 index 00000000000..3ef8e26dd99 --- /dev/null +++ b/modelopt_recipes/general/ptq/iq2_xxs.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_XXS weight-only PTQ. + +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + preset: configs/ptq/presets/model/iq2_xxs + +metadata: + description: >- + Applies uniform GGML-compatible IQ2_XXS 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 fa493a5c945..e7d59fa506e 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 28 general/ptq/ recipes (click to expand) +All 29 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -59,7 +59,8 @@ supported combinations. | `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) | +| `iq2_xxs` | IQ2_XXS W2A16 (2.06 bpw), eligible linears | none | none (no calibration) | +| `iq2_xs` | IQ2_XS W2A16 (2.31 bpw), eligible linears | none | none (no calibration) |
@@ -138,9 +139,10 @@ 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 +- **`iq1_s` / `iq2_xxs` / `iq2_xs`** — GGML-compatible IQ 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. The five + formats trade size against accuracy in order: 1.56, 2.06 and 2.31 bits per weight. 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 diff --git a/tests/_test_utils/torch/quantization/iq_llama_cpp_vectors.py b/tests/_test_utils/torch/quantization/iq_llama_cpp_vectors.py new file mode 100644 index 00000000000..6fe3a9ccd06 --- /dev/null +++ b/tests/_test_utils/torch/quantization/iq_llama_cpp_vectors.py @@ -0,0 +1,162 @@ +# 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. + +"""Conformance vectors captured from a real llama.cpp-quantized checkpoint. + +Each entry holds packed block bytes lifted verbatim from +unsloth/Qwen3.8-27B-GGUF (Qwen3.8-27B-UD-IQ1_S.gguf) together with the values +llama.cpp's own dequantize_row_* produces for them, computed from +ggml-quants.c revision 9b05354ec6fb58b4e665e9a39ebc40285c015638. + +These pin our decoders against bytes we did not produce. A decoder that drifts +from the GGML layout -- a mis-set high bit, a swapped scale nibble, a sign +parity mistake -- fails here even though a round-trip test against our own +encoder would still pass. +""" + +import base64 +import zlib + +import numpy as np + +_VECTORS = { + "iq1_s": { + "source": "blk.0.attn_qkv.weight", + "block_bytes": 50, + "blocks": ( + "eNoBLAHT/sgYPIuVHR5unIa25z3nTMBrrDsowEZpEmhZYzdCh6ybWSOReiY2MD9xWx/BZF4J3Z/LTxnt0w2T+DYRwT0N" + "IIwJoKkC4mejy6P3sqbJWgv0FgfWRjO1zs0h+ZpZymO6MyQ9Ms9aF+DLPrjwY54mJRLDZo2emJfnx3B7rQMfrk8bWdoJ" + "Hc7n5ue6duVmzL9SXOzWw0SmcOsXFEu0ESDUFxuXCM0YGhrzGpYf6sJ0hTc5bGk4hFUG9NtWPKTaovj92oteeEPl9Fc9" + "8BYIR2j/8Hv38eeyZ9fzxt60f8/DEfda/DyI7XWEZ6CXKNxq2nbhUvHVwm4FckftMeVoGA2YAjcSngU6wBLrdmQNavPg" + "js7fDIkqa+Domf9cZoy23tsEc05IMcBUzWTS8th0Vso1lo8=" + ), + "expected": ( + "eNqFVzGIHVUUHTCRgGCaLQJugoFUQYLoFuK8T3YtxDKIBASRDRZiRLCwsRtCYrGsWBjEwuAvAiJJJTEE5wW+pAsJNmJC" + "QNlOkG3EJiiic+7OmT3v/jv/F4d77znn3vf+vPdgtzrwT64+/KwO4TXWiKpFPvVGnrePTyy+dj5XX32RDMhZL+N3DtYF" + "6EGuPeQYfc41xtb0Pl1zzBftTXldc+VWNmy+XlfXd5NFQjXWqsPPHkSds2gWc3iefa4dcPXH/d+PWnPW9CBqr86gx/cQ" + "9AD+d3Nf0d49H30zP4t9GpnveVvD9d1sceVWyZFnVIBTv+/3oKZr+W+v39p/Z+qei7DsnIZZK23z0jdZowJcc/hC8rxp" + "wo96+n6dbdz+mqlaPd0SWjNv/vslA+rzfOFl7me7Gaa/8nU7hvWTL95mPnvzqckiL/Qxj+e1Xn/niUl1+b2asTl3L7Nm" + "ztp7fO7ryD+HJ7dCNFt/5oifvfrMxHsiL3zqJQcvwLy7AzVh54F4+ALuTNbc0OlzUJ/Uxbye1xnao371RXvyPdEcj2Kf" + "fr3+vkZ3Uzn1Rve88DpfMYO/n/WhB7k60ekdmtMd3+daw0MMNTXxsFYtrHXNv2e52n6rBpqjj9OQ93yhCwcv/Vb3nNZz" + "fehhL7lDD7C32kP5bs/ZMOKln76C7/vIz3mOriXDw/s/DIjqz9/NA88Ijh7q5HSu5k5vLn6aq3y8bv76PiES5KmhZlTd" + "uB5+TjEjyPt56M2IHdda3eesmVvdeSNedc4Mda/he7xxoGZsfrucgIEXrYiEr5XX/ojDeu6eF3eYuYsew/0feyvCz72t" + "G7++3Py7sRfPf5wAcqwLKN/3RX7lOLuI1F441QKz979MAGsFtfXHdzJqRPWS9z2LfAO/9cEcplcO3gZYzx7+lCKf6gT7" + "tJ95OEf2HO5Pfgt1avwurKPvsGx+xyHWjECnZc1Zw+M18tqvmu/X9Wze2WmqPnm6RmQOdHcD9yOzVg/4KNd+nct5Omvg" + "5Pz0nPSsNIfHa9pLzp+5+vR+8T0U70jekmpzb1XeV+Qr3lrgtXpzra527+7FZaAP8ecr2XJGArVyWus64FCv3uz+PzlT" + "G078kYacvOoRxx5E5RfNVNy4lgzbl7KBuUafA9Mj9QBomvs68nDm9Ejb5dmwfSkZmCNCVx41AY45Z6hfdeXU68+G5+XP" + "NTpP9fBORPfBz9d7FJ0TzyriNKfH3we9J9GZ64zdu9hzKqDc5lo7RAI6o9foJx95yY316jrarzW9vh6by171oN7ZaAtc" + "/TZZXH2ULQeQs6bm+9hLD2ud4zXE5y92fw/9ng03VyZWK8D7HD72eH+EMR/4e1vdOzhWGzQHPvouGVTzfuiLIn0E5+6v" + "AT5ZnB5rLUcEyCmoax+i9lFnv/d43q9NnlzUp7nfH3X2Rr+NNe+S3o1liLy8o57Tu+x7VY/m+Xfg1/Wzovehb0i1/f1k" + "w85GPeSrj5LVBGpy6iev/eS19nOE/x+gaS2i" + ), + }, + "iq2_xxs": { + "source": "blk.0.attn_gate.weight", + "block_bytes": 66, + "blocks": ( + "eNoBjAFz/h4NhxlMgsOHqpTEeFmmMeZw4V0WDG0g4Cb0pTZTgvOM3r2oaw9MQEyjfPeuF/uVJwiZksqBaESGaYExuAFT" + "KMEBk0MNCEkDIWIGt8+8CRrWkLA2d7ZC/0m3vOOCDi0WUT6Mgv8VisTSWq625wLzKca0Lpy+i2qXT7x15YoTItH5Rzf0" + "hmMMII18mY4im+1dZIo77zVm6MUboI8D79T31BB1QJuOy5qNKSG6TWwbwyDmyNH/Kx6iY4bTDpD+o9HE4ACHp1TGvQwM" + "yGBr0/KEtaKcEJDJ4GDO/xjDFkyPmje7B5AlhYYbWP8QBfPTejik75531FEJNlH+mGkNxp5IHM3KzeC8Sr+o1pkM9A8t" + "sqJ02ZIZDS8LwnHA040JCXza+pedDYUCN94HQvK5kLGPRoUCkZtRJ4dVbQm0J+VwazuGl3caymqF/qftjjYMkfIiNBP2" + "4c7htG3cwS21qXiRkebCWsy4lJiFhE8AP6HXO76mrJpLsAYAURyv5J/4toWGKjGmD6mo9CS+VbBguEDGxLA=" + ), + "expected": ( + "eNp1VE2IVmUUvhhhCMlQQi7Shv60gppVi865NauiRSBRKUg0i4wkSBfiJnJRhn9BhgsJUobQfikGixx9z4XRLHSiEnUR" + "ljpEmJnJV0EOktY5r+fceeb1tng45zznec77c9/vq1btT72dM5rhs1u5WrWfAEk5sRiwukNH3ifXh5fBx9bTdWrr23qe" + "p+q5OTJ06hfWGBBD3/JHauUtZ++LI9fab4AT8Oe+e6Xw5Wizo+4/cYFH7nqyrt6Zz5Z7FM8F+SKXgOnDF56Azm6wDm+s" + "WT37WzL0n95tOVk+sGB+45x4n7yfe6gFT42ewisGn0sWY41q7VuEGDy5goce/ZY1T5qL1cFbrT0DAy8WjY8ckBzk8xjy" + "K1p/O/Ym4g3Z27A341wqQKGD90bxLuHdte8w3h28y6k3+M/HZBge/1UsTpxbx86lyLXHmls/WQ5cqw1/zHJdArQzwZvi" + "rHAm7jhXe444C/ymmvJuyt+ln7eG33R7Xwf3D9TVpmvEIzsEwAMbvkGOgcuIGYWmjYXf6lY3vGM1V8e3UaB38zExziHK" + "JeXYALoEINeFV0CTOZyNPpvZt6wn1eldyTDyxSZWCObeJ+yhx3vTZmg/ELOiT+DhXM+7jSfvXdosOnFZDFoHGGJAXIcc" + "uy/6XHgYZrVrROy9cJ6rrXNFkaPWbe6Rix5qc6/w5PzMrMeacnaBrDl8U8OGatafSUEe23x0+62N5R5zzz1inHvJNbX3" + "pZhD4Qu+9eo3sG9VfCf8Rsm/ef6e/n2lQys+S/w7l+9CYh3smT7eifcJ30nk+IZwLnpiv7hXXL/rPL0vv+LqxcsJQM4R" + "cmf+uLsBPve0FqydY+NN7746OOfr/9G280LbsQ7uRaIOlByuizrcQ3VpphgWvaT/L5dmtphM99fBad5oLhFdk33God/6" + "nk/jUOvI/MC/pzKqXe+SIjksJ+XZ6j56hU2jUbAPCF8Kj8fcD5/NQU3uP/VDUtDBp29sBpZ+JJY7ktZsCM40wbsvddTk" + "c3Ktnto5Rn9gZe8NqR4cTYgubtGWC6I8RwxOIwFaL+g5+uDLvGk39+3h6qdzBEjfL7ildj45SGuBOsM0qm0Kvp1lc2Ce" + "+Ixpa1Xjh9LhndvEooIcaXTLzFp5Dt5r01FE5ZrS11FnXbHG1LqTH5AieaSxQ2uNS4PLfhbN2aLxGtl6xgVcbzFFz/Sh" + "dV+5RrtW7p25SBNbFtYWA4PXHjA+aeSIpvE6dAmjaxqvW84hCOubNvdWvCmBauHepJE1EiCN3XHRuGS15mIwXWjD634J" + "H4AKHbU6P4PvnTqQ/A7E7wnPJ8W9Md5fcZ8E3nZmNXqWxp7ZLcOfzqstap0QyrFrciz75jM+YngMyjWGLl/rhX3GHcA5" + "GL75tHOH1u4E7mXavcG7Sh1v68qdLBmn/r/0d7RkPI3N2S6eZ15rNs4icKapgUugZa8DhDBveGydXK8/TkNHb2gG937I" + "mifNa+DEeIvWC1gNHgpf1K4n5yQ0HXO5+vwTGlizXjy2teWbD5zl6GGuSA4qo/lcF3zuqV86dQ9PkiJZHPl7HxuQ85g6" + "em3/uu/urD2/aobnMUPAdwVT+yv3nfw+cq/Y/zRYz/td90Loi3O389bNZoX0dp3MsDy4APDRy1C+BXgYgLOy7vnXHmpA" + "z3BW2xPjt8Hv7e9B8Eyuy3zcVdnDd9Xhl2rN0aQgjCsf2MkQy555KKJx78+dXaMeexatD7NaPuvOb0iHF++xSBOvLs61" + "R7LcIwWvWtbIxkGeZ4TfezGLY7b7xbWZrx5/m/oHV4nGNDD7iETuoJLXnI0HrtVYzzjvcfRgJsHslDWvH0kOMgy9vEO6" + "ctQZb+hbfX0dfOHj8DifQh/eds0fR5KCAhP3bGSD8cPv/R5RrOcxuU4MrsE5qdCmmAfeqfmf7aPhJ7aKxYDW3Lt9RuN1" + "wujayBn61NGPucnm+dw6fHmNqX3lc/reyPcd5+bYN5w3zkHF+RKcW2JGcUdT/unnSHgm37MU50xwNwR3kwofu4+K+7J+" + "3d7PxhkyfOHrDMsVXMTIuSO/CjGnmNd6Vy6/r8Y63ha8PYk3E3dXvp/y+0CeineG/niTOK/zffn9xLegjiigQw+hrvwu" + "AX+DecZ/P39gQw==" + ), + }, + "iq2_xs": { + "source": "blk.2.ssm_out.weight", + "block_bytes": 74, + "blocks": ( + "eNoBvAFD/ooQA6AAiCN6lmuU9VeqmJ0A5LvON22xpJNoA/QAWAGcDZaWQKyDoBePECMoCX4R2nrRr2MDSLwulHU1HAI0" + "CJIA7qhpU211f1h7XxAAzCTiWUXSqgJfCmAn5Z2fsULmOlXie4P+2Hs6K4JXIkNMR9BMRi0ZH7TGIaAzha65Ev8svrLU" + "1ifnCjik2AR9VzVCZmaZ9mj0FNMeyOaKjppuTNaS4wbCEv+Gcmfqfj5EnRP95Ooudk/83EyYhB7dAMgAhGrkfgw9mb+0" + "iVCEeyeZATYADOqTF9s0MyQ0hTxVjwANNM9T0JRdV5xbUJAyTRtpVWhT61gGwqYlAPxNcKRXB3CIfKqRCfyjf2kPFAEX" + "LmoX/nqif6ITk2ALpVceb5N4Tbpntlv4epaLABIJsJWJzbQAtif3Y3IAVhfoXfpC8GUkXYI35fkDxVr/SmoyAAJJMjgb" + "H0hKjGpgJVmrJd94HzoDLeSgUQN6OSomyKRERK+Iw6ljEmOmdWwqnHp6sDpbIynDWhCVbdJGRFvGr1757kJ49dKyb1ZJ" + "vCaNS9cwmkYajOIXLCaNWoLNhkaeOkv/ABP8eMYiEiIh58/YfcDzwlE=" + ), + "expected": ( + "eNp9Vl+IncUVHymxoUVYIZYktXpRCLQNuiw0LX5n9BZDjSS1C81Dgwveh6ASKElpK5GK3hppI0RYTRM01nYfUtiaNmyb" + "kj/OmeU++CeBPIS2lH0odNOq9UFwzUsjVew5M+fMd+bbmz78+J3zO7+Zb+bMuey6nRthefAd74gF2GH4v/nxu5EAvZV1" + "kWNl1oRTPPmFYRyj+eXR9d49sRmVLebefCRKDMrki8YL/YfPglv7fmBQjMSNIAyO3OI1trpB4yafDMMj30JmQkMxcN7R" + "sH/+VdRYask7+uwvwZ19Kwg3GhPQ6pyr1+jB3fxjuHTyeGQm4BhApwaW5Y71ncydpSc2rupzd17vFd07CKNoBVrjNW7z" + "Why9/lckBolBYuj9d0vUmGF8ID50pw8BY/D3B6NlA/bg7LoYNRbkOu8he9tzGNTfkzqfjf00+5HnuH/7h2WuKU6zrTX5" + "feDY30L7nmjf2bw7VO/dIvXeXZwGAiom7j+kWtIHt97uCbHjK3W+A93Fd+7X7YHtRZ3f9hkGClblC1+bjMy9dy6CMFb+" + "naOGEOa+eTTFwikXoOioXuZ937gumjxIPVgPa+Tz4vfWL3Hjzv+IEYRT3Id7UBiYB48dR4X1cT56kfqwdEdDDMwSo8mD" + "ahIHrvWv7Mn1mVfC6B/0W5x5pZn74V9AcmRmzXCjPvEi54PDfwC3418NIQhKLLXCEuPEmfVefA3FUWP2dfazrB40nqJZ" + "n+RprX7LnMH6LdRTrWeN97DfMD6Yfp3edeMmJAAz5VFjZqmDQdFXph6KssaPWQ9clxztN2R91s29tZedPnb7U/fppz9n" + "gAC/d9PpOHvuu54QVRPdG2/ypfxX6xQgwOUt76Hy9J/u86wTx3Fe8478VtHMDJq74Jg3y+efnsLhA5+PDI4JwBi+vNdL" + "jJXO3pf3Fv9w5q7oFj6H7tQSEkNBm6OB6uy/Vq316L7ZC9396Nu+7HUNjzlbe1ar5/t379rVQDV794QLx9D13gyM4ZfW" + "RM6JPeWNaN7WWFeNcjB7p752zyDf6p6j9dp+5l61/RTuvFHtWdnFwOG9X4kcE3vO+99/IbFoqSYaWM09Os9Axmj7P2P/" + "+QNec4qj1NVT+RPm/wMEZQu8Vry8e9Ni0cwcpXmgXOYCVs0jQT2lD3IXc+8Ul/u19y39kH5pL9Awai+5V5090PTTq6+3" + "84FFd2hDJHhhjS2upXvuue0r5dx/YF1qqumb1O8zQX8T3Z+DoKEcipa5SVh/gy+5XbNwCpx7lxkTu3ebwgc3epMH8UFZ" + "w7UB/R11R5VD4sEeTBpj9kIbZz1U/t52+k1tp/jpIGgSWOMaxwsnUfRQ9LberhlXXzgJsh+O0bt3albd1/Yj9ySa2I/x" + "12tbT1N6nDmvW/4qxTPMkGJmzt1M4y4dhoRcUz2UOvOlZ8Dte5v2eDVUYD3HTcK+t7OPOWshx19koNs2HRPXgCqfvJKh" + "tW3TPs3Q+htiZ/7AzGFTgfUqTvOQMXsByny0CEXjWcrzBWXOdEba+WnnqJ2PUGI7M2kGFilepP78IseZm4TpT6HkzOqp" + "fGYm7O+i/k2N+/3IrJk30pjf5dIzKADzrk31dlx3OyBh4ll6ix0ogCoe3OoJsXi1ntYc5Z6j4RzPXoirtLHxUhA0JR7N" + "ozBI3JR678ZYa0+gAMawBf39/K03fvHIt3lf/bZ+n3U+g/2+PZ/V1ZfP113Dui/75Xq7nvdr61jd3+6t69mbEaveaV/q" + "M4RV32vvRv7XGCDI+cpHvsrbehdY+rq8te1tjmHV2yxvjammdc55ttp5A8nbudR84tm4SrO+ei5rL89p1tDMOeSZnqc+" + "XKaeXwYBJi3r4PoHMljvH8hw8yHV2DPYU895m4P8DnxHx1JjrTdF9d1ggEmbfCmWnHnypdrHnt5ULPPNvbT9zm9Sa7n3" + "0f4mRkfuQ7f/qUBoGHPHT7AGVuOYfaqzx9Qa8VqPao3sX3nM+iD76ffUA919ZC/Q/dQzHH4QCI3btqGANeIg3Iyufhuk" + "lrzGX87YPaf0Qr8XzHkae2a9i/StxJ3zBnOP8s3ky+fVswU6K2qsOZ9fzyxx0Dt13iqYnqK9l/EFe8bphZuiO3YPEntm" + "AuwPQ9ZAwbn6Ol50Xz8HBGa8+sdZ/97lq141yqPWjK/UUv2TN5AAH69Zszj/+2OR4DlXXZlrJm49J36N55/7m5/8yeOe" + "YuCcQXmUvGjsk7jVN9H/5pt2MePCb77sD77zvFeN8shsNGAPQ9aA+8HvUAAWS2f+HU0Nt539meZAsbc1WYOyzosXDPvu" + "/qrzW/DbKMq7yFtqTG/oGeZ9U/1/TYe49w==" + ), + }, +} + + +def packed_blocks(name: str) -> np.ndarray: + """Packed bytes for ``name`` as ``(n_blocks, block_bytes)`` uint8.""" + entry = _VECTORS[name] + raw = zlib.decompress(base64.b64decode(entry["blocks"])) + return np.frombuffer(raw, dtype=np.uint8).reshape(-1, entry["block_bytes"]).copy() + + +def expected_values(name: str) -> np.ndarray: + """llama.cpp's dequantized output for ``packed_blocks(name)``, as ``(n, 256)``.""" + raw = zlib.decompress(base64.b64decode(_VECTORS[name]["expected"])) + return np.frombuffer(raw, dtype=np.float32).reshape(-1, 256).copy() + + +def formats() -> list[str]: + """Every format with captured vectors.""" + return sorted(_VECTORS) diff --git a/tests/examples/hf_ptq/test_llm_ptq.py b/tests/examples/hf_ptq/test_llm_ptq.py index 6ae697b3c09..98906428372 100644 --- a/tests/examples/hf_ptq/test_llm_ptq.py +++ b/tests/examples/hf_ptq/test_llm_ptq.py @@ -76,12 +76,14 @@ 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. + # GGML IQ weight-only, recipe-driven: three formats between 1.56 and 2.31 bits + # per weight. These encoders require every weight's input dimension to be a + # multiple of 256; TinyLlama's 2048 and 5632 both are. None of them calibrates -- + # every recipe sets algorithm: null -- so the only IQ-specific cost is packing each + # weight once on a CUDA encoder and decoding it on each forward, which fits the + # 300s tests/examples default. PTQCommand(recipe="general/ptq/iq1_s", kv_cache_quant="none"), + PTQCommand(recipe="general/ptq/iq2_xxs", kv_cache_quant="none"), PTQCommand(recipe="general/ptq/iq2_xs", kv_cache_quant="none"), PTQCommand(quant="nvfp4"), PTQCommand(quant="nvfp4_awq_lite"), diff --git a/tests/gpu/torch/quantization/test_iq_formats_cuda.py b/tests/gpu/torch/quantization/test_iq_formats_cuda.py new file mode 100644 index 00000000000..894cd1bc51d --- /dev/null +++ b/tests/gpu/torch/quantization/test_iq_formats_cuda.py @@ -0,0 +1,179 @@ +# 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. + +"""CUDA encoders for every GGML IQ format. + +Parametrized over all five rather than per format, so a behaviour asserted for one is +asserted for all. IQ1_S is the odd one out only in its entry point: it derives the block +scale in its own kernel instead of taking a precomputed one. +""" + +import pytest +import torch + +import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module +import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module +import modelopt.torch.quantization.ggml.iq2_xxs as iq2_xxs_module +from modelopt.torch.quantization.extensions import get_cuda_ext_ggml +from modelopt.torch.quantization.ggml import ( + IQ1_S_BLOCK_BYTES, + IQ2_XS_BLOCK_BYTES, + IQ2_XXS_BLOCK_BYTES, +) + +# module, packer name, per-block payload size, whether the packer takes precomputed scales +FORMATS = { + "iq1_s": (iq1_s_module, "iq1_s_pack", IQ1_S_BLOCK_BYTES, False), + "iq2_xxs": (iq2_xxs_module, "iq2_xxs_pack", IQ2_XXS_BLOCK_BYTES, True), + "iq2_xs": (iq2_xs_module, "iq2_xs_pack", IQ2_XS_BLOCK_BYTES, True), +} + + +def _extension(): + extension = get_cuda_ext_ggml(raise_if_failed=True) + assert extension is not None + return extension + + +def _pack(name, weight): + module, packer, _, takes_scales = FORMATS[name] + grid = getattr(module, f"{name}_grid")("cuda") + if not takes_scales: + return getattr(_extension(), packer)(weight, grid) + blocks = weight.contiguous().reshape(-1, 256) + scales = getattr(module, f"_predict_{name}_scales")(blocks) + return getattr(_extension(), packer)(weight, grid, scales) + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_pack_matches_pytorch_encoder_and_is_decodable(monkeypatch, name): + """Byte parity with the reference encoder on a fixed small weight. + + Exact parity is asserted on this input, not in general: the two encoders evaluate the same + squared error with different floating-point fusion, so where two local scales fall within a + float32 ULP they can round to different sides. That happens in roughly one block in several + thousand, costs under 1e-8 of relative reconstruction error, and favours neither encoder -- + see ``test_cuda_pack_reconstruction_matches_pytorch_at_scale``. + """ + module, _, block_bytes, _ = FORMATS[name] + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((8, 512), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed = _pack(name, weight).reshape(8, 2, block_bytes) + packed_again = _pack(name, weight).reshape(8, 2, block_bytes) + monkeypatch.setattr(module, "get_cuda_ext_ggml", lambda: None) + reference, shape = getattr(module, f"quantize_{name}")(weight) + reconstructed = getattr(module, f"dequantize_{name}")(packed, shape) + + assert packed.shape == (8, 2, block_bytes) + assert torch.equal(packed, packed_again) + assert torch.equal(packed, reference) + assert shape.device.type == "cpu" + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.25 + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_pack_is_deterministic_on_one_device(name): + """The contract reproducibility actually needs: same machine, same bytes.""" + generator = torch.Generator(device="cuda").manual_seed(7) + weight = torch.randn((64, 1024), generator=generator, device="cuda", dtype=torch.bfloat16) + first = _pack(name, weight) + for _ in range(3): + assert torch.equal(_pack(name, weight), first) + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_pack_reconstruction_matches_pytorch_at_scale(monkeypatch, name): + """Over many blocks the encoders may disagree on a near-tied scale, but not on quality.""" + module, _, block_bytes, _ = FORMATS[name] + generator = torch.Generator(device="cuda").manual_seed(11) + weight = torch.randn((128, 2048), generator=generator, device="cuda", dtype=torch.bfloat16) + blocks = weight.shape[0] * weight.shape[1] // 256 + + packed = _pack(name, weight).reshape(weight.shape[0], weight.shape[1] // 256, block_bytes) + monkeypatch.setattr(module, "get_cuda_ext_ggml", lambda: None) + reference, shape = getattr(module, f"quantize_{name}")(weight) + + dequantize = getattr(module, f"dequantize_{name}") + target = weight.float() + denominator = target.square().sum() + cuda_error = ( + dequantize(packed, shape, dtype=torch.float32) - target + ).square().sum() / denominator + torch_error = ( + dequantize(reference, shape, dtype=torch.float32) - target + ).square().sum() / denominator + + differing = int((packed != reference).any(dim=-1).sum()) + assert differing <= blocks // 1000, f"{differing} of {blocks} blocks differ" + assert torch.isclose(cuda_error, torch_error, rtol=1e-5) + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_zero_encoding_matches_ggml_block_layout(name): + module, _, block_bytes, _ = FORMATS[name] + weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16) + packed = _pack(name, weight).reshape(1, 1, block_bytes) + shape = torch.tensor(weight.shape, device="cuda") + + assert not packed.any() + assert torch.equal(getattr(module, f"dequantize_{name}")(packed, shape), weight) + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_nonfinite_policy_matches_pytorch_encoder(monkeypatch, name): + module, _, _, _ = FORMATS[name] + weight = torch.zeros((1, 256), device="cuda", dtype=torch.float32) + weight[0, 0] = float("nan") + weight[0, 1] = float("inf") + weight[0, 2] = float("-inf") + + packed, _ = getattr(module, f"quantize_{name}")(weight) + monkeypatch.setattr(module, "get_cuda_ext_ggml", lambda: None) + reference, _ = getattr(module, f"quantize_{name}")(weight) + assert torch.equal(packed, reference) + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_falls_back_to_pytorch_encoder(monkeypatch, name): + """Without the extension the format still packs, through the torch search.""" + module, _, block_bytes, _ = FORMATS[name] + monkeypatch.setattr(module, "get_cuda_ext_ggml", lambda: None) + generator = torch.Generator(device="cuda").manual_seed(5) + weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed, shape = getattr(module, f"quantize_{name}")(weight) + reconstructed = getattr(module, f"dequantize_{name}")(packed, shape) + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + + assert packed.shape == (2, 1, block_bytes) + assert normalized_mse < 0.25 + + +@pytest.mark.parametrize("name", sorted(FORMATS)) +def test_cuda_float64_matches_pytorch_encoder(monkeypatch, name): + module, _, _, _ = FORMATS[name] + generator = torch.Generator(device="cuda").manual_seed(99) + weight = torch.randn((2, 256), generator=generator, device="cuda", dtype=torch.float64) + + packed, _ = getattr(module, f"quantize_{name}")(weight) + monkeypatch.setattr(module, "get_cuda_ext_ggml", lambda: None) + reference, _ = getattr(module, f"quantize_{name}")(weight) + assert torch.equal(reference, packed) diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index 2bf899cf650..bf9ca840b10 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -38,6 +38,8 @@ IQ1_S_EFFECTIVE_BITS, IQ2_XS_BLOCK_SIZE, IQ2_XS_EFFECTIVE_BITS, + IQ2_XXS_BLOCK_SIZE, + IQ2_XXS_EFFECTIVE_BITS, ) @@ -135,6 +137,7 @@ def test_mlp_weight_only_recipe_matches_its_mtq_cfg(recipe_name, cfg_name): ("qformat", "block_size", "effective_bits"), [ ("iq1_s", IQ1_S_BLOCK_SIZE, IQ1_S_EFFECTIVE_BITS), + ("iq2_xxs", IQ2_XXS_BLOCK_SIZE, IQ2_XXS_EFFECTIVE_BITS), ("iq2_xs", IQ2_XS_BLOCK_SIZE, IQ2_XS_EFFECTIVE_BITS), ], ) diff --git a/tests/unit/torch/quantization/test_ggml_backend.py b/tests/unit/torch/quantization/test_ggml_backend.py index 3cfc522d915..a9f93a7aee5 100644 --- a/tests/unit/torch/quantization/test_ggml_backend.py +++ b/tests/unit/torch/quantization/test_ggml_backend.py @@ -82,7 +82,8 @@ def fake_quant(inputs, _quantizer, **kwargs): received.update(kwargs) return inputs - monkeypatch.setattr(backend_module, "iq1_s_fake_quant", fake_quant) + # The dispatcher resolves through its registry, so that is the seam to patch. + monkeypatch.setitem(backend_module._FAKE_QUANTS, "iq1_s", fake_quant) inputs = torch.ones(1, 256) quantizer = SimpleNamespace(num_bits="iq1_s", backend_extra_args=extra_args) diff --git a/tests/unit/torch/quantization/test_iq_formats.py b/tests/unit/torch/quantization/test_iq_formats.py new file mode 100644 index 00000000000..16a6c767615 --- /dev/null +++ b/tests/unit/torch/quantization/test_iq_formats.py @@ -0,0 +1,274 @@ +# 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. + +"""Behaviour every GGML IQ format shares, exercised identically for all five. + +The formats differ only in codebook size, payload layout and bits per weight. Anything that +should hold for one should hold for all, so the contract lives here once and is parametrized +rather than duplicated per format -- a new format is a row in ``FORMATS``. +""" + +import numpy as np +import pytest +import torch +from _test_utils.torch.quantization.iq_llama_cpp_vectors import ( + expected_values, + formats, + packed_blocks, +) + +import modelopt.torch.quantization.ggml.iq1_s as iq1_s_module +import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module +import modelopt.torch.quantization.ggml.iq2_xxs as iq2_xxs_module +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer + +# name -> (module, packed bytes per block, codebook entries, bits per weight) +FORMATS = { + "iq1_s": (iq1_s_module, 50, 2048, 1.5625), + "iq2_xxs": (iq2_xxs_module, 66, 256, 2.0625), + "iq2_xs": (iq2_xs_module, 74, 512, 2.3125), +} +NAMES = sorted(FORMATS) +# IQ1 grids are ternary; IQ2 grids hold the magnitudes 8, 25 and 43. +TERNARY = {"iq1_s"} + + +def _parts(name): + module, block_bytes, entries, bits = FORMATS[name] + return ( + module, + getattr(module, f"quantize_{name}"), + getattr(module, f"dequantize_{name}"), + getattr(module, f"{name}_grid"), + block_bytes, + entries, + bits, + ) + + +@pytest.mark.parametrize("name", NAMES) +def test_canonical_grid(name): + _, _, _, grid_fn, _, entries, _ = _parts(name) + grid = grid_fn() + + assert grid.shape == (entries, 8) + assert grid.dtype == torch.float32 + if name in TERNARY: + assert set(grid.unique().tolist()) == {-1.0, 0.0, 1.0} + assert grid[0].tolist() == [-1.0] * 8 + else: + assert set(grid.unique().tolist()) <= {8.0, 25.0, 43.0} + assert grid[0].tolist() == [8.0] * 8 + + +@pytest.mark.parametrize("name", NAMES) +def test_grid_normalizes_unindexed_cuda_device(monkeypatch, name): + module, _, _, grid_fn, _, _, _ = _parts(name) + cached = torch.empty(0) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 7) + monkeypatch.setitem(module._GRID_CACHE, torch.device("cuda", 7), cached) + + assert grid_fn("cuda") is cached + + +@pytest.mark.parametrize("name", NAMES) +def test_effective_bits_matches_the_payload(name): + module, _, _, _, block_bytes, _, bits = _parts(name) + assert getattr(module, f"{name.upper()}_BLOCK_BYTES") == block_bytes + assert getattr(module, f"{name.upper()}_EFFECTIVE_BITS") == pytest.approx(bits) + assert getattr(module, f"{name.upper()}_BLOCK_SIZE") == 256 + + +@pytest.mark.parametrize("name", NAMES) +def test_round_trip_and_payload_fields(name): + _, quantize, dequantize, _, block_bytes, _, _ = _parts(name) + generator = torch.Generator().manual_seed(1234) + weight = torch.randn((2, 512), generator=generator, dtype=torch.bfloat16) + + packed, shape = quantize(weight) + reconstructed = dequantize(packed, shape) + chunked = dequantize(packed, shape, block_chunk_size=1) + + assert packed.shape == (2, 2, block_bytes) + assert packed.dtype == torch.uint8 + assert reconstructed.shape == weight.shape + assert reconstructed.dtype == torch.bfloat16 + assert torch.equal(reconstructed, chunked) + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.25 + + +@pytest.mark.parametrize("name", NAMES) +def test_decode_is_invariant_to_chunk_size(name): + """Chunking the decode is a memory bound, not a numerical choice.""" + _, quantize, dequantize, _, _, _, _ = _parts(name) + torch.manual_seed(0) + weight = torch.randn(3, 1024, dtype=torch.bfloat16) + packed, shape = quantize(weight) + + 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 + ) + + +@pytest.mark.parametrize("name", NAMES) +def test_zero_block_has_canonical_zero_encoding(name): + _, quantize, dequantize, _, _, _, _ = _parts(name) + weight = torch.zeros((2, 256), dtype=torch.bfloat16) + packed, shape = quantize(weight) + + assert not packed.any() + assert torch.equal(dequantize(packed, shape), weight) + + +@pytest.mark.parametrize("name", NAMES) +def test_underflowed_scale_has_canonical_zero_encoding(name): + """A block whose scale rounds to zero in FP16 packs as all zero bytes. + + The magnitude has to clear every format's threshold at once: the IQ1 formats divide by a + native max of 16.875 against the IQ2 formats' 166.6, so a weight that underflows an IQ2 + scale still lands on an FP16 subnormal for IQ1. + """ + _, quantize, dequantize, _, _, _, _ = _parts(name) + weight = torch.full((1, 256), -1e-8, dtype=torch.bfloat16) + packed, shape = quantize(weight) + + assert not packed.any() + assert torch.equal(dequantize(packed, shape), torch.zeros_like(weight)) + + +@pytest.mark.parametrize("name", NAMES) +def test_requires_complete_last_dimension_blocks(name): + _, quantize, _, _, _, _, _ = _parts(name) + with pytest.raises(ValueError, match="last weight dimension"): + quantize(torch.ones(2, 257)) + + +@pytest.mark.parametrize("name", NAMES) +def test_treats_nonfinite_values_as_zero(name): + _, quantize, _, _, _, _, _ = _parts(name) + weight = torch.zeros(1, 256) + weight[0, 0] = float("nan") + weight[0, 1] = float("inf") + weight[0, 2] = float("-inf") + + packed, _ = quantize(weight) + assert not packed.any() + + +@pytest.mark.parametrize("name", NAMES) +def test_saturates_finite_values_above_the_float32_range(name): + """float64 weights are accepted, so a finite value too large for float32 must saturate. + + Converting before sanitizing would turn it into infinity and then zero, which silently + encodes a large weight as nothing and diverges from the CUDA ``load_float`` policy. + """ + _, quantize, _, _, _, _, _ = _parts(name) + torch.manual_seed(0) + weight = torch.randn(1, 256, dtype=torch.float64) + weight[0, 7] = 1e100 + saturated = weight.clone() + saturated[0, 7] = torch.finfo(torch.float32).max + zeroed = weight.clone() + zeroed[0, 7] = 0.0 + + packed, _ = quantize(weight) + assert torch.equal(packed, quantize(saturated)[0]) + assert not torch.equal(packed, quantize(zeroed)[0]) + + +@pytest.mark.parametrize("name", NAMES) +def test_rejects_invalid_shape_metadata(name): + _, _, dequantize, _, block_bytes, _, _ = _parts(name) + packed = torch.zeros((1, 1, block_bytes), dtype=torch.uint8) + for weight_shape in (torch.tensor([[1, 256]]), torch.tensor([1.0, 256.0]), torch.tensor([257])): + with pytest.raises(ValueError, match=r"weight_shape|logical weight shape"): + dequantize(packed, weight_shape) + + +@pytest.mark.parametrize("name", NAMES) +def test_rejects_scalar_packed_payload(name): + _, _, dequantize, _, _, _, _ = _parts(name) + with pytest.raises(ValueError, match="packed_weights"): + dequantize(torch.tensor(0, dtype=torch.uint8), torch.tensor([1, 256])) + + +@pytest.mark.parametrize("name", NAMES) +def test_fake_quant_has_pass_through_gradient(name): + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=name, block_sizes={-1: 256}, backend="ggml") + ) + weight = torch.randn(2, 256, requires_grad=True) + quantizer(weight).sum().backward() + assert torch.equal(weight.grad, torch.ones_like(weight)) + + +@pytest.mark.parametrize("name", NAMES) +def test_search_is_independent_of_default_dtype(name): + """The encoder must not inherit a global default dtype; it works in float32 throughout.""" + _, quantize, _, _, _, _, _ = _parts(name) + torch.manual_seed(0) + weight = torch.randn(2, 256) + expected, _ = quantize(weight) + try: + torch.set_default_dtype(torch.float64) + assert torch.equal(quantize(weight)[0], expected) + finally: + torch.set_default_dtype(torch.float32) + + +@pytest.mark.parametrize("name", formats()) +def test_decoder_matches_llama_cpp_on_captured_blocks(name): + """Decode bytes we did not produce and match llama.cpp's own output exactly. + + A round-trip against our own encoder cannot catch a layout error that the encoder makes + symmetrically; these blocks come from a real checkpoint, so they can. + """ + _, _, dequantize, _, block_bytes, _, _ = _parts(name) + blocks = packed_blocks(name) + assert blocks.shape[1] == block_bytes + + count = blocks.shape[0] + decoded = dequantize( + torch.from_numpy(blocks).reshape(count, 1, block_bytes), + torch.tensor([count, 256]), + dtype=torch.float32, + ).reshape(count, 256) + + assert np.array_equal(decoded.numpy(), expected_values(name)) + + +def test_every_format_has_conformance_vectors(): + """A new format must arrive with blocks captured from a real llama.cpp checkpoint.""" + assert sorted(formats()) == NAMES + + +def test_error_decreases_with_bit_width(): + """More bits must buy less error, or a format's scale handling is wrong.""" + generator = torch.Generator().manual_seed(7) + weight = torch.randn((4, 1024), generator=generator) + errors = [] + for name in sorted(NAMES, key=lambda n: FORMATS[n][3]): + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=name, block_sizes={-1: 256}, backend="ggml") + ) + errors.append(float((quantizer(weight) - weight).square().mean())) + + assert errors == sorted(errors, reverse=True), dict(zip(sorted(NAMES), errors)) From f6f658c801a050d6907cfdbb8a3f1d839e75ff75 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 22 Sep 2026 22:45:59 +0000 Subject: [PATCH 2/4] Carry IQ block metadata through HF config conversion for every format convert_hf_config.py spelled the IQ family as a tuple of upper-case algorithm strings, at two sites the registry refactor did not reach because it looked for the QUANTIZATION_* constants. IQ2_XXS therefore fell through to the generic branch: its checkpoint config lost group_size, block_payload_bytes, effective_bits and the ggml packing marker, and a mismatched group size was accepted instead of rejected. A consumer needs those fields to walk the payload, so the checkpoint would not decode. Fix it where the family is already defined rather than adding a third spelling. quant_format.py gains IQ_BLOCK_METADATA and IQ_PACKERS next to IQ_FORMATS; convert_hf_config, quant_utils and both exporters resolve through them. That also removes the duplicated packer tables the HF and Megatron exporters were each carrying. The conversion path had no IQ coverage at all, which is why the gap went unnoticed. Add three parametrized tests over the family -- block metadata present, mismatched group size rejected, and the exported geometry matching the codec's own constants -- so a format that stops describing itself fails here. Reverting the fix fails exactly the two IQ2_XXS cases. Also drop the "five formats" wording from ptq.md and the shared test docstrings, which described the whole stack rather than what this change ships; the replacements name no count, so the later formats need not edit them again. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/export/convert_hf_config.py | 22 ++------ modelopt/torch/export/quant_format.py | 34 +++++++++++++ modelopt/torch/export/quant_utils.py | 25 +-------- modelopt/torch/export/unified_export_hf.py | 15 +----- .../torch/export/unified_export_megatron.py | 14 +---- modelopt_recipes/ptq.md | 4 +- .../quantization/test_iq_formats_cuda.py | 6 +-- .../torch/export/test_convert_hf_config.py | 51 +++++++++++++++++++ .../torch/quantization/test_iq_formats.py | 2 +- 9 files changed, 101 insertions(+), 72 deletions(-) diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index 24931b05137..8bf8cb8ef5d 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -19,14 +19,7 @@ from collections import defaultdict from typing import Any -from modelopt.torch.quantization.ggml import ( - IQ1_S_BLOCK_BYTES, - IQ1_S_BLOCK_SIZE, - IQ1_S_EFFECTIVE_BITS, - IQ2_XS_BLOCK_BYTES, - IQ2_XS_BLOCK_SIZE, - IQ2_XS_EFFECTIVE_BITS, -) +from .quant_format import IQ_BLOCK_METADATA, IQ_FORMATS def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) -> dict[str, Any]: @@ -127,15 +120,8 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) }, "weights": {"dynamic": False, "num_bits": 8, "type": "float", "group_size": gs}, } - elif quant_algo in ("IQ1_S", "IQ2_XS"): - if quant_algo == "IQ1_S": - block_size = IQ1_S_BLOCK_SIZE - payload_bytes = IQ1_S_BLOCK_BYTES - effective_bits = IQ1_S_EFFECTIVE_BITS - else: - block_size = IQ2_XS_BLOCK_SIZE - payload_bytes = IQ2_XS_BLOCK_BYTES - effective_bits = IQ2_XS_EFFECTIVE_BITS + elif quant_algo.lower() in IQ_FORMATS: + block_size, payload_bytes, effective_bits = IQ_BLOCK_METADATA[quant_algo.lower()] if group_size not in (None, block_size): raise ValueError(f"{quant_algo} requires group size {block_size}, got {group_size}") # IQ payloads are self-contained blocks, not compressed-tensors integer groups. @@ -239,7 +225,7 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An "targets": ["Linear"], } new_config["config_groups"] = {"group_0": config_group_details} - elif quant_algo_value in ("IQ1_S", "IQ2_XS"): + elif str(quant_algo_value).lower() in IQ_FORMATS: # Forward the caller's group size so a mismatched one is rejected rather than rewritten # to the format's block size. iq_metadata = _quant_algo_to_group_config( diff --git a/modelopt/torch/export/quant_format.py b/modelopt/torch/export/quant_format.py index dea3148404d..44fb4cb76eb 100644 --- a/modelopt/torch/export/quant_format.py +++ b/modelopt/torch/export/quant_format.py @@ -19,6 +19,21 @@ constants, for example, are in :mod:`modelopt.torch.export.trtllm.model_config`. """ +from modelopt.torch.quantization.ggml import ( + IQ1_S_BLOCK_BYTES, + IQ1_S_BLOCK_SIZE, + IQ1_S_EFFECTIVE_BITS, + IQ2_XS_BLOCK_BYTES, + IQ2_XS_BLOCK_SIZE, + IQ2_XS_EFFECTIVE_BITS, + IQ2_XXS_BLOCK_BYTES, + IQ2_XXS_BLOCK_SIZE, + IQ2_XXS_EFFECTIVE_BITS, + quantize_iq1_s, + quantize_iq2_xs, + quantize_iq2_xxs, +) + QUANTIZATION_NONE = None QUANTIZATION_FP8 = "fp8" QUANTIZATION_INT8_SQ = "int8_sq" @@ -51,6 +66,25 @@ } ) +# Block geometry per IQ format: (block size, packed bytes per block, bits per weight). Checkpoint +# metadata spells the algorithm in upper case, so consumers look up +# ``IQ_BLOCK_METADATA[algo.lower()]`` rather than carrying a second spelling of the family. +IQ_BLOCK_METADATA = { + QUANTIZATION_IQ1_S: (IQ1_S_BLOCK_SIZE, IQ1_S_BLOCK_BYTES, IQ1_S_EFFECTIVE_BITS), + QUANTIZATION_IQ2_XXS: (IQ2_XXS_BLOCK_SIZE, IQ2_XXS_BLOCK_BYTES, IQ2_XXS_EFFECTIVE_BITS), + QUANTIZATION_IQ2_XS: (IQ2_XS_BLOCK_SIZE, IQ2_XS_BLOCK_BYTES, IQ2_XS_EFFECTIVE_BITS), +} + + +# The packer each format's checkpoint weights are written with. Both exporters resolve through +# this one mapping so they cannot drift apart. +IQ_PACKERS = { + QUANTIZATION_IQ1_S: quantize_iq1_s, + QUANTIZATION_IQ2_XXS: quantize_iq2_xxs, + QUANTIZATION_IQ2_XS: quantize_iq2_xs, +} + + # Formats whose scales are purely per-module, so export never merges them across the q/k/v # and gate/up groups that share an input. Every other format unifies input_amax (and, for # NVFP4, weight_scale_2) across such a group, which only a whole-model forward can discover. diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 8caa4ce2c19..fb25d23c5e4 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -27,17 +27,6 @@ from modelopt import __version__ from modelopt.torch.models import get_spec, list_all_possible -from modelopt.torch.quantization.ggml import ( - IQ1_S_BLOCK_BYTES, - IQ1_S_BLOCK_SIZE, - IQ1_S_EFFECTIVE_BITS, - IQ2_XS_BLOCK_BYTES, - IQ2_XS_BLOCK_SIZE, - IQ2_XS_EFFECTIVE_BITS, - IQ2_XXS_BLOCK_BYTES, - IQ2_XXS_BLOCK_SIZE, - IQ2_XXS_EFFECTIVE_BITS, -) from modelopt.torch.quantization.model_calib import ( enable_stats_collection, finish_stats_collection, @@ -62,6 +51,7 @@ from ..quantization.nn import NVFP4StaticQuantizer, SequentialQuantizer, TensorQuantizer from .model_utils import TiedWeightMap, get_language_model_from_vl from .quant_format import ( + IQ_BLOCK_METADATA, IQ_FORMATS, KV_CACHE_FP8, KV_CACHE_FP8_K_NVFP4_V, @@ -75,9 +65,6 @@ QUANTIZATION_INT4_AWQ, QUANTIZATION_INT8_SQ, QUANTIZATION_INT8_WO, - QUANTIZATION_IQ1_S, - QUANTIZATION_IQ2_XS, - QUANTIZATION_IQ2_XXS, QUANTIZATION_MXFP4, QUANTIZATION_MXFP8, QUANTIZATION_NONE, @@ -93,14 +80,6 @@ logger = logging.getLogger(__name__) -# Block geometry per GGML IQ format, used to describe the packed payload in the checkpoint. -_IQ_BLOCK_METADATA = { - QUANTIZATION_IQ1_S: (IQ1_S_BLOCK_SIZE, IQ1_S_BLOCK_BYTES, IQ1_S_EFFECTIVE_BITS), - QUANTIZATION_IQ2_XXS: (IQ2_XXS_BLOCK_SIZE, IQ2_XXS_BLOCK_BYTES, IQ2_XXS_EFFECTIVE_BITS), - QUANTIZATION_IQ2_XS: (IQ2_XS_BLOCK_SIZE, IQ2_XS_BLOCK_BYTES, IQ2_XS_EFFECTIVE_BITS), -} - - def _has_large_fp8_scale(value: torch.Tensor) -> bool: """Return whether an FP8 scale exceeds the recommended threshold.""" return bool(torch.any(value > 0.5)) @@ -794,7 +773,7 @@ def process_layer_quant_config(layer_config_dict): "group_size": block_size_value, } elif v in IQ_FORMATS: - block_size, payload_bytes, effective_bits = _IQ_BLOCK_METADATA[v] + block_size, payload_bytes, effective_bits = IQ_BLOCK_METADATA[v] if block_size_value != block_size: raise ValueError( f"{v.upper()} requires block size {block_size}, got {block_size_value}" diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index bced8866c71..1ec13cccb97 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -67,7 +67,6 @@ from modelopt.torch.opt.conversion import ModeloptStateManager, modelopt_state from modelopt.torch.opt.plugins.huggingface import _MODELOPT_STATE_SAVE_NAME from modelopt.torch.quantization import set_quantizer_by_cfg_context -from modelopt.torch.quantization.ggml import quantize_iq1_s, quantize_iq2_xs, quantize_iq2_xxs from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper @@ -102,12 +101,10 @@ from .quant_format import ( FUSION_FREE_FORMATS, IQ_FORMATS, + IQ_PACKERS, QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, QUANTIZATION_FP8_PC_PT, - QUANTIZATION_IQ1_S, - QUANTIZATION_IQ2_XS, - QUANTIZATION_IQ2_XXS, QUANTIZATION_MXFP8, QUANTIZATION_NONE, QUANTIZATION_NVFP4, @@ -136,14 +133,6 @@ ) from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry -# One packer per GGML IQ format; the checkpoint stores the packed blocks directly. -_IQ_PACKERS = { - QUANTIZATION_IQ1_S: quantize_iq1_s, - QUANTIZATION_IQ2_XXS: quantize_iq2_xxs, - QUANTIZATION_IQ2_XS: quantize_iq2_xs, -} - - __all__ = ["export_hf_checkpoint", "export_speculative_decoding"] @@ -646,7 +635,7 @@ def _export_quantized_weight( "IQ unified export currently supports modules with a standard 'weight' " f"attribute, got {weight_name!r} on {type(sub_module).__name__}" ) - quantize_iq = _IQ_PACKERS[quantization_format] + quantize_iq = IQ_PACKERS[quantization_format] packed_weight, _ = quantize_iq(weight.to(dtype)) setattr(sub_module, weight_name, nn.Parameter(packed_weight, requires_grad=False)) maybe_clear_cuda_cache() diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 39be2118651..67564853ce8 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -35,7 +35,6 @@ from safetensors.torch import save_file from modelopt import __version__ -from modelopt.torch.quantization.ggml import quantize_iq1_s, quantize_iq2_xs, quantize_iq2_xxs from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer from modelopt.torch.utils import import_plugin, warn_rank_0 from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -58,14 +57,12 @@ from .plugins.megatron_importer import GPTModelImporter, _get_mamba_conv1d from .quant_format import ( IQ_FORMATS, + IQ_PACKERS, KV_CACHE_FP8, KV_CACHE_NVFP4, QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, QUANTIZATION_FP8_PB_WO, - QUANTIZATION_IQ1_S, - QUANTIZATION_IQ2_XS, - QUANTIZATION_IQ2_XXS, QUANTIZATION_NONE, QUANTIZATION_NVFP4, QUANTIZATION_W4A16_NVFP4, @@ -87,13 +84,6 @@ import transformers from transformers import AutoProcessor -# One packer per GGML IQ format, mirroring the HF exporter's table. -_IQ_PACKERS = { - QUANTIZATION_IQ1_S: quantize_iq1_s, - QUANTIZATION_IQ2_XXS: quantize_iq2_xxs, - QUANTIZATION_IQ2_XS: quantize_iq2_xs, -} - has_mcore = False with import_plugin("megatron"): @@ -1195,7 +1185,7 @@ def _get_weight_scales(self, quantized_state: dict[str, Any], qformat: str): @staticmethod def _pack_iq_weight(weight: torch.Tensor, qformat: str) -> torch.Tensor: """Pack one ``[out, in]`` weight and return its CPU payload.""" - quantize_iq = _IQ_PACKERS[qformat] + quantize_iq = IQ_PACKERS[qformat] packed_weight, _ = quantize_iq(weight) return packed_weight.detach().cpu() diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index e7d59fa506e..ed1d1e003cb 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -141,8 +141,8 @@ activations and tensor-core math are what deliver the throughput. GPT-OSS family (see `examples/gpt-oss`). - **`iq1_s` / `iq2_xxs` / `iq2_xs`** — GGML-compatible IQ 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. The five - formats trade size against accuracy in order: 1.56, 2.06 and 2.31 bits per weight. No calibration data is + `conv1d` and the vision branch stay in BF16 like every other preset. The formats + trade size against accuracy in order: 1.56, 2.06 and 2.31 bits per weight. 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 diff --git a/tests/gpu/torch/quantization/test_iq_formats_cuda.py b/tests/gpu/torch/quantization/test_iq_formats_cuda.py index 894cd1bc51d..70d85c1d0bd 100644 --- a/tests/gpu/torch/quantization/test_iq_formats_cuda.py +++ b/tests/gpu/torch/quantization/test_iq_formats_cuda.py @@ -15,9 +15,9 @@ """CUDA encoders for every GGML IQ format. -Parametrized over all five rather than per format, so a behaviour asserted for one is -asserted for all. IQ1_S is the odd one out only in its entry point: it derives the block -scale in its own kernel instead of taking a precomputed one. +Parametrized over the family rather than written per format, so a behaviour asserted +for one is asserted for all. IQ1_S is the odd one out only in its entry point: it +derives the block scale in its own kernel instead of taking a precomputed one. """ import pytest diff --git a/tests/unit/torch/export/test_convert_hf_config.py b/tests/unit/torch/export/test_convert_hf_config.py index facee1e3a12..8da7092fab2 100644 --- a/tests/unit/torch/export/test_convert_hf_config.py +++ b/tests/unit/torch/export/test_convert_hf_config.py @@ -16,6 +16,7 @@ import pytest from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format +from modelopt.torch.export.quant_format import IQ_BLOCK_METADATA, IQ_FORMATS from modelopt.torch.export.unified_export_hf import _revert_hf_quant_config_names @@ -110,3 +111,53 @@ def failing_mapper(name): "kv_cache_quantized_layers": {"model.bad": {"quant_algo": "FP8"}}, } } + + +@pytest.mark.parametrize("fmt", sorted(IQ_FORMATS)) +def test_iq_config_carries_block_metadata(fmt): + """Every IQ format must describe its packed block, not just name itself. + + A consumer reads group_size and block_payload_bytes to walk the payload, so a format + that falls through to the generic branch produces a checkpoint that cannot be decoded. + """ + block_size, payload_bytes, effective_bits = IQ_BLOCK_METADATA[fmt] + converted = convert_hf_quant_config_format( + { + "producer": {"name": "modelopt", "version": "test"}, + "quantization": {"quant_algo": fmt.upper()}, + } + ) + + assert converted["quant_algo"] == fmt.upper() + assert converted["group_size"] == block_size + assert converted["block_payload_bytes"] == payload_bytes + assert converted["effective_bits"] == pytest.approx(effective_bits) + assert converted["packing"] == "ggml" + # IQ payloads are self-contained blocks, not compressed-tensors integer groups. + assert "config_groups" not in converted + + +@pytest.mark.parametrize("fmt", sorted(IQ_FORMATS)) +def test_iq_config_rejects_mismatched_group_size(fmt): + """A caller's group size is rejected rather than silently rewritten to the block size.""" + block_size, _, _ = IQ_BLOCK_METADATA[fmt] + with pytest.raises(ValueError, match=f"requires group size {block_size}"): + convert_hf_quant_config_format( + { + "producer": {"name": "modelopt", "version": "test"}, + "quantization": {"quant_algo": fmt.upper(), "group_size": block_size // 2}, + } + ) + + +@pytest.mark.parametrize("fmt", sorted(IQ_FORMATS)) +def test_iq_block_metadata_matches_the_codec(fmt): + """The exported geometry is the codec's own, so a checkpoint cannot claim a wrong layout.""" + import modelopt.torch.quantization.ggml as ggml + + block_size, payload_bytes, effective_bits = IQ_BLOCK_METADATA[fmt] + upper = fmt.upper() + assert block_size == getattr(ggml, f"{upper}_BLOCK_SIZE") + assert payload_bytes == getattr(ggml, f"{upper}_BLOCK_BYTES") + assert effective_bits == pytest.approx(getattr(ggml, f"{upper}_EFFECTIVE_BITS")) + assert effective_bits == pytest.approx(payload_bytes * 8 / block_size) diff --git a/tests/unit/torch/quantization/test_iq_formats.py b/tests/unit/torch/quantization/test_iq_formats.py index 16a6c767615..f59c4566531 100644 --- a/tests/unit/torch/quantization/test_iq_formats.py +++ b/tests/unit/torch/quantization/test_iq_formats.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Behaviour every GGML IQ format shares, exercised identically for all five. +"""Behaviour every GGML IQ format shares, exercised identically for each of them. The formats differ only in codebook size, payload layout and bits per weight. Anything that should hold for one should hold for all, so the contract lives here once and is parametrized From 37ae39c322cbf95d2c9193ce8f004c6f90caa8be Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 22 Sep 2026 22:56:13 +0000 Subject: [PATCH 3/4] Cover the mixed IQ export path and correct the changelog Three follow-ups from review. The regression tests added with the metadata fix only exercised a uniform export, where quant_algo sits at the top level. A mixed export groups layers by their distinct configs and routes each through the same helper, so the fix already covers it -- but nothing proved that. Add per-layer equivalents of both tests: block metadata present in the matching config group, and an invalid per-layer group size rejected. Reverting the fix now fails four IQ2_XXS cases rather than two. The changelog still advertised IQ1_M and IQ2_S recipes, which this branch does not ship. My earlier check missed it because the grep that was meant to find stale format names filtered out every line mentioning iq2_xxs, and the offending line mentions both. Move the ggml import in test_convert_hf_config.py to module scope. It was function-local for no reason: the module already imports the export package, which pulls in ggml anyway, so there is no cycle to avoid and nothing to defer. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 2 +- .../torch/export/test_convert_hf_config.py | 56 ++++++++++++++++++- 2 files changed, 55 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1f2ad30b05e..36404803b24 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -14,7 +14,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. -- Add ``iq1_m``, ``iq2_xxs`` and ``iq2_s`` weight-only quantization with CUDA encoders and ``general/ptq`` recipes, completing the GGML IQ formats at one and two bits. The same 256-value block constraint applies as for ``iq1_s`` and ``iq2_xs``. +- Add ``iq2_xxs`` weight-only quantization with a CUDA encoder and a ``general/ptq`` recipe, at 2.0625 bits per weight between ``iq1_s`` and ``iq2_xs``. The same 256-value block constraint applies. - 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/tests/unit/torch/export/test_convert_hf_config.py b/tests/unit/torch/export/test_convert_hf_config.py index 8da7092fab2..bc446c9e042 100644 --- a/tests/unit/torch/export/test_convert_hf_config.py +++ b/tests/unit/torch/export/test_convert_hf_config.py @@ -15,6 +15,7 @@ import pytest +import modelopt.torch.quantization.ggml as ggml from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format from modelopt.torch.export.quant_format import IQ_BLOCK_METADATA, IQ_FORMATS from modelopt.torch.export.unified_export_hf import _revert_hf_quant_config_names @@ -153,11 +154,62 @@ def test_iq_config_rejects_mismatched_group_size(fmt): @pytest.mark.parametrize("fmt", sorted(IQ_FORMATS)) def test_iq_block_metadata_matches_the_codec(fmt): """The exported geometry is the codec's own, so a checkpoint cannot claim a wrong layout.""" - import modelopt.torch.quantization.ggml as ggml - block_size, payload_bytes, effective_bits = IQ_BLOCK_METADATA[fmt] upper = fmt.upper() assert block_size == getattr(ggml, f"{upper}_BLOCK_SIZE") assert payload_bytes == getattr(ggml, f"{upper}_BLOCK_BYTES") assert effective_bits == pytest.approx(getattr(ggml, f"{upper}_EFFECTIVE_BITS")) assert effective_bits == pytest.approx(payload_bytes * 8 / block_size) + + +@pytest.mark.parametrize("fmt", sorted(IQ_FORMATS)) +def test_iq_mixed_precision_config_group_carries_block_metadata(fmt): + """A per-layer IQ config must describe its block too, not only a uniform one. + + Mixed exports route each distinct layer config through the same helper, so a format + missing there loses its geometry for exactly the layers that use it. + """ + block_size, payload_bytes, effective_bits = IQ_BLOCK_METADATA[fmt] + converted = convert_hf_quant_config_format( + { + "producer": {"name": "modelopt", "version": "test"}, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "model.layers.0.mlp.gate_proj": {"quant_algo": fmt.upper()}, + "model.layers.1.mlp.up_proj": {"quant_algo": "FP8"}, + }, + }, + } + ) + + groups = converted["config_groups"].values() + iq_group = next(g for g in groups if g.get("quant_algo") == fmt.upper()) + assert iq_group["group_size"] == block_size + assert iq_group["block_payload_bytes"] == payload_bytes + assert iq_group["effective_bits"] == pytest.approx(effective_bits) + assert iq_group["packing"] == "ggml" + assert iq_group["targets"] == ["model.layers.0.mlp.gate_proj"] + # The FP8 layer keeps its own compressed-tensors scheme. + assert any("weights" in g for g in groups) + + +@pytest.mark.parametrize("fmt", sorted(IQ_FORMATS)) +def test_iq_mixed_precision_rejects_bad_per_layer_group_size(fmt): + """A per-layer group size is validated, not silently rewritten to the block size.""" + block_size, _, _ = IQ_BLOCK_METADATA[fmt] + with pytest.raises(ValueError, match=f"requires group size {block_size}"): + convert_hf_quant_config_format( + { + "producer": {"name": "modelopt", "version": "test"}, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "model.layers.0.mlp.gate_proj": { + "quant_algo": fmt.upper(), + "group_size": block_size // 2, + } + }, + }, + } + ) From cecb5d603cadc59e62abb4267b7415b1b4868347 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Wed, 23 Sep 2026 16:47:01 +0000 Subject: [PATCH 4/4] Run the Megatron IQ export tests for every format The nine IQ tests in the Megatron export suite were hard-wired to IQ1_S and IQ2_XS, so the exporter's IQ2_XXS path -- which this change routes through IQ_FORMATS at nine sites and a shared packer table -- had no Megatron coverage at all. Parametrize all nine over IQ_FORMATS so each format inherits them, and the next two formats get this coverage without touching the file again. Only the list of formats comes from the export tables. Each test resolves what it expects from the codec module itself -- quantize_, dequantize_, _BLOCK_BYTES -- so a wrong entry in IQ_PACKERS or IQ_BLOCK_METADATA cannot make both sides of an assertion agree. Pointing IQ2_XXS at the IQ2_XS packer fails exactly the five payload tests, and leaves the four rejection tests passing, since those raise before anything is packed. This also fixes a flake already on main. The name-remapping test compared the exported payload, decoded, against the fake quantizer's forward output. That output is the straight-through form a + (r - a), which in bf16 never equals r exactly; wherever r is small next to a, as IQ1_S's grid near zero often makes it, the difference exceeds bf16's relative tolerance on r. The test builds an unseeded Linear, so it passed or failed by luck: over 300 seeds it fails 135 times for IQ1_S on main before this change, and at the same rate after. It now checks exact payload bytes against the format's own packer and exact decode against the decoded reference, which fails 0 of 300 seeds for every format. Verified in nvcr.io/nvidia/nemo:26.08, the image CI uses for this suite, with Megatron-Core 0.19.1: 27 passed. Megatron-Core in 26.04 lacks megatron.core.models.hybrid, which the suite's test utilities import. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- .../export/test_unified_export_megatron.py | 115 ++++++++++-------- 1 file changed, 66 insertions(+), 49 deletions(-) diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py index 5659e0481e7..d3f23c7f4f8 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -43,11 +43,12 @@ import modelopt.torch.export.unified_export_megatron as uem import modelopt.torch.quantization as mtq +import modelopt.torch.quantization.ggml as ggml import modelopt.torch.speculative as mtsp from modelopt.torch.export import KV_CACHE_FP8, export_mcore_gpt_to_hf, import_mcore_gpt_from_hf +from modelopt.torch.export.quant_format import IQ_FORMATS from modelopt.torch.export.unified_export_megatron import GPTModelExporter from modelopt.torch.quantization.config import QuantizerAttributeConfig -from modelopt.torch.quantization.ggml import dequantize_iq1_s, dequantize_iq2_xs, quantize_iq2_xs from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.speculative.eagle.default_config import default_eagle_config from modelopt.torch.speculative.plugins.megatron_eagle import _DynamicEagleGPTModel @@ -89,12 +90,17 @@ def _verify_model_quant_config( assert quant_config_dict["kv_cache_quant_algo"] == KV_CACHE_FP8 -@pytest.mark.parametrize( - ("qformat", "payload_bytes", "dequantize"), - [("iq1_s", 50, dequantize_iq1_s), ("iq2_xs", 74, dequantize_iq2_xs)], -) -def test_megatron_name_remapping_exports_iq_payload(qformat, payload_bytes, dequantize): +# Every IQ format the exporter accepts. Only the list of formats comes from the export +# tables; each test resolves what it expects from the codec module itself, so a wrong entry +# in IQ_PACKERS or IQ_BLOCK_METADATA cannot make both sides of an assertion agree. +IQ_FORMAT_NAMES = sorted(IQ_FORMATS) + + +@pytest.mark.parametrize("qformat", IQ_FORMAT_NAMES) +def test_megatron_name_remapping_exports_iq_payload(qformat): """Megatron export writes the same scale-free IQ representation as HF export.""" + payload_bytes = getattr(ggml, f"{qformat.upper()}_BLOCK_BYTES") + dequantize = getattr(ggml, f"dequantize_{qformat}") linear = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16) linear.weight_quantizer = TensorQuantizer( QuantizerAttributeConfig( @@ -112,20 +118,23 @@ def test_megatron_name_remapping_exports_iq_payload(qformat, payload_bytes, dequ exporter._name_remapping(linear, "model.layers.0.mlp.down_proj.") packed_key = "model.layers.0.mlp.down_proj.weight" - assert exporter._state_dict[packed_key].shape == (2, 1, payload_bytes) - assert exporter._state_dict[packed_key].dtype == torch.uint8 - logical_shape = torch.tensor( - [ - *exporter._state_dict[packed_key].shape[:-2], - exporter._state_dict[packed_key].shape[-2] * 256, - ] - ) - reconstructed = dequantize( - exporter._state_dict[packed_key], - logical_shape, - dtype=torch.bfloat16, + packed = exporter._state_dict[packed_key] + assert packed.shape == (2, 1, payload_bytes) + assert packed.dtype == torch.uint8 + # Exact bytes against the format's own packer, as the slicing tests below check. + _assert_iq_payload_matches(qformat, packed, linear.weight) + # And the payload decodes to exactly what the fake quantizer reconstructs. Compare with the + # decoded reference, not the fake-quant forward: that returns the straight-through form + # a + (r - a), which in bf16 differs from r by up to one ULP of a -- enough to fail a + # relative tolerance wherever r is small next to a, as IQ1_S's grid near zero often is. + logical_shape = torch.tensor([*packed.shape[:-2], packed.shape[-2] * 256]) + reference, _ = getattr(ggml, f"quantize_{qformat}")(linear.weight) + torch.testing.assert_close( + dequantize(packed, logical_shape, dtype=torch.bfloat16), + dequantize(reference, logical_shape, dtype=torch.bfloat16), + rtol=0, + atol=0, ) - torch.testing.assert_close(reconstructed, linear.weight_quantizer(linear.weight)) assert exporter.layer_config_dict == { "model.layers.0.mlp.down_proj.quantization": qformat, "model.layers.0.mlp.down_proj.awq_block_size": 256, @@ -167,28 +176,30 @@ def _make_iq_weight(rows): return torch.linspace(-1, 1, rows * 256, dtype=torch.float32).reshape(rows, 256).bfloat16() -def _assert_iq2_payload_matches(packed, logical_weight): - expected, _ = quantize_iq2_xs(logical_weight) +def _assert_iq_payload_matches(qformat, packed, logical_weight): + expected, _ = getattr(ggml, f"quantize_{qformat}")(logical_weight) torch.testing.assert_close(packed, expected.cpu(), rtol=0, atol=0) -def test_megatron_gated_mlp_slicing_exports_iq_payloads(): +@pytest.mark.parametrize("qformat", IQ_FORMAT_NAMES) +def test_megatron_gated_mlp_slicing_exports_iq_payloads(qformat): weight = _make_iq_weight(8) module = SimpleNamespace(config=SimpleNamespace(ffn_hidden_size=4)) exporter = _make_iq_exporter() - exporter._get_quantized_state = lambda *a, **k: ({"weight": weight}, "iq2_xs", 256) + exporter._get_quantized_state = lambda *a, **k: ({"weight": weight}, qformat, 256) exporter._gated_mlp_slicing(module, "model.layers.0.mlp.") - _assert_iq2_payload_matches( - exporter._state_dict["model.layers.0.mlp.gate_proj.weight"], weight[:4] + _assert_iq_payload_matches( + qformat, exporter._state_dict["model.layers.0.mlp.gate_proj.weight"], weight[:4] ) - _assert_iq2_payload_matches( - exporter._state_dict["model.layers.0.mlp.up_proj.weight"], weight[4:] + _assert_iq_payload_matches( + qformat, exporter._state_dict["model.layers.0.mlp.up_proj.weight"], weight[4:] ) -def test_megatron_grouped_mlp_slicing_exports_iq_payloads(): +@pytest.mark.parametrize("qformat", IQ_FORMAT_NAMES) +def test_megatron_grouped_mlp_slicing_exports_iq_payloads(qformat): weight = _make_iq_weight(8) module = SimpleNamespace( num_gemms=1, @@ -199,7 +210,7 @@ def test_megatron_grouped_mlp_slicing_exports_iq_payloads(): exporter = _make_iq_exporter() exporter._get_quantized_state = lambda *a, **k: ( {"weight": module.weight}, - "iq2_xs", + qformat, 256, ) @@ -210,15 +221,16 @@ def test_megatron_grouped_mlp_slicing_exports_iq_payloads(): up_proj_name="up_proj", ) - _assert_iq2_payload_matches( - exporter._state_dict["model.layers.0.mlp.experts.0.gate_proj.weight"], weight[:4] + _assert_iq_payload_matches( + qformat, exporter._state_dict["model.layers.0.mlp.experts.0.gate_proj.weight"], weight[:4] ) - _assert_iq2_payload_matches( - exporter._state_dict["model.layers.0.mlp.experts.0.up_proj.weight"], weight[4:] + _assert_iq_payload_matches( + qformat, exporter._state_dict["model.layers.0.mlp.experts.0.up_proj.weight"], weight[4:] ) -def test_megatron_qkv_slicing_exports_iq_payloads(): +@pytest.mark.parametrize("qformat", IQ_FORMAT_NAMES) +def test_megatron_qkv_slicing_exports_iq_payloads(qformat): weight = _make_iq_weight(8) module = SimpleNamespace( config=SimpleNamespace( @@ -230,7 +242,7 @@ def test_megatron_qkv_slicing_exports_iq_payloads(): ) ) exporter = _make_iq_exporter() - exporter._get_quantized_state = lambda *a, **k: ({"weight": weight}, "iq2_xs", 256) + exporter._get_quantized_state = lambda *a, **k: ({"weight": weight}, qformat, 256) exporter._qkv_slicing(module, "model.layers.0.self_attn.") @@ -241,13 +253,15 @@ def test_megatron_qkv_slicing_exports_iq_payloads(): "v_proj": reshaped[3].reshape(2, 256), } for projection, logical_weight in expected.items(): - _assert_iq2_payload_matches( + _assert_iq_payload_matches( + qformat, exporter._state_dict[f"model.layers.0.self_attn.{projection}.weight"], logical_weight, ) -def test_megatron_gated_delta_net_slicing_exports_iq_payloads(): +@pytest.mark.parametrize("qformat", IQ_FORMAT_NAMES) +def test_megatron_gated_delta_net_slicing_exports_iq_payloads(qformat): weight = _make_iq_weight(12) module = SimpleNamespace( in_proj=object(), @@ -255,15 +269,15 @@ def test_megatron_gated_delta_net_slicing_exports_iq_payloads(): in_proj_split_sections=(2, 2, 2, 2, 2, 2), ) exporter = _make_iq_exporter() - exporter._get_quantized_state = lambda *a, **k: ({"weight": weight}, "iq2_xs", 256) + exporter._get_quantized_state = lambda *a, **k: ({"weight": weight}, qformat, 256) exporter._gated_delta_net_slicing(module, "model.layers.0.mixer.") - _assert_iq2_payload_matches( - exporter._state_dict["model.layers.0.mixer.in_proj_qkv.weight"], weight[:6] + _assert_iq_payload_matches( + qformat, exporter._state_dict["model.layers.0.mixer.in_proj_qkv.weight"], weight[:6] ) - _assert_iq2_payload_matches( - exporter._state_dict["model.layers.0.mixer.in_proj_z.weight"], weight[6:8] + _assert_iq_payload_matches( + qformat, exporter._state_dict["model.layers.0.mixer.in_proj_z.weight"], weight[6:8] ) torch.testing.assert_close( exporter._state_dict["model.layers.0.mixer.in_proj_b.weight"], weight[8:10] @@ -273,7 +287,7 @@ def test_megatron_gated_delta_net_slicing_exports_iq_payloads(): ) -@pytest.mark.parametrize("qformat", ["iq1_s", "iq2_xs"]) +@pytest.mark.parametrize("qformat", IQ_FORMAT_NAMES) def test_megatron_packed_experts_reject_iq_without_deployment_loader(qformat): experts = _make_iq_experts(qformat, "linear_fc2") exporter = _make_iq_exporter() @@ -287,8 +301,9 @@ def test_megatron_packed_experts_reject_iq_without_deployment_loader(qformat): assert exporter._state_dict == {} -def test_megatron_gpt_oss_packed_experts_reject_iq_without_deployment_loader(): - experts = _make_iq_experts("iq2_xs", "linear_fc1", bias=True) +@pytest.mark.parametrize("qformat", IQ_FORMAT_NAMES) +def test_megatron_gpt_oss_packed_experts_reject_iq_without_deployment_loader(qformat): + experts = _make_iq_experts(qformat, "linear_fc1", bias=True) exporter = _make_iq_exporter() with pytest.raises(NotImplementedError, match="Fused-MoE IQ export requires"): @@ -300,12 +315,13 @@ def test_megatron_gpt_oss_packed_experts_reject_iq_without_deployment_loader(): assert exporter._state_dict == {} -def test_megatron_iq_export_rejects_tensor_parallelism(): +@pytest.mark.parametrize("qformat", IQ_FORMAT_NAMES) +def test_megatron_iq_export_rejects_tensor_parallelism(qformat): """IQ packing is intentionally limited to complete TP=1 weights.""" linear = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16) linear.weight_quantizer = TensorQuantizer( QuantizerAttributeConfig( - num_bits="iq2_xs", + num_bits=qformat, block_sizes={-1: 256}, backend="ggml", ) @@ -324,7 +340,8 @@ def test_megatron_iq_export_rejects_tensor_parallelism(): exporter.save_pretrained("unused", "unused") -def test_megatron_iq_export_rejects_pipeline_parallelism(): +@pytest.mark.parametrize("qformat", IQ_FORMAT_NAMES) +def test_megatron_iq_export_rejects_pipeline_parallelism(qformat): """IQ packing requires PP=1 so the fused-MoE rejection reaches every rank. The rejection raises from inside the per-expert loops, so a stage owning no expert would @@ -334,7 +351,7 @@ def test_megatron_iq_export_rejects_pipeline_parallelism(): linear = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16) linear.weight_quantizer = TensorQuantizer( QuantizerAttributeConfig( - num_bits="iq2_xs", + num_bits=qformat, block_sizes={-1: 256}, backend="ggml", )