[OMNIML-5899] Export IQ checkpoints from HF and Megatron

docs/source/deployment/3_unified_hf.rst

@@ -50,6 +50,54 @@ 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 codebook quantization using the GGML block layout
+8. IQ2_XS - 2-bit codebook 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.
+
+.. note::
+   Megatron IQ export currently requires tensor and pipeline model parallel sizes of 1. Packing
+   happens during export, so a tensor-parallel shard would be packed as if it were a whole
+   weight, and a pipeline stage holding no IQ layer would not reach the same rejection as its
+   peers. Expert parallelism is supported, assuming every expert uses the same format.
+
+.. warning::
+   Megatron fused-MoE IQ export is not currently supported. Its packed tensor would require the
+   deployment consumer to understand
+   ``[num_experts, out_features, in_features // 256, payload_bytes]`` rather than the ordinary HF
+   fused-expert order. The exporter raises ``NotImplementedError`` until a deployment loader owns
+   this layout and is covered by an integration test. Dense and individually named expert weights
+   continue to use the representation above.
+
+The generated configuration records ``quant_method: modelopt``, ``packing: ggml``, the 256-value
+block size, and the payload byte count. IQ payloads are not represented as compressed-tensors
+integer ``weights`` groups because all scales and indices are embedded in each packed block.
+
+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

@@ -19,6 +19,15 @@
 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,
+)
+
 
 def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) -> dict[str, Any]:
     """Map a per-layer quant_algo string to compressed-tensors config group details.
@@ -29,7 +38,8 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None)
 
     Returns:
         Dictionary with ``input_activations`` and ``weights`` entries suitable for
-        a compressed-tensors ``config_groups`` entry.
+        a compressed-tensors ``config_groups`` entry, or ModelOpt-owned metadata for
+        self-contained IQ payloads.
     """
     if quant_algo == "FP8":
         return {
@@ -117,6 +127,26 @@ 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
+        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.
+        # Keep their format marker outside a ``weights`` quantization scheme.
+        return {
+            "quant_algo": quant_algo,
+            "effective_bits": effective_bits,
+            "group_size": block_size,
+            "packing": "ggml",
+            "block_payload_bytes": payload_bytes,
+        }
     else:
         warnings.warn(
             f"Unsupported quantization algorithm '{quant_algo}' in "
@@ -209,6 +239,13 @@ 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"):
+        # 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(
+            quant_algo_value, original_quantization_details.get("group_size")
+        )
+        new_config.update(iq_metadata)
     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/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_FP8_K_NVFP4_V = "FP8_K_NVFP4_V"

modelopt/torch/export/quant_utils.py

@@ -26,6 +26,14 @@
 
 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,
+)
 from modelopt.torch.quantization.model_calib import (
     enable_stats_collection,
     finish_stats_collection,
@@ -62,6 +70,8 @@
     QUANTIZATION_INT4_AWQ,
     QUANTIZATION_INT8_SQ,
     QUANTIZATION_INT8_WO,
+    QUANTIZATION_IQ1_S,
+    QUANTIZATION_IQ2_XS,
     QUANTIZATION_MXFP4,
     QUANTIZATION_MXFP8,
     QUANTIZATION_NONE,
@@ -440,6 +450,36 @@ def get_weight_block_size(module: nn.Module, weight_name: str = "weight") -> int
     return 0
 
 
+def uses_iq_quantization(module) -> bool:
+    """Whether any weight quantizer in ``module`` or its children targets an IQ format.
+
+    ``get_quantization_format`` returns the *first* non-``NONE`` format it finds, so in a
+    mixed-format model IQ layers sitting behind, say, an FP8 layer are invisible to it. Callers
+    that must reject IQ specifically need to see every layer.
+
+    This reads ``num_bits`` directly rather than resolving each layer's full format, so an
+    unrelated unsupported quantizer elsewhere in the model cannot turn the check into an error.
+
+    Known gap, shared with ``get_quantization_format``: ``weight_attr_names`` yields nothing for
+    a TEGroupedLinear, whose parameters are ``weight0..N`` while its quantizer is a single
+    ``GroupedQuantizer`` under ``weight_quantizer``. Neither function sees such a module, so an
+    experts-only IQ model reports no format at all -- not just here. Closing it belongs in
+    ``weight_attr_names``, where it affects every format, rather than in this helper.
+    """
+    for weight_name in weight_attr_names(module):
+        weight_quantizer = representative_weight_quantizer(module, weight_name)
+        # getattr: a SequentialQuantizer has is_enabled but no num_bits, and is never IQ --
+        # IQ is a single quantizer with backend="ggml".
+        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)
+        ):
+            return True
+    return any(uses_iq_quantization(child) for _, child in module.named_children())
+
+
 def get_quantization_format(module) -> str | None:
     """Gets the quantization string.
 
@@ -474,6 +514,24 @@ 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 != "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
+            # handling below, so an enabled activation quantizer would be dropped without a trace
+            # and the checkpoint would load as weight-only. Refuse instead.
+            if input_quantizer is not None and input_quantizer.is_enabled:
+                raise NotImplementedError(
+                    "IQ1_S/IQ2_XS export is weight-only, but this layer has an enabled input "
+                    "quantizer. The GGML block payload carries no activation scale, so the "
+                    "activation quantization would be silently lost."
+                )
+            if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"):
+                raise NotImplementedError(
+                    "IQ1_S/IQ2_XS export does not support an AWQ-style pre_quant_scale."
+                )
+            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"
@@ -722,6 +780,26 @@ 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
+            if block_size_value != block_size:
+                raise ValueError(
+                    f"{v.upper()} requires block size {block_size}, got {block_size_value}"
+                )
+            layer_config = {
+                "quant_algo": v.upper(),
+                "group_size": block_size,
+                "effective_bits": effective_bits,
+                "block_payload_bytes": payload_bytes,
+                "packing": "ggml",
+            }
         else:
             layer_config = {"quant_algo": v}
 

modelopt/torch/export/unified_export_hf.py

@@ -62,6 +62,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
@@ -97,6 +98,8 @@
     QUANTIZATION_FP8,
     QUANTIZATION_FP8_PB_REAL,
     QUANTIZATION_FP8_PC_PT,
+    QUANTIZATION_IQ1_S,
+    QUANTIZATION_IQ2_XS,
     QUANTIZATION_MXFP8,
     QUANTIZATION_NONE,
     QUANTIZATION_NVFP4,
@@ -621,6 +624,20 @@ 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))
+        setattr(sub_module, weight_name, nn.Parameter(packed_weight, requires_grad=False))
+        maybe_clear_cuda_cache()
+        return
+
     weight_quantizer: TensorQuantizer | SequentialQuantizer = getattr(
         sub_module, quantizer_attrs.weight_quantizer
     )