[OMNIML-5899] Add IQ unified checkpoint export

CHANGELOG.rst

@@ -8,8 +8,8 @@ Changelog
 
 *Quantization*
 
-- Add IQ1_S and IQ2_XS weight-only fake quantization with GGML-compatible 256-value block encoders and built-in ``iq1_s`` / ``iq2_xs`` PTQ recipes.
-- 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.
+- Add IQ1_S and IQ2_XS weight-only fake quantization with GGML-compatible 256-value block encoders and built-in ``iq1_s`` / ``iq2_xs`` PTQ recipes. Unified checkpoint export stores supported Hugging Face and TP=1 Megatron weights as packed ``uint8`` tensors.
+- 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, NVFP4, IQ1_S and IQ2_XS on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts.
 - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration.
 - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/onnx_ptq/bevformer>`_ for details.
 

docs/source/deployment/3_unified_hf.rst

@@ -50,6 +50,47 @@ 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 - GGML-compatible block format, 1.5625 bits per weight
+8. IQ2_XS - GGML-compatible block format, 2.3125 bits per weight
+
+.. 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 50-byte IQ1_S block represents 256 logical weights:
+
+* bytes 0--1 are the little-endian FP16 super-block scale ``d``;
+* bytes 2--33 are the low eight bits of 32 codebook indices, one per group of eight weights; and
+* bytes 34--49 are eight little-endian ``uint16`` metadata words, one per group of 32 weights. Each
+  word stores the remaining three bits of four codebook indices, a 3-bit local-scale code, and one
+  delta-shift bit.
+
+The canonical 2048-by-8 IQ1_S codebook is part of the implementation rather than the checkpoint.
+The complete block therefore costs ``50 * 8 / 256 = 1.5625`` bits per logical weight.
+
+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,8 @@
 from collections import defaultdict
 from typing import Any
 
+from .quant_format import IQ_FORMATS, iq_format_spec
+
 
 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.
@@ -117,6 +119,20 @@ 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.lower() in IQ_FORMATS:
+        spec = iq_format_spec(quant_algo)
+        # ``num_bits`` and ``type`` are nominal format labels, not a scalar integer decode
+        # contract. Consumers must honor ``packing`` and ``block_payload_bytes``.
+        return {
+            "weights": {
+                "dynamic": False,
+                "num_bits": spec["num_bits"],
+                "type": "int",
+                "group_size": spec["group_size"],
+                "packing": spec["packing"],
+                "block_payload_bytes": spec["block_payload_bytes"],
+            }
+        }
     else:
         warnings.warn(
             f"Unsupported quantization algorithm '{quant_algo}' in "
@@ -209,6 +225,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 and quant_algo_value.lower() in IQ_FORMATS:
+        config_group_details = _quant_algo_to_group_config(quant_algo_value)
+        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/hf_export_handlers.py

@@ -38,12 +38,14 @@ def _export_weight(
     module: nn.Module,
     ctx: ExportContext,
     weight_name: str = "weight",
+    module_name: str | None = None,
 ) -> None:
     # Imported lazily to avoid a cycle: unified_export_hf imports this module to
     # install the built-in handlers while retaining this legacy helper's import path.
     from .unified_export_hf import _export_quantized_weight
 
-    _export_quantized_weight(module, ctx.dtype, weight_name)
+    describe_as = f"{module_name}.{weight_name}" if module_name else None
+    _export_quantized_weight(module, ctx.dtype, weight_name, describe_as)
 
 
 # Preparation handlers are registered in the same precedence as the legacy MoE prepass.
@@ -155,7 +157,7 @@ def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContex
     Tied experts are packed independently and their duplicate keys are dropped by name
     in postprocess_state_dict; no per-module dedup cache is used.
     """
-    _export_fused_experts(module, ctx.dtype)
+    _export_fused_experts(module, ctx.dtype, module_name=name)
 
 
 @ExportModuleRegistry.register(predicate=is_quantlinear)
@@ -168,7 +170,7 @@ def _export_quant_linear(name: str, module: nn.Module, ctx: ExportContext) -> No
     if get_quantization_format(module) == QUANTIZATION_NONE:
         return
     try:
-        _export_weight(module, ctx)
+        _export_weight(module, ctx, module_name=name)
     except AssertionError as e:
         raise AssertionError(
             f"Failed to export module '{name}' (type={type(module).__name__}): {e}"
@@ -198,7 +200,7 @@ def _export_quant_embedding(name: str, module: nn.Module, ctx: ExportContext) ->
         )
         return
     try:
-        _export_weight(module, ctx)
+        _export_weight(module, ctx, module_name=name)
     except AssertionError as e:
         raise AssertionError(
             f"Failed to export embedding '{name}' (type={type(module).__name__}): {e}"
@@ -224,4 +226,4 @@ def _export_bmm_experts(name: str, module: nn.Module, ctx: ExportContext) -> Non
         quantizer_attrs=["gate_up_proj_input_quantizer", "down_proj_input_quantizer"],
     )
     for weight_name in ["gate_up_proj", "down_proj"]:
-        _export_weight(module, ctx, weight_name)
+        _export_weight(module, ctx, weight_name, module_name=name)

modelopt/torch/export/layerwise_export.py

@@ -42,6 +42,7 @@
 from .quant_utils import (
     _get_kv_cache_postprocess_config,
     _postprocess_single_tensor,
+    _validate_iq_export_support,
     get_quant_config,
     get_quantization_format,
 )
@@ -137,6 +138,7 @@ def assert_formats_supported(module: nn.Module, scope: str) -> None:
 def assert_layerwise_export_supported(model: nn.Module) -> None:
     """Raise unless per-layer export is valid for this model."""
     assert_formats_supported(model, "before calibration")
+    _validate_iq_export_support(model)
 
     tied = _tied_quantized_modules(model)
     if tied:

modelopt/torch/export/moe_utils.py

@@ -15,13 +15,15 @@
 
 """Utilities for Mixture-of-Experts (MoE) model export."""
 
-import copy
 import warnings
+from copy import deepcopy
 from pathlib import Path
 
 import torch
 import torch.nn as nn
 
+from .quant_format import IQ_FORMATS
+
 
 def _delete_fused_moe_source_attrs(module: nn.Module) -> None:
     """Remove the 3-D fused source params and per-expert quantizer ModuleLists.
@@ -48,6 +50,7 @@ def _delete_fused_moe_source_attrs(module: nn.Module) -> None:
 def _export_fused_experts(
     module: nn.Module,
     dtype: torch.dtype,
+    module_name: str | None = None,
 ) -> None:
     """Split fused MoE expert weights and export per-expert quantization scales.
 
@@ -114,6 +117,7 @@ def _export_fused_experts(
         first_proj_q = first_proj_weight_quantizers[idx]
         if (
             is_gated
+            and getattr(first_proj_q, "num_bits", None) not in IQ_FORMATS
             and getattr(first_proj_q, "is_enabled", False)
             and (
                 not hasattr(first_proj_q, "_amax")
@@ -156,17 +160,23 @@ def _export_fused_experts(
                 if uses_first_proj_quantizers
                 else module.down_proj_weight_quantizers[idx]
             )
+            is_iq = getattr(w_quantizer_src, "num_bits", None) in IQ_FORMATS
             i_quantizer = first_proj_input_q if uses_first_proj_quantizers else down_input_q
 
-            # gate/up share a weight quantizer — clone so each gets independent amax.
+            # Formats that export amax need independent gate/up copies before slicing it. IQ
+            # packing derives metadata from each weight directly, so retaining the source avoids
+            # duplicating a potentially populated reconstruction cache.
             w_quantizer = (
-                copy.deepcopy(w_quantizer_src) if uses_first_proj_quantizers else w_quantizer_src
+                deepcopy(w_quantizer_src)
+                if uses_first_proj_quantizers and not is_iq
+                else w_quantizer_src
             )
 
             # For per-channel amax (dim >= 1), proportionally slice dim-0
             # to match the split weight.
             if (
-                hasattr(w_quantizer, "_amax")
+                not is_iq
+                and hasattr(w_quantizer, "_amax")
                 and w_quantizer._amax is not None
                 and w_quantizer._amax.dim() >= 1
             ):
@@ -194,7 +204,8 @@ def _export_fused_experts(
 
             # If the weight quantizer was never calibrated, compute amax from weights.
             if (
-                hasattr(w_quantizer, "is_enabled")
+                not is_iq
+                and hasattr(w_quantizer, "is_enabled")
                 and w_quantizer.is_enabled
                 and (
                     not hasattr(w_quantizer, "_amax")
@@ -215,7 +226,10 @@ def _export_fused_experts(
             wrapper.weight_quantizer = w_quantizer
             wrapper.input_quantizer = i_quantizer
 
-            _export_quantized_weight(wrapper, dtype)
+            describe_as = (
+                f"{module_name}.{idx}.{proj_name}.weight" if module_name is not None else None
+            )
+            _export_quantized_weight(wrapper, dtype, describe_as=describe_as)
 
             proj = nn.Module()
             proj.weight = wrapper.weight

modelopt/torch/export/plugins/vllm_fakequant_megatron.py

@@ -71,6 +71,8 @@ def _merge_quantizer_states(objs: list) -> dict:
 class VllmFqGPTModelExporter(GPTModelExporter):
     """VLLM fakequant GPTModel exporter."""
 
+    _packs_iq_weights = False
+
     @staticmethod
     def _pop_quantizer_keys(state_dict: dict) -> None:
         """Remove quantizer tensors from an export shard (OrderedDict-safe)."""
@@ -114,6 +116,7 @@ def _get_quantized_state(
         module: torch.nn.Module,
         dtype: torch.dtype = torch.float16,
         prefix: str = "",
+        describe_as: str | None = None,
     ) -> tuple[dict[str, torch.Tensor], str, int]:
         """Return a state_dict, quantization format, and block_size of the module.
 
@@ -124,6 +127,8 @@ def _get_quantized_state(
         Args:
             module: The target module to perform real quantization.
             dtype: The default data type.
+            prefix: The exported module prefix.
+            describe_as: Optional diagnostic label accepted for base-exporter compatibility.
 
         Returns:
             Tuple: state_dict, quantization format, and block_size of the module.

modelopt/torch/export/quant_format.py

@@ -19,6 +19,15 @@
 constants, for example, are in :mod:`modelopt.torch.export.trtllm.model_config`.
 """
 
+from typing import Any
+
+from modelopt.torch.quantization.ggml import (
+    IQ1_S_BLOCK_BYTES,
+    IQ1_S_BLOCK_SIZE,
+    IQ2_XS_BLOCK_BYTES,
+    IQ2_XS_BLOCK_SIZE,
+)
+
 QUANTIZATION_NONE = None
 QUANTIZATION_FP8 = "fp8"
 QUANTIZATION_INT8_SQ = "int8_sq"
@@ -36,11 +45,47 @@
 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"
+
+IQ_FORMATS = frozenset({QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS})
+IQ_FORMAT_SPECS: dict[str, dict[str, Any]] = {
+    QUANTIZATION_IQ1_S: {
+        "quant_algo": "IQ1_S",
+        "num_bits": 1,
+        "group_size": IQ1_S_BLOCK_SIZE,
+        "block_payload_bytes": IQ1_S_BLOCK_BYTES,
+        "packing": "ggml",
+    },
+    QUANTIZATION_IQ2_XS: {
+        "quant_algo": "IQ2_XS",
+        "num_bits": 2,
+        "group_size": IQ2_XS_BLOCK_SIZE,
+        "block_payload_bytes": IQ2_XS_BLOCK_BYTES,
+        "packing": "ggml",
+    },
+}
+
+
+def iq_format_spec(quantization_format: str) -> dict[str, Any]:
+    """Return a copy of the canonical export metadata for one IQ format."""
+    try:
+        return IQ_FORMAT_SPECS[quantization_format.lower()].copy()
+    except KeyError:
+        raise ValueError(f"Unsupported IQ quantization format: {quantization_format}") from None
+
 
 # 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,
+        *IQ_FORMATS,
+        QUANTIZATION_NONE,
+        QUANTIZATION_FP8_PB_REAL,
+    }
+)
 
 KV_CACHE_FP8 = "FP8"
 KV_CACHE_FP8_K_NVFP4_V = "FP8_K_NVFP4_V"