[OMNIML-5899] Add IQ1_S and IQ2_XS quantization and unified checkpoint export

.pre-commit-config.yaml

@@ -103,6 +103,8 @@ repos:
         exclude: >
           (?x)^(
               modelopt/torch/quantization/utils/calib_utils.py|
+              modelopt/torch/quantization/ggml/iq1_s.py|
+              modelopt/torch/quantization/ggml/iq2_xs.py|
               modelopt/onnx/quantization/operators.py|
               modelopt/onnx/quantization/ort_patching.py|
               modelopt/torch/_deploy/utils/onnx_utils.py|

CHANGELOG.rst

@@ -8,6 +8,7 @@ Changelog
 
 *Quantization*
 
+- Add IQ1_S and IQ2_XS weight quantization and unified checkpoint export from Hugging Face and TP=1 Megatron models. Supported weights are encoded in GGML-compatible 256-value blocks and stored as packed ``uint8`` weight tensors in safetensors.
 - 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.
 
 *Misc*

LICENSE

@@ -250,6 +250,7 @@ the following copyright holders, licensed under the MIT License:
   Copyright (c) 2023 DeepSeek
   Copyright (c) 2025 sgl-project
   Copyright (c) 2026 The DeepSpec Authors
+  Copyright (c) 2023-2026 The ggml authors
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

docs/source/deployment/3_unified_hf.rst

@@ -50,6 +50,36 @@ The unified HF export API supports the following quantization formats:
 4. NVFP4_AWQ - NVIDIA 4-bit floating point with AWQ optimization
 5. INT4_AWQ - 4-bit integer with AWQ optimization
 6. W4A8_AWQ - 4-bit weights and 8-bit activations with AWQ optimization
+7. IQ1_S - 1-bit importance-aware quantization using the GGML block layout
+8. IQ2_XS - 2-bit importance-aware quantization using the GGML block layout
+
+.. note::
+   GGML has no equivalent for ModelOpt's per-tensor FP8 weight-and-activation format. In particular,
+   GGML does not define a first-class FP8 tensor type with the corresponding per-tensor weight and
+   activation scale semantics. Converting a ModelOpt FP8 checkpoint to GGUF therefore requires
+   conversion to another GGML-supported tensor type rather than a lossless FP8 encoding.
+
+IQ weight representation
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+For IQ1_S and IQ2_XS, unified export replaces each floating-point ``<module>.weight`` with a
+``uint8`` tensor containing byte-exact GGML blocks. Its shape is
+``[*logical_shape[:-1], logical_shape[-1] // 256, payload_bytes]``, where ``payload_bytes`` is 50
+for IQ1_S and 74 for IQ2_XS. No separate shape tensor is stored: a loader recovers the logical
+shape as ``[*weight.shape[:-2], weight.shape[-2] * 256]``. This is unambiguous because IQ export
+requires the logical last dimension to be divisible by 256.
+
+Each 74-byte IQ2_XS block represents 256 logical weights:
+
+* bytes 0--1 are the little-endian FP16 super-block scale ``d``;
+* bytes 2--65 are 32 little-endian ``uint16`` codes, one per group of eight weights. Each code
+  contains a 9-bit codebook index and seven stored sign bits; the eighth sign bit is derived from
+  parity; and
+* bytes 66--73 contain sixteen 4-bit local-scale codes, packed two per byte. Each local scale is
+  shared by two adjacent eight-weight groups.
+
+The canonical 512-by-8 IQ2_XS codebook is part of the implementation rather than the checkpoint.
+The complete block therefore costs ``74 * 8 / 256 = 2.3125`` bits per logical weight.
 
 Minimum Framework Versions
 --------------------------

modelopt/torch/export/convert_hf_config.py

@@ -117,6 +117,19 @@ 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"):
+        effective_bits, payload_bytes = (1.5625, 50) if quant_algo == "IQ1_S" else (2.3125, 74)
+        return {
+            "weights": {
+                "dynamic": False,
+                "num_bits": 1 if quant_algo == "IQ1_S" else 2,
+                "effective_bits": effective_bits,
+                "type": "int",
+                "group_size": 256,
+                "packing": "ggml",
+                "block_payload_bytes": payload_bytes,
+            }
+        }
     else:
         warnings.warn(
             f"Unsupported quantization algorithm '{quant_algo}' in "
@@ -209,6 +222,10 @@ 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"):
+        config_group_details = _quant_algo_to_group_config(quant_algo_value, 256)
+        config_group_details["targets"] = ["Linear"]
+        new_config["config_groups"] = {"group_0": config_group_details}
     elif quant_algo_value == "NVFP4_SVD":
         # NVFP4 + SVDQuant: NVFP4 weights/activations plus an AWQ-style
         # pre_quant_scale and a low-rank residual (svdquant_lora_a/b) stored as

modelopt/torch/export/moe_utils.py

@@ -218,7 +218,10 @@ def _export_fused_experts(
             _export_quantized_weight(wrapper, dtype)
 
             proj = nn.Module()
-            proj.weight = wrapper.weight
+            if isinstance(wrapper.weight, nn.Parameter):
+                proj.weight = wrapper.weight
+            else:
+                proj.register_buffer("weight", wrapper.weight)
             for attr in ("weight_scale", "weight_scale_2", "input_scale"):
                 if hasattr(wrapper, attr):
                     proj.register_buffer(attr, getattr(wrapper, attr))

modelopt/torch/export/quant_format.py

@@ -36,11 +36,21 @@
 QUANTIZATION_FP8_PB_REAL = "fp8_pb_real"
 QUANTIZATION_FP8_PB_WO = "fp8_pb_wo"
 QUANTIZATION_FP8_PC_PT = "fp8_pc_pt"
+QUANTIZATION_IQ1_S = "iq1_s"
+QUANTIZATION_IQ2_XS = "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({QUANTIZATION_FP8, QUANTIZATION_NONE, QUANTIZATION_FP8_PB_REAL})
+FUSION_FREE_FORMATS = frozenset(
+    {
+        QUANTIZATION_FP8,
+        QUANTIZATION_IQ1_S,
+        QUANTIZATION_IQ2_XS,
+        QUANTIZATION_NONE,
+        QUANTIZATION_FP8_PB_REAL,
+    }
+)
 
 KV_CACHE_FP8 = "FP8"
 KV_CACHE_INT8 = "INT8"

modelopt/torch/export/quant_utils.py

@@ -61,6 +61,8 @@
     QUANTIZATION_INT4_AWQ,
     QUANTIZATION_INT8_SQ,
     QUANTIZATION_INT8_WO,
+    QUANTIZATION_IQ1_S,
+    QUANTIZATION_IQ2_XS,
     QUANTIZATION_MXFP4,
     QUANTIZATION_MXFP8,
     QUANTIZATION_NONE,
@@ -434,6 +436,11 @@ 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.backend != "psx_luts":
+                raise ValueError("IQ formats require the built-in 'psx_luts' quantization backend")
+            return weight_quantizer.num_bits
+
         if weight_quantizer.num_bits == 4:
             assert len(weight_quantizer.block_sizes) > 0 and weight_quantizer.block_sizes[-1] > 0, (
                 "Invalid block_sizes for INT4 quantizer"
@@ -682,6 +689,14 @@ 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):
+            payload_bytes = 50 if v == QUANTIZATION_IQ1_S else 74
+            layer_config = {
+                "quant_algo": v.upper(),
+                "group_size": 256,
+                "block_payload_bytes": payload_bytes,
+                "packing": "ggml",
+            }
         else:
             layer_config = {"quant_algo": v}
 
@@ -1075,6 +1090,7 @@ def _export_key(key: str) -> str:
         # (pre_quant_scale is the AWQ / NVFP4_AWQ / SVDQuant companion, renamed in the KV-cache pass.)
         weight_suffixes = (
             "weight",
+            "weight_shape",
             "weight_scale",
             "weight_scale_2",
             "input_scale",

modelopt/torch/export/unified_export_hf.py

@@ -59,6 +59,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.nn import SequentialQuantizer, TensorQuantizer
 from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor
 from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper
@@ -100,6 +101,8 @@
     QUANTIZATION_FP8,
     QUANTIZATION_FP8_PB_REAL,
     QUANTIZATION_FP8_PC_PT,
+    QUANTIZATION_IQ1_S,
+    QUANTIZATION_IQ2_XS,
     QUANTIZATION_MXFP8,
     QUANTIZATION_NONE,
     QUANTIZATION_NVFP4,
@@ -622,6 +625,21 @@ 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 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
+        )
+        packed_weight, _ = quantize_iq(weight.to(dtype))
+        delattr(sub_module, weight_name)
+        sub_module.register_buffer("weight", packed_weight)
+        maybe_clear_cuda_cache()
+        return
+
     weight_quantizer: TensorQuantizer | SequentialQuantizer = getattr(
         sub_module, quantizer_attrs.weight_quantizer
     )