From f7f460821c0bbc160bc44edf117f8582a8e39777 Mon Sep 17 00:00:00 2001
From: "Chenhan D. Yu" <5185878+ChenhanYu@users.noreply.github.com>
Date: Tue, 15 Sep 2026 16:51:48 -0700
Subject: [PATCH 01/32] Add IQ unified checkpoint export

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 docs/source/deployment/3_unified_hf.rst       | 30 +++++++
 modelopt/torch/export/convert_hf_config.py    | 17 ++++
 modelopt/torch/export/moe_utils.py            |  5 +-
 modelopt/torch/export/quant_format.py         | 12 ++-
 modelopt/torch/export/quant_utils.py          | 16 ++++
 modelopt/torch/export/unified_export_hf.py    | 18 ++++
 .../torch/export/unified_export_megatron.py   | 90 ++++++++++++++++---
 .../export/test_unified_export_megatron.py    | 72 +++++++++++++++
 tests/unit/torch/export/test_export_weight.py | 24 +++++
 .../torch/export/test_get_quantization.py     | 41 +++++++++
 .../plugins/test_fused_experts.py             | 31 +++++++
 11 files changed, 343 insertions(+), 13 deletions(-)

diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst
index ccef639d00e..2923c227a79 100644
--- a/docs/source/deployment/3_unified_hf.rst
+++ b/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
 --------------------------
diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py
index 251a37cd076..29a9822c5db 100644
--- a/modelopt/torch/export/convert_hf_config.py
+++ b/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
diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py
index 734302690f1..e8ae02090d2 100644
--- a/modelopt/torch/export/moe_utils.py
+++ b/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))
diff --git a/modelopt/torch/export/quant_format.py b/modelopt/torch/export/quant_format.py
index f241880f26f..d1bb213b13d 100644
--- a/modelopt/torch/export/quant_format.py
+++ b/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"
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 35a9cd29d29..4d8a0e47f19 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -62,6 +62,8 @@
     QUANTIZATION_INT4_AWQ,
     QUANTIZATION_INT8_SQ,
     QUANTIZATION_INT8_WO,
+    QUANTIZATION_IQ1_S,
+    QUANTIZATION_IQ2_XS,
     QUANTIZATION_MXFP4,
     QUANTIZATION_MXFP8,
     QUANTIZATION_NONE,
@@ -474,6 +476,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 != "ggml":
+                raise ValueError("IQ formats require the built-in 'ggml' 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"
@@ -722,6 +729,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}
 
@@ -1152,6 +1167,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",
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 8670af77403..8bdd125421f 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/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,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
     )
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 37c896fd609..87eadeb12d1 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -35,6 +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.nn.modules.tensor_quantizer import GroupedQuantizer
 from modelopt.torch.utils import import_plugin, warn_rank_0
 
@@ -61,6 +62,8 @@
     QUANTIZATION_FP8,
     QUANTIZATION_FP8_PB_REAL,
     QUANTIZATION_FP8_PB_WO,
+    QUANTIZATION_IQ1_S,
+    QUANTIZATION_IQ2_XS,
     QUANTIZATION_NONE,
     QUANTIZATION_NVFP4,
     QUANTIZATION_W4A16_NVFP4,
@@ -94,6 +97,7 @@
         get_pipeline_model_parallel_rank,
         get_pipeline_model_parallel_world_size,
         get_tensor_model_parallel_rank,
+        get_tensor_model_parallel_world_size,
     )
     from megatron.core.ssm.mamba_layer import MambaLayer
     from megatron.core.transformer.identity_op import IdentityOp
@@ -309,10 +313,19 @@ def save_pretrained(
         is_last_stage_main_rank = pp_rank == pp_size - 1 and tp_rank == 0
         is_writer_rank = self._is_sidecar_writer_rank(is_last_stage_main_rank)
 
+        quantization_format = self._get_quantization_format(self.model)
+        if (
+            quantization_format in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS)
+            and get_tensor_model_parallel_world_size() != 1
+        ):
+            raise NotImplementedError(
+                "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model "
+                "parallel size 1"
+            )
+
         # Main export process
         layer_state_dicts = self.layer_state_dicts
 
-        quantization_format = self._get_quantization_format(self.model)
         quantization = None
         if quantization_format in (
             QUANTIZATION_FP8_PB_REAL,
@@ -325,6 +338,8 @@ def save_pretrained(
             quantization = "NVFP4"
         elif quantization_format == QUANTIZATION_W4A16_NVFP4:
             quantization = "W4A16_NVFP4"
+        elif quantization_format in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS):
+            quantization = quantization_format.upper()
 
         if is_last_stage_main_rank:
             if is_writer_rank:
@@ -1027,6 +1042,7 @@ def _get_weight_bias(
         module: torch.nn.Module,
         dtype: torch.dtype = torch.float16,
         name_to_value: dict[str, torch.Tensor] | None = None,
+        keep_weight_device: bool = False,
     ) -> dict[str, torch.Tensor]:
         """Get the weight and bias of the module.
 
@@ -1035,6 +1051,7 @@ def _get_weight_bias(
             dtype: The data type of the weight and bias.
             name_to_value: The dictionary to store the weight and bias. A new dict is created
                 if not provided.
+            keep_weight_device: Keep the weight on its current device instead of moving it to CPU.
 
         Returns:
             The dictionary containing the weight and bias.
@@ -1045,7 +1062,9 @@ def _get_weight_bias(
         # layers whose weight is a placeholder) so callers can use "weight" in name_to_value
         # as a reliable guard without re-inspecting module.weight.
         if hasattr(module, "weight") and module.weight is not None and module.weight.numel() > 0:
-            weight = module.weight.to(dtype).cpu()
+            weight = module.weight.to(dtype)
+            if not keep_weight_device:
+                weight = weight.cpu()
             name_to_value["weight"] = weight
 
         if hasattr(module, "bias") and module.bias is not None and module.bias.numel() > 0:
@@ -1082,13 +1101,21 @@ def _get_quantized_state(
             self._record_excluded_module(prefix)
         block_size = get_weight_block_size(module)
 
-        name_to_value = self._get_weight_bias(module, dtype, name_to_value)
+        is_iq = qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS)
+        name_to_value = self._get_weight_bias(
+            module, dtype, name_to_value, keep_weight_device=is_iq
+        )
 
         if "weight" not in name_to_value:
             return name_to_value, qformat, block_size
 
         if qformat == QUANTIZATION_NONE:
             return name_to_value, qformat, block_size
+        # IQ formats derive all block metadata directly from the weight and do not use amax or
+        # separately exported scaling tensors. Keep the weight on-device until its final HF layout
+        # has been produced, so the CUDA packer can be used.
+        if is_iq:
+            return name_to_value, qformat, block_size
         # Getting the weight scales
         weight_scale = get_weight_scaling_factor(module)
         weight_scale_2 = get_weight_scaling_factor_2(module)
@@ -1124,6 +1151,15 @@ def _get_weight_scales(self, quantized_state: dict[str, Any], qformat: str):
 
         return weight_scale, weight_scale_2
 
+    @staticmethod
+    def _get_iq_weight_state(
+        weight_key: str, weight: torch.Tensor, qformat: str
+    ) -> dict[str, torch.Tensor]:
+        """Pack one final-layout weight into the IQ unified-checkpoint representation."""
+        quantize_iq = quantize_iq1_s if qformat == QUANTIZATION_IQ1_S else quantize_iq2_xs
+        packed_weight, _ = quantize_iq(weight)
+        return {weight_key: packed_weight.detach().cpu()}
+
     def _record_layer_quant_config(self, prefix: str, qformat: str | None, block_size: int | None):
         """Record per-HF-layer quantization metadata for mixed precision exports."""
         if qformat in (None, QUANTIZATION_NONE):
@@ -1188,7 +1224,9 @@ def _name_remapping(
             weight = weight + 1.0
         weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat)
 
-        if weight_scale is None:
+        if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS):
+            self._state_dict.update(self._get_iq_weight_state(prefix + "weight", weight, qformat))
+        elif weight_scale is None:
             self._state_dict[prefix + "weight"] = weight
         else:
             self._state_dict[prefix + "weight"] = to_quantized_weight(
@@ -1233,7 +1271,14 @@ def _gated_mlp_slicing(
         gate_proj_weight = weight[:ffn_hidden_size, :]
         up_proj_weight = weight[ffn_hidden_size:, :]
 
-        if weight_scale is None:
+        if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS):
+            self._state_dict.update(
+                self._get_iq_weight_state(gate_proj_prefix + "weight", gate_proj_weight, qformat)
+            )
+            self._state_dict.update(
+                self._get_iq_weight_state(up_proj_prefix + "weight", up_proj_weight, qformat)
+            )
+        elif weight_scale is None:
             self._state_dict[gate_proj_prefix + "weight"] = gate_proj_weight
             self._state_dict[up_proj_prefix + "weight"] = up_proj_weight
         else:
@@ -1399,7 +1444,9 @@ def _grouped_mlp_slicing(
                 name_to_value.pop("weight", None)
                 seen_qformat, seen_block_size = qformat, block_size
 
-                weight = state_dict[weight_key].to(self.dtype).cpu()
+                weight = state_dict[weight_key].to(self.dtype)
+                if qformat not in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS):
+                    weight = weight.cpu()
                 weight_scale_cpu = (
                     weight_scale.detach().cpu().clone() if weight_scale is not None else None
                 )
@@ -1430,7 +1477,13 @@ def _grouped_mlp_slicing(
                     ]
 
                 for shard_prefix, shard_weight, shard_scale in shards:
-                    if shard_scale is None:
+                    if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS):
+                        local_expert_state.update(
+                            self._get_iq_weight_state(
+                                shard_prefix + "weight", shard_weight, qformat
+                            )
+                        )
+                    elif shard_scale is None:
                         local_expert_state[shard_prefix + "weight"] = shard_weight
                     else:
                         local_expert_state[shard_prefix + "weight"] = to_quantized_weight(
@@ -1593,7 +1646,10 @@ 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 weight_scale is None:
+        if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS):
+            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:
             for key, weight in zip(proj_keys, proj_weights):
                 self._state_dict[key] = weight
         else:
@@ -1708,7 +1764,15 @@ 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 weight_scale is None:
+        if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS):
+            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()
+                else:
+                    self._state_dict.update(
+                        self._get_iq_weight_state(proj_prefix + "weight", proj_weight, qformat)
+                    )
+        elif weight_scale is None:
             for key, proj_weight in zip(proj_keys, proj_weights):
                 self._state_dict[key] = proj_weight
         else:
@@ -1842,7 +1906,9 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False, tr
             merged_input_scale = None
 
         # Save the merged weights
-        if merged_weight_scale is None:
+        if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS):
+            self._state_dict.update(self._get_iq_weight_state(prefix, merged_weight, qformat))
+        elif merged_weight_scale is None:
             self._state_dict[prefix] = merged_weight
         else:
             self._state_dict[prefix] = to_quantized_weight(
@@ -1953,7 +2019,9 @@ def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None, is_mtp=F
             merged_input_scale = None
 
         # Save the merged weights
-        if merged_weight_scale is None:
+        if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS):
+            self._state_dict.update(self._get_iq_weight_state(prefix, merged_weight, qformat))
+        elif merged_weight_scale is None:
             # TODO: May need to modify the key name later.
             self._state_dict[prefix] = merged_weight
         else:
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 2da25e90e15..735a7a22de8 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -46,6 +46,9 @@
 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.unified_export_megatron import GPTModelExporter
+from modelopt.torch.quantization.config import QuantizerAttributeConfig
+from modelopt.torch.quantization.ggml import dequantize_iq1_s, dequantize_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
 from modelopt.torch.speculative.plugins.megatron_medusa import _DynamicMedusaGPTModel
@@ -86,6 +89,75 @@ 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):
+    """Megatron export writes the same scale-free IQ representation as HF export."""
+    linear = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16)
+    linear.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits=qformat,
+            block_sizes={-1: 256},
+            backend="ggml",
+            backend_extra_args={"search_impl": "auto"},
+        )
+    )
+    exporter = object.__new__(GPTModelExporter)
+    exporter.dtype = torch.bfloat16
+    exporter._state_dict = {}
+    exporter.exclude_modules = []
+    exporter.layer_config_dict = {}
+
+    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,
+    )
+    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,
+    }
+
+
+def test_megatron_iq_export_rejects_tensor_parallelism():
+    """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",
+            block_sizes={-1: 256},
+            backend="ggml",
+            backend_extra_args={"search_impl": "auto"},
+        )
+    )
+    exporter = object.__new__(GPTModelExporter)
+    exporter.model = torch.nn.Sequential(linear)
+
+    with (
+        patch.object(exporter, "_is_sidecar_writer_rank", return_value=False),
+        patch.object(uem, "get_pipeline_model_parallel_rank", return_value=0),
+        patch.object(uem, "get_pipeline_model_parallel_world_size", return_value=1),
+        patch.object(uem, "get_tensor_model_parallel_rank", return_value=0),
+        patch.object(uem, "get_tensor_model_parallel_world_size", return_value=2),
+        pytest.raises(NotImplementedError, match="tensor model parallel size 1"),
+    ):
+        exporter.save_pretrained("unused", "unused")
+
+
 def _test_unified_export_megatron(
     tmp_path,
     model_type,
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 6fc17d982e8..f15377dc7a4 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -20,10 +20,13 @@
 from _test_utils.torch.export.utils import ToyModel, partial_fp8_config, partial_w4a8_config
 
 import modelopt.torch.quantization as mtq
+from modelopt.torch.export.quant_utils import postprocess_state_dict
 from modelopt.torch.export.unified_export_hf import (
     _export_quantized_weight,
     _process_quantized_modules,
 )
+from modelopt.torch.quantization.config import QuantizerAttributeConfig
+from modelopt.torch.quantization.nn import TensorQuantizer
 from modelopt.torch.quantization.utils import quantizer_attr_names
 
 
@@ -102,6 +105,27 @@ def test_export_per_block_quantized_weight():
     assert not hasattr(model.linears[2], quantizer_attrs.output_scale)
 
 
+@pytest.mark.parametrize(("num_bits", "payload_bytes"), [("iq1_s", 50), ("iq2_xs", 74)])
+def test_export_iq_payload_as_weight(num_bits, payload_bytes):
+    linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)
+    linear.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits=num_bits,
+            block_sizes={-1: 256},
+            backend="ggml",
+            backend_extra_args={"search_impl": "auto"},
+        )
+    )
+
+    _export_quantized_weight(linear, torch.bfloat16)
+    state_dict = postprocess_state_dict(linear.state_dict(), maxbound=448, quantization=None)
+
+    assert state_dict["weight"].shape == (4, 1, payload_bytes)
+    assert state_dict["weight"].dtype == torch.uint8
+    assert "packed_weights" not in state_dict
+    assert "weight_shape" not in state_dict
+
+
 class QuantMoELinear(nn.Module):
     def __init__(self):
         super().__init__()
diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py
index 8b9670e5576..15f9657dffa 100644
--- a/tests/unit/torch/export/test_get_quantization.py
+++ b/tests/unit/torch/export/test_get_quantization.py
@@ -33,6 +33,8 @@
     KV_CACHE_FP8_K_NVFP4_V,
     KV_CACHE_NVFP4,
     QUANTIZATION_FP8,
+    QUANTIZATION_IQ1_S,
+    QUANTIZATION_IQ2_XS,
     QUANTIZATION_NVFP4,
     QUANTIZATION_W4A8_AWQ,
 )
@@ -53,6 +55,45 @@ def __init__(self):
         self.v_bmm_quantizer = TensorQuantizer()
 
 
+@pytest.mark.parametrize(
+    ("num_bits", "quantization_format", "payload_bytes", "effective_bits"),
+    [
+        ("iq1_s", QUANTIZATION_IQ1_S, 50, 1.5625),
+        ("iq2_xs", QUANTIZATION_IQ2_XS, 74, 2.3125),
+    ],
+)
+def test_iq_quantization_config(num_bits, quantization_format, payload_bytes, effective_bits):
+    model = torch.nn.Sequential(torch.nn.Linear(256, 256, bias=False))
+    mtq.quantize(
+        model,
+        {
+            "quant_cfg": [
+                {"quantizer_name": "*", "enable": False},
+                {
+                    "quantizer_name": "*weight_quantizer",
+                    "cfg": {
+                        "num_bits": num_bits,
+                        "block_sizes": {-1: 256},
+                        "backend": "ggml",
+                        "backend_extra_args": {"search_impl": "auto"},
+                    },
+                },
+            ],
+            "algorithm": None,
+        },
+    )
+
+    assert get_quantization_format(model) == quantization_format
+    config = get_quant_config(model)
+    assert config["quantization"]["quant_algo"] == num_bits.upper()
+    assert config["quantization"]["block_payload_bytes"] == payload_bytes
+    hf_config = convert_hf_quant_config_format(config)
+    weights = hf_config["config_groups"]["group_0"]["weights"]
+    assert weights["group_size"] == 256
+    assert weights["effective_bits"] == effective_bits
+    assert weights["packing"] == "ggml"
+
+
 class _FakeKVCacheQuantizer(torch.nn.Module):
     """Minimal FP8 KV cache quantizer for scaling-factor tests."""
 
diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 49cd999a205..f6fe438a2d3 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -502,6 +502,37 @@ def forward_loop(m):
 
         self._cleanup_registry(expert_type)
 
+    def test_export_registers_packed_weight_buffers(self, monkeypatch):
+        """Packed expert weights must remain present in the exported state dict."""
+        model = _TinyMoEModel()
+        expert_type = type(model.moe.experts)
+        self._cleanup_registry(expert_type)
+        register_fused_experts_on_the_fly(model)
+
+        try:
+            converted = QuantModuleRegistry.convert(model.moe.experts)
+
+            def _pack_weight_as_buffer(wrapper, dtype):
+                packed = torch.zeros((*wrapper.weight.shape, 1), dtype=torch.uint8)
+                del wrapper.weight
+                wrapper.register_buffer("weight", packed)
+
+            monkeypatch.setattr(
+                "modelopt.torch.export.unified_export_hf._export_quantized_weight",
+                _pack_weight_as_buffer,
+            )
+
+            _export_fused_experts(converted, torch.float16)
+
+            state_dict = converted.state_dict()
+            for idx in range(NUM_EXPERTS):
+                for projection in ("gate_proj", "up_proj", "down_proj"):
+                    key = f"{idx}.{projection}.weight"
+                    assert key in state_dict
+                    assert state_dict[key].dtype == torch.uint8
+        finally:
+            self._cleanup_registry(expert_type)
+
     def test_uncalibrated_expert_gate_up_share_amax(self, monkeypatch):
         """gate_proj and up_proj must share weight_scale_2 even when an expert
         was never routed during calibration.

From 6f6372ca686477dec027513bf5a3947eaf027e07 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Tue, 15 Sep 2026 18:21:20 -0700
Subject: [PATCH 02/32] Address IQ export review feedback

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/convert_hf_config.py    | 17 +++---
 modelopt/torch/export/moe_utils.py            | 10 +++-
 modelopt/torch/export/quant_format.py         | 43 ++++++++++++++-
 modelopt/torch/export/quant_utils.py          | 52 ++++++++++++++++---
 modelopt/torch/export/unified_export_hf.py    | 12 ++---
 .../torch/export/unified_export_megatron.py   | 29 +++++++----
 .../gpu/torch/export/test_layerwise_export.py | 52 +++++++++++++++++++
 .../export/test_unified_export_megatron.py    | 11 ++++
 tests/unit/torch/export/test_export_weight.py | 16 ++++++
 .../torch/export/test_get_quantization.py     | 17 +++---
 .../plugins/test_fused_experts.py             | 37 +++++++++++++
 11 files changed, 252 insertions(+), 44 deletions(-)

diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py
index 29a9822c5db..cebb2cfe205 100644
--- a/modelopt/torch/export/convert_hf_config.py
+++ b/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,17 +119,14 @@ 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)
+    elif quant_algo.lower() in IQ_FORMATS:
+        spec = iq_format_spec(quant_algo)
         return {
             "weights": {
                 "dynamic": False,
-                "num_bits": 1 if quant_algo == "IQ1_S" else 2,
-                "effective_bits": effective_bits,
+                "num_bits": spec["num_bits"],
                 "type": "int",
-                "group_size": 256,
-                "packing": "ggml",
-                "block_payload_bytes": payload_bytes,
+                "group_size": spec["group_size"],
             }
         }
     else:
@@ -222,8 +221,8 @@ 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)
+    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":
diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py
index e8ae02090d2..f46c0f55dce 100644
--- a/modelopt/torch/export/moe_utils.py
+++ b/modelopt/torch/export/moe_utils.py
@@ -22,6 +22,8 @@
 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.
@@ -114,6 +116,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,6 +159,7 @@ 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.
@@ -166,7 +170,8 @@ def _export_fused_experts(
             # 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 +199,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")
diff --git a/modelopt/torch/export/quant_format.py b/modelopt/torch/export/quant_format.py
index d1bb213b13d..01067a2d59a 100644
--- a/modelopt/torch/export/quant_format.py
+++ b/modelopt/torch/export/quant_format.py
@@ -19,6 +19,17 @@
 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,
+    IQ1_S_EFFECTIVE_BITS,
+    IQ2_XS_BLOCK_BYTES,
+    IQ2_XS_BLOCK_SIZE,
+    IQ2_XS_EFFECTIVE_BITS,
+)
+
 QUANTIZATION_NONE = None
 QUANTIZATION_FP8 = "fp8"
 QUANTIZATION_INT8_SQ = "int8_sq"
@@ -39,14 +50,42 @@
 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,
+        "effective_bits": IQ1_S_EFFECTIVE_BITS,
+        "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,
+        "effective_bits": IQ2_XS_EFFECTIVE_BITS,
+        "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_IQ1_S,
-        QUANTIZATION_IQ2_XS,
+        *IQ_FORMATS,
         QUANTIZATION_NONE,
         QUANTIZATION_FP8_PB_REAL,
     }
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 4d8a0e47f19..51568c5fe61 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -26,6 +26,7 @@
 
 from modelopt import __version__
 from modelopt.torch.models import get_spec, list_all_possible
+from modelopt.torch.quantization.ggml import quantize_iq1_s, quantize_iq2_xs
 from modelopt.torch.quantization.model_calib import (
     enable_stats_collection,
     finish_stats_collection,
@@ -50,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_FORMATS,
     KV_CACHE_FP8,
     KV_CACHE_FP8_K_NVFP4_V,
     KV_CACHE_INT8,
@@ -74,11 +76,49 @@
     QUANTIZATION_W4A8_MXFP4_FP8,
     QUANTIZATION_W4A8_NVFP4_FP8,
     QUANTIZATION_W4A16_NVFP4,
+    iq_format_spec,
 )
 
 logger = logging.getLogger(__name__)
 
 
+def _validate_iq_quantizer_config(
+    module: nn.Module, quantization_format: str, weight_name: str = "weight"
+) -> None:
+    """Reject IQ export settings that the checkpoint encoder cannot reproduce."""
+    quantizer = representative_weight_quantizer(module, weight_name)
+    if not isinstance(quantizer, TensorQuantizer):
+        raise ValueError(f"{quantization_format.upper()} export requires one TensorQuantizer")
+    if quantizer.num_bits != quantization_format or quantizer.backend != "ggml":
+        raise ValueError(
+            f"{quantization_format.upper()} export requires the matching ggml quantizer"
+        )
+    extra_args = quantizer.backend_extra_args or {}
+    search_impl = extra_args.get("search_impl", extra_args.get("iq_search_impl", "auto"))
+    if search_impl != "auto":
+        raise NotImplementedError(
+            f"{quantization_format.upper()} export only supports search_impl='auto'"
+        )
+
+
+def _pack_iq_weight(
+    weight: torch.Tensor,
+    quantization_format: str,
+    module: nn.Module | None = None,
+    weight_name: str = "weight",
+) -> torch.Tensor:
+    """Pack one IQ weight after checking any attached quantizer configuration."""
+    if module is not None:
+        _validate_iq_quantizer_config(module, quantization_format, weight_name)
+    if quantization_format == QUANTIZATION_IQ1_S:
+        packed_weight, _ = quantize_iq1_s(weight)
+    elif quantization_format == QUANTIZATION_IQ2_XS:
+        packed_weight, _ = quantize_iq2_xs(weight)
+    else:
+        raise ValueError(f"Unsupported IQ quantization format: {quantization_format}")
+    return packed_weight
+
+
 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))
@@ -729,13 +769,11 @@ 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
+        elif v in IQ_FORMATS:
+            spec = iq_format_spec(v)
             layer_config = {
-                "quant_algo": v.upper(),
-                "group_size": 256,
-                "block_payload_bytes": payload_bytes,
-                "packing": "ggml",
+                key: spec[key]
+                for key in ("quant_algo", "group_size", "block_payload_bytes", "packing")
             }
         else:
             layer_config = {"quant_algo": v}
@@ -958,7 +996,6 @@ def from_quantized_weight(
     "_amax",
     "_bias_value",
     "input_quantizer._pre_quant_scale",
-    "weight_shape",
 )
 
 
@@ -1167,7 +1204,6 @@ 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",
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 8bdd125421f..a8b25fac73b 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -62,7 +62,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
 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
@@ -95,11 +94,10 @@
 )
 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_MXFP8,
     QUANTIZATION_NONE,
     QUANTIZATION_NVFP4,
@@ -111,6 +109,7 @@
 )
 from .quant_utils import (
     _get_kv_cache_postprocess_config,
+    _pack_iq_weight,
     fuse_prequant_layernorm,
     fuse_prequant_to_linear,
     get_activation_scaling_factor,
@@ -624,16 +623,15 @@ 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
+        packed_weight = _pack_iq_weight(
+            weight.to(dtype), quantization_format, sub_module, weight_name
         )
-        packed_weight, _ = quantize_iq(weight.to(dtype))
         delattr(sub_module, weight_name)
         sub_module.register_buffer("weight", packed_weight)
         maybe_clear_cuda_cache()
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 87eadeb12d1..9d214884653 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
 from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer
 from modelopt.torch.utils import import_plugin, warn_rank_0
 
@@ -57,6 +56,7 @@
 )
 from .plugins.megatron_importer import GPTModelImporter, _get_mamba_conv1d
 from .quant_format import (
+    IQ_FORMATS,
     KV_CACHE_FP8,
     KV_CACHE_NVFP4,
     QUANTIZATION_FP8,
@@ -67,8 +67,11 @@
     QUANTIZATION_NONE,
     QUANTIZATION_NVFP4,
     QUANTIZATION_W4A16_NVFP4,
+    iq_format_spec,
 )
 from .quant_utils import (
+    _pack_iq_weight,
+    _validate_iq_quantizer_config,
     get_activation_scaling_factor,
     get_kv_cache_dtype,
     get_kv_cache_scaling_factor,
@@ -314,10 +317,7 @@ def save_pretrained(
         is_writer_rank = self._is_sidecar_writer_rank(is_last_stage_main_rank)
 
         quantization_format = self._get_quantization_format(self.model)
-        if (
-            quantization_format in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS)
-            and get_tensor_model_parallel_world_size() != 1
-        ):
+        if quantization_format in IQ_FORMATS and get_tensor_model_parallel_world_size() != 1:
             raise NotImplementedError(
                 "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model "
                 "parallel size 1"
@@ -338,7 +338,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:
@@ -403,6 +403,11 @@ def save_pretrained(
                 }
                 if quantization in ("NVFP4", "W4A16_NVFP4"):  # update block size
                     quantization_config["group_size"] = 16
+                elif quantization_format in IQ_FORMATS:
+                    spec = iq_format_spec(quantization_format)
+                    quantization_config.update(
+                        {key: spec[key] for key in ("group_size", "block_payload_bytes", "packing")}
+                    )
 
             if gathered_kv_cache_dtype is not None:
                 quantization_config["kv_cache_quant_algo"] = gathered_kv_cache_dtype
@@ -1101,7 +1106,9 @@ 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
+        if is_iq:
+            _validate_iq_quantizer_config(module, qformat)
         name_to_value = self._get_weight_bias(
             module, dtype, name_to_value, keep_weight_device=is_iq
         )
@@ -1156,8 +1163,12 @@ def _get_iq_weight_state(
         weight_key: str, weight: torch.Tensor, qformat: str
     ) -> dict[str, torch.Tensor]:
         """Pack one final-layout weight into the IQ unified-checkpoint representation."""
-        quantize_iq = quantize_iq1_s if qformat == QUANTIZATION_IQ1_S else quantize_iq2_xs
-        packed_weight, _ = quantize_iq(weight)
+        if get_tensor_model_parallel_world_size() != 1:
+            raise NotImplementedError(
+                "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model "
+                "parallel size 1"
+            )
+        packed_weight = _pack_iq_weight(weight, qformat)
         return {weight_key: packed_weight.detach().cpu()}
 
     def _record_layer_quant_config(self, prefix: str, qformat: str | None, block_size: int | None):
diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py
index 84244a8fec1..396dda8ce36 100644
--- a/tests/gpu/torch/export/test_layerwise_export.py
+++ b/tests/gpu/torch/export/test_layerwise_export.py
@@ -139,6 +139,24 @@ def _fp8_cfg():
     return copy.deepcopy(mtq.FP8_DEFAULT_CFG)
 
 
+def _iq2_cfg():
+    return {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*weight_quantizer",
+                "cfg": {
+                    "num_bits": "iq2_xs",
+                    "block_sizes": {-1: 256},
+                    "backend": "ggml",
+                    "backend_extra_args": {"search_impl": "auto"},
+                },
+            },
+        ],
+        "algorithm": None,
+    }
+
+
 def _kv_cache_cfg():
     return mtq.update_quant_cfg_with_kv_cache_quant(
         copy.deepcopy(mtq.FP8_DEFAULT_CFG), copy.deepcopy(mtq.FP8_KV_CFG["quant_cfg"])
@@ -180,6 +198,40 @@ def test_layerwise_export_resolves_mixed_kv_cache_formats(tmp_path):
     }
 
 
+def test_iq2_layerwise_export_matches_whole_model_export(tmp_path):
+    """The per-layer path writes the same packed IQ tensors and metadata."""
+
+    def build_model():
+        model = (
+            get_tiny_llama(
+                hidden_size=256,
+                intermediate_size=256,
+                num_hidden_layers=2,
+                vocab_size=256,
+            )
+            .cuda()
+            .eval()
+        )
+        model.config.architectures = ["LlamaForCausalLM"]
+        return model
+
+    baseline_dir = tmp_path / "baseline"
+    baseline = mtq.quantize(build_model(), _iq2_cfg())
+    export_hf_checkpoint(baseline, export_dir=baseline_dir)
+
+    layerwise_dir = tmp_path / "layerwise"
+    layerwise = mtq.quantize(build_model(), _iq2_cfg())
+    exporter = LayerwiseExporter(layerwise, layerwise_dir)
+    layers = list(layerwise.model.layers)
+    exporter.bind(layers)
+    for layer_idx, layer in enumerate(layers):
+        exporter.export_layer(layer_idx, layer)
+    exporter.finalize()
+
+    _assert_same_checkpoint(_load_checkpoint(baseline_dir), _load_checkpoint(layerwise_dir))
+    _assert_same_quant_config(baseline_dir, layerwise_dir)
+
+
 def _nvfp4_cfg():
     """NVFP4 with o_proj left unquantized.
 
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 735a7a22de8..01a7686ca96 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -158,6 +158,17 @@ def test_megatron_iq_export_rejects_tensor_parallelism():
         exporter.save_pretrained("unused", "unused")
 
 
+def test_megatron_iq_packer_rejects_tensor_parallelism():
+    """Internal packing keeps the TP=1 guard for extra-module export paths."""
+    weight = torch.randn(2, 256, dtype=torch.bfloat16)
+
+    with (
+        patch.object(uem, "get_tensor_model_parallel_world_size", return_value=2),
+        pytest.raises(NotImplementedError, match="tensor model parallel size 1"),
+    ):
+        GPTModelExporter._get_iq_weight_state("weight", weight, "iq2_xs")
+
+
 def _test_unified_export_megatron(
     tmp_path,
     model_type,
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index f15377dc7a4..08aed32264e 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -126,6 +126,22 @@ def test_export_iq_payload_as_weight(num_bits, payload_bytes):
     assert "weight_shape" not in state_dict
 
 
+@pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"])
+def test_export_iq_rejects_an_unhandled_search_impl(num_bits):
+    linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)
+    linear.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits=num_bits,
+            block_sizes={-1: 256},
+            backend="ggml",
+            backend_extra_args={"search_impl": "reference"},
+        )
+    )
+
+    with pytest.raises(NotImplementedError, match="only supports search_impl='auto'"):
+        _export_quantized_weight(linear, torch.bfloat16)
+
+
 class QuantMoELinear(nn.Module):
     def __init__(self):
         super().__init__()
diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py
index 15f9657dffa..d6e4727f8f6 100644
--- a/tests/unit/torch/export/test_get_quantization.py
+++ b/tests/unit/torch/export/test_get_quantization.py
@@ -56,13 +56,13 @@ def __init__(self):
 
 
 @pytest.mark.parametrize(
-    ("num_bits", "quantization_format", "payload_bytes", "effective_bits"),
+    ("num_bits", "quantization_format", "payload_bytes", "storage_bits"),
     [
-        ("iq1_s", QUANTIZATION_IQ1_S, 50, 1.5625),
-        ("iq2_xs", QUANTIZATION_IQ2_XS, 74, 2.3125),
+        ("iq1_s", QUANTIZATION_IQ1_S, 50, 1),
+        ("iq2_xs", QUANTIZATION_IQ2_XS, 74, 2),
     ],
 )
-def test_iq_quantization_config(num_bits, quantization_format, payload_bytes, effective_bits):
+def test_iq_quantization_config(num_bits, quantization_format, payload_bytes, storage_bits):
     model = torch.nn.Sequential(torch.nn.Linear(256, 256, bias=False))
     mtq.quantize(
         model,
@@ -89,9 +89,12 @@ def test_iq_quantization_config(num_bits, quantization_format, payload_bytes, ef
     assert config["quantization"]["block_payload_bytes"] == payload_bytes
     hf_config = convert_hf_quant_config_format(config)
     weights = hf_config["config_groups"]["group_0"]["weights"]
-    assert weights["group_size"] == 256
-    assert weights["effective_bits"] == effective_bits
-    assert weights["packing"] == "ggml"
+    assert weights == {
+        "dynamic": False,
+        "num_bits": storage_bits,
+        "type": "int",
+        "group_size": 256,
+    }
 
 
 class _FakeKVCacheQuantizer(torch.nn.Module):
diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index f6fe438a2d3..41a6c398ccd 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -15,6 +15,8 @@
 
 """Tests for _QuantFusedExperts: generic fused MoE quantization and export."""
 
+import warnings
+
 import pytest
 import torch
 import torch.nn as nn
@@ -533,6 +535,41 @@ def _pack_weight_as_buffer(wrapper, dtype):
         finally:
             self._cleanup_registry(expert_type)
 
+    def test_iq_export_does_not_create_or_warn_about_amax(self, monkeypatch):
+        """IQ expert packing is weight-only and must not enter amax fallback paths."""
+        model = _TinyMoEModel()
+        expert_type = type(model.moe.experts)
+        self._cleanup_registry(expert_type)
+        register_fused_experts_on_the_fly(model)
+
+        try:
+            converted = QuantModuleRegistry.convert(model.moe.experts)
+            quantizers = list(converted.gate_up_proj_weight_quantizers) + list(
+                converted.down_proj_weight_quantizers
+            )
+            config = QuantizerAttributeConfig(
+                num_bits="iq2_xs",
+                block_sizes={-1: 256},
+                backend="ggml",
+                backend_extra_args={"search_impl": "auto"},
+            )
+            for quantizer in quantizers:
+                quantizer.set_from_attribute_config(config)
+                assert not hasattr(quantizer, "_amax")
+
+            monkeypatch.setattr(
+                "modelopt.torch.export.unified_export_hf._export_quantized_weight",
+                lambda *_args, **_kwargs: None,
+            )
+            with warnings.catch_warnings(record=True) as caught:
+                warnings.simplefilter("always")
+                _export_fused_experts(converted, torch.float16)
+
+            assert not [warning for warning in caught if "amax" in str(warning.message)]
+            assert all(not hasattr(quantizer, "_amax") for quantizer in quantizers)
+        finally:
+            self._cleanup_registry(expert_type)
+
     def test_uncalibrated_expert_gate_up_share_amax(self, monkeypatch):
         """gate_proj and up_proj must share weight_scale_2 even when an expert
         was never routed during calibration.

From fe58bc1a10d8e8d08a590330b2073453bf722b49 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Tue, 15 Sep 2026 21:29:12 -0700
Subject: [PATCH 03/32] Guard unsupported IQ expert export

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/quant_format.py         |  4 ---
 .../torch/export/unified_export_megatron.py   | 18 +++++++----
 .../export/test_unified_export_megatron.py    | 30 +++++++++++++++++++
 3 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/modelopt/torch/export/quant_format.py b/modelopt/torch/export/quant_format.py
index 01067a2d59a..27ab3e7de50 100644
--- a/modelopt/torch/export/quant_format.py
+++ b/modelopt/torch/export/quant_format.py
@@ -24,10 +24,8 @@
 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,
 )
 
 QUANTIZATION_NONE = None
@@ -55,7 +53,6 @@
     QUANTIZATION_IQ1_S: {
         "quant_algo": "IQ1_S",
         "num_bits": 1,
-        "effective_bits": IQ1_S_EFFECTIVE_BITS,
         "group_size": IQ1_S_BLOCK_SIZE,
         "block_payload_bytes": IQ1_S_BLOCK_BYTES,
         "packing": "ggml",
@@ -63,7 +60,6 @@
     QUANTIZATION_IQ2_XS: {
         "quant_algo": "IQ2_XS",
         "num_bits": 2,
-        "effective_bits": IQ2_XS_EFFECTIVE_BITS,
         "group_size": IQ2_XS_BLOCK_SIZE,
         "block_payload_bytes": IQ2_XS_BLOCK_BYTES,
         "packing": "ggml",
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 9d214884653..f842b0d08a5 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -1879,6 +1879,11 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False, tr
             name_to_value, qformat, block_size = self._get_quantized_state(
                 getattr(expert, layer_type), self.dtype, prefix=prefix
             )
+            if qformat in IQ_FORMATS:
+                raise NotImplementedError(
+                    "Megatron IQ1_S/IQ2_XS export does not yet support packed expert weights "
+                    "because the HF transpose changes the 256-value block axis"
+                )
             weight = name_to_value.pop("weight")
             weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat)
             input_scale = (
@@ -1917,9 +1922,7 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False, tr
             merged_input_scale = None
 
         # Save the merged weights
-        if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS):
-            self._state_dict.update(self._get_iq_weight_state(prefix, merged_weight, qformat))
-        elif merged_weight_scale is None:
+        if merged_weight_scale is None:
             self._state_dict[prefix] = merged_weight
         else:
             self._state_dict[prefix] = to_quantized_weight(
@@ -1950,6 +1953,11 @@ def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None, is_mtp=F
             name_to_value, qformat, block_size = self._get_quantized_state(
                 getattr(expert, layer_type), self.dtype, prefix=prefix
             )
+            if qformat in IQ_FORMATS:
+                raise NotImplementedError(
+                    "Megatron IQ1_S/IQ2_XS export does not yet support packed expert weights "
+                    "because the HF transpose changes the 256-value block axis"
+                )
             weight = name_to_value.pop("weight")
             bias = name_to_value.pop("bias", None)
             weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat)
@@ -2030,9 +2038,7 @@ def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None, is_mtp=F
             merged_input_scale = None
 
         # Save the merged weights
-        if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS):
-            self._state_dict.update(self._get_iq_weight_state(prefix, merged_weight, qformat))
-        elif merged_weight_scale is None:
+        if merged_weight_scale is None:
             # TODO: May need to modify the key name later.
             self._state_dict[prefix] = merged_weight
         else:
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 01a7686ca96..85ba9784e26 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -169,6 +169,36 @@ def test_megatron_iq_packer_rejects_tensor_parallelism():
         GPTModelExporter._get_iq_weight_state("weight", weight, "iq2_xs")
 
 
+@pytest.mark.parametrize("method_name", ["_pack_name_remapping", "_pack_name_remapping_gpt_oss"])
+def test_megatron_iq_export_rejects_packed_experts(method_name):
+    """Packed expert layouts cannot preserve the calibrated IQ block axis."""
+    linear = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16)
+    linear.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+            backend_extra_args={"search_impl": "auto"},
+        )
+    )
+    expert = torch.nn.Module()
+    expert.linear_fc1 = linear
+    exporter = object.__new__(GPTModelExporter)
+    exporter.dtype = torch.bfloat16
+    exporter._state_dict = {}
+    exporter.exclude_modules = []
+    exporter.layer_config_dict = {}
+
+    with pytest.raises(NotImplementedError, match="packed expert weights"):
+        getattr(exporter, method_name)(
+            torch.nn.ModuleList([expert]),
+            "model.layers.0.mlp.experts.weight",
+            layer_type="linear_fc1",
+        )
+
+    assert exporter._state_dict == {}
+
+
 def _test_unified_export_megatron(
     tmp_path,
     model_type,

From 040afe8bb1b4a1da0822edeb497e308158c79ac2 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Tue, 15 Sep 2026 21:44:50 -0700
Subject: [PATCH 04/32] Reject unsupported IQ export configs

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          |  9 ++++++++
 .../export/test_unified_export_megatron.py    |  3 ++-
 tests/unit/torch/export/test_export_weight.py | 22 +++++++++++++++++++
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 51568c5fe61..1796561bc3e 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -99,6 +99,15 @@ def _validate_iq_quantizer_config(
         raise NotImplementedError(
             f"{quantization_format.upper()} export only supports search_impl='auto'"
         )
+    input_quantizer = getattr(module, quantizer_attr_names(weight_name).input_quantizer, None)
+    if input_quantizer is not None and (
+        getattr(input_quantizer, "is_enabled", False)
+        or getattr(input_quantizer, "pre_quant_scale", None) is not None
+    ):
+        raise NotImplementedError(
+            f"{quantization_format.upper()} export is weight-only; input quantization and "
+            "pre_quant_scale are not represented in the checkpoint"
+        )
 
 
 def _pack_iq_weight(
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 85ba9784e26..8fa1d69be8e 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -110,7 +110,8 @@ def test_megatron_name_remapping_exports_iq_payload(qformat, payload_bytes, dequ
     exporter.exclude_modules = []
     exporter.layer_config_dict = {}
 
-    exporter._name_remapping(linear, "model.layers.0.mlp.down_proj.")
+    with patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1):
+        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)
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 08aed32264e..03a406185ff 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -142,6 +142,28 @@ def test_export_iq_rejects_an_unhandled_search_impl(num_bits):
         _export_quantized_weight(linear, torch.bfloat16)
 
 
+@pytest.mark.parametrize("input_mode", ["enabled", "pre_quant_scale"])
+@pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"])
+def test_export_iq_rejects_non_weight_only_config(num_bits, input_mode):
+    linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)
+    linear.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits=num_bits,
+            block_sizes={-1: 256},
+            backend="ggml",
+            backend_extra_args={"search_impl": "auto"},
+        )
+    )
+    linear.input_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(num_bits=8, axis=None, enable=input_mode == "enabled")
+    )
+    if input_mode == "pre_quant_scale":
+        linear.input_quantizer.pre_quant_scale = torch.ones(256)
+
+    with pytest.raises(NotImplementedError, match="export is weight-only"):
+        _export_quantized_weight(linear, torch.bfloat16)
+
+
 class QuantMoELinear(nn.Module):
     def __init__(self):
         super().__init__()

From ebb77510b58e70cd0aec9555ad2cbb9accc79b8c Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Tue, 15 Sep 2026 22:11:32 -0700
Subject: [PATCH 05/32] Complete IQ export validation and docs

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 docs/source/deployment/3_unified_hf.rst       | 15 +++++++++++--
 modelopt/torch/export/quant_utils.py          | 22 ++++++++++---------
 modelopt/torch/export/unified_export_hf.py    |  3 +--
 tests/unit/torch/export/test_export_weight.py | 21 +++++++++++++-----
 4 files changed, 41 insertions(+), 20 deletions(-)

diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst
index 2923c227a79..970aa8fc25e 100644
--- a/docs/source/deployment/3_unified_hf.rst
+++ b/docs/source/deployment/3_unified_hf.rst
@@ -50,8 +50,8 @@ 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
+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,
@@ -69,6 +69,17 @@ for IQ1_S and 74 for IQ2_XS. No separate shape tensor is stored: a loader recove
 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``;
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 1796561bc3e..555dcbb9b39 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -94,20 +94,22 @@ def _validate_iq_quantizer_config(
             f"{quantization_format.upper()} export requires the matching ggml quantizer"
         )
     extra_args = quantizer.backend_extra_args or {}
-    search_impl = extra_args.get("search_impl", extra_args.get("iq_search_impl", "auto"))
+    search_impl = extra_args.get("search_impl", "auto")
     if search_impl != "auto":
         raise NotImplementedError(
             f"{quantization_format.upper()} export only supports search_impl='auto'"
         )
-    input_quantizer = getattr(module, quantizer_attr_names(weight_name).input_quantizer, None)
-    if input_quantizer is not None and (
-        getattr(input_quantizer, "is_enabled", False)
-        or getattr(input_quantizer, "pre_quant_scale", None) is not None
-    ):
-        raise NotImplementedError(
-            f"{quantization_format.upper()} export is weight-only; input quantization and "
-            "pre_quant_scale are not represented in the checkpoint"
-        )
+    quantizer_names = quantizer_attr_names(weight_name)
+    for quantizer_name in (quantizer_names.input_quantizer, quantizer_names.output_quantizer):
+        activation_quantizer = getattr(module, quantizer_name, None)
+        if activation_quantizer is not None and (
+            getattr(activation_quantizer, "is_enabled", False)
+            or getattr(activation_quantizer, "pre_quant_scale", None) is not None
+        ):
+            raise NotImplementedError(
+                f"{quantization_format.upper()} export is weight-only; activation quantization "
+                "and pre_quant_scale are not represented in the checkpoint"
+            )
 
 
 def _pack_iq_weight(
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index a8b25fac73b..d61ac8a60cb 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -632,8 +632,7 @@ def _export_quantized_weight(
         packed_weight = _pack_iq_weight(
             weight.to(dtype), quantization_format, sub_module, weight_name
         )
-        delattr(sub_module, weight_name)
-        sub_module.register_buffer("weight", packed_weight)
+        setattr(sub_module, weight_name, nn.Parameter(packed_weight, requires_grad=False))
         maybe_clear_cuda_cache()
         return
 
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 03a406185ff..7d8a7d5621b 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -122,6 +122,7 @@ def test_export_iq_payload_as_weight(num_bits, payload_bytes):
 
     assert state_dict["weight"].shape == (4, 1, payload_bytes)
     assert state_dict["weight"].dtype == torch.uint8
+    assert isinstance(linear.weight, nn.Parameter)
     assert "packed_weights" not in state_dict
     assert "weight_shape" not in state_dict
 
@@ -142,9 +143,16 @@ def test_export_iq_rejects_an_unhandled_search_impl(num_bits):
         _export_quantized_weight(linear, torch.bfloat16)
 
 
-@pytest.mark.parametrize("input_mode", ["enabled", "pre_quant_scale"])
+@pytest.mark.parametrize(
+    ("quantizer_name", "mode"),
+    [
+        ("input_quantizer", "enabled"),
+        ("input_quantizer", "pre_quant_scale"),
+        ("output_quantizer", "enabled"),
+    ],
+)
 @pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"])
-def test_export_iq_rejects_non_weight_only_config(num_bits, input_mode):
+def test_export_iq_rejects_non_weight_only_config(num_bits, quantizer_name, mode):
     linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)
     linear.weight_quantizer = TensorQuantizer(
         QuantizerAttributeConfig(
@@ -154,11 +162,12 @@ def test_export_iq_rejects_non_weight_only_config(num_bits, input_mode):
             backend_extra_args={"search_impl": "auto"},
         )
     )
-    linear.input_quantizer = TensorQuantizer(
-        QuantizerAttributeConfig(num_bits=8, axis=None, enable=input_mode == "enabled")
+    activation_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(num_bits=8, axis=None, enable=mode == "enabled")
     )
-    if input_mode == "pre_quant_scale":
-        linear.input_quantizer.pre_quant_scale = torch.ones(256)
+    setattr(linear, quantizer_name, activation_quantizer)
+    if mode == "pre_quant_scale":
+        activation_quantizer.pre_quant_scale = torch.ones(256)
 
     with pytest.raises(NotImplementedError, match="export is weight-only"):
         _export_quantized_weight(linear, torch.bfloat16)

From a0470b4ca2fd70d39ec476ff2f8536cf25daef8e Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Tue, 15 Sep 2026 22:39:25 -0700
Subject: [PATCH 06/32] Harden IQ export format dispatch

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/moe_utils.py            |  5 +---
 .../torch/export/unified_export_megatron.py   | 14 +++++-----
 .../plugins/test_fused_experts.py             | 27 +++++++++++++------
 3 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py
index f46c0f55dce..f029fa7fadb 100644
--- a/modelopt/torch/export/moe_utils.py
+++ b/modelopt/torch/export/moe_utils.py
@@ -224,10 +224,7 @@ def _export_fused_experts(
             _export_quantized_weight(wrapper, dtype)
 
             proj = nn.Module()
-            if isinstance(wrapper.weight, nn.Parameter):
-                proj.weight = wrapper.weight
-            else:
-                proj.register_buffer("weight", wrapper.weight)
+            proj.weight = wrapper.weight
             for attr in ("weight_scale", "weight_scale_2", "input_scale"):
                 if hasattr(wrapper, attr):
                     proj.register_buffer(attr, getattr(wrapper, attr))
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index f842b0d08a5..4686a9cc574 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -62,8 +62,6 @@
     QUANTIZATION_FP8,
     QUANTIZATION_FP8_PB_REAL,
     QUANTIZATION_FP8_PB_WO,
-    QUANTIZATION_IQ1_S,
-    QUANTIZATION_IQ2_XS,
     QUANTIZATION_NONE,
     QUANTIZATION_NVFP4,
     QUANTIZATION_W4A16_NVFP4,
@@ -1235,7 +1233,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
@@ -1282,7 +1280,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)
             )
@@ -1456,7 +1454,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
@@ -1488,7 +1486,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
@@ -1657,7 +1655,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:
@@ -1775,7 +1773,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/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 41a6c398ccd..04738134504 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -504,7 +504,7 @@ def forward_loop(m):
 
         self._cleanup_registry(expert_type)
 
-    def test_export_registers_packed_weight_buffers(self, monkeypatch):
+    def test_export_preserves_packed_iq_weights(self, monkeypatch):
         """Packed expert weights must remain present in the exported state dict."""
         model = _TinyMoEModel()
         expert_type = type(model.moe.experts)
@@ -513,15 +513,23 @@ def test_export_registers_packed_weight_buffers(self, monkeypatch):
 
         try:
             converted = QuantModuleRegistry.convert(model.moe.experts)
-
-            def _pack_weight_as_buffer(wrapper, dtype):
-                packed = torch.zeros((*wrapper.weight.shape, 1), dtype=torch.uint8)
-                del wrapper.weight
-                wrapper.register_buffer("weight", packed)
+            config = QuantizerAttributeConfig(
+                num_bits="iq2_xs",
+                block_sizes={-1: 256},
+                backend="ggml",
+                backend_extra_args={"search_impl": "auto"},
+            )
+            quantizers = list(converted.gate_up_proj_weight_quantizers) + list(
+                converted.down_proj_weight_quantizers
+            )
+            for quantizer in quantizers:
+                quantizer.set_from_attribute_config(config)
 
             monkeypatch.setattr(
-                "modelopt.torch.export.unified_export_hf._export_quantized_weight",
-                _pack_weight_as_buffer,
+                "modelopt.torch.export.unified_export_hf._pack_iq_weight",
+                lambda weight, *_args: torch.zeros(
+                    (weight.numel() // 256, 74), dtype=torch.uint8, device=weight.device
+                ),
             )
 
             _export_fused_experts(converted, torch.float16)
@@ -532,6 +540,9 @@ def _pack_weight_as_buffer(wrapper, dtype):
                     key = f"{idx}.{projection}.weight"
                     assert key in state_dict
                     assert state_dict[key].dtype == torch.uint8
+                    assert isinstance(
+                        getattr(getattr(converted, str(idx)), projection).weight, nn.Parameter
+                    )
         finally:
             self._cleanup_registry(expert_type)
 

From 1e01776bf8e4be45a4f464184c58fc27e2efbd28 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Tue, 15 Sep 2026 22:53:38 -0700
Subject: [PATCH 07/32] Skip unused IQ expert amax

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/unified_export_megatron.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 4686a9cc574..8d49e9a8c27 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -1442,7 +1442,11 @@ def _grouped_mlp_slicing(
                     # weight_scale_2 derivation asserts one. Max-calibration weight amax
                     # is exactly max(|W|), so compute it from this expert's weight.
                     _wq = module.weight_quantizer
-                    if getattr(_wq, "_amax", None) is None and getattr(_wq, "is_enabled", False):
+                    if (
+                        getattr(_wq, "num_bits", None) not in IQ_FORMATS
+                        and getattr(_wq, "_amax", None) is None
+                        and getattr(_wq, "is_enabled", False)
+                    ):
                         _wq.amax = module.weight.detach().abs().max().float()
                         temp_amax_wqs.append(_wq)
 

From bc0faf587f6fa760de0ad5cd5f8a0665fac6e4b8 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Tue, 15 Sep 2026 23:23:09 -0700
Subject: [PATCH 08/32] Address final IQ export review feedback

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/convert_hf_config.py            |  2 ++
 modelopt/torch/export/moe_utils.py                    |  8 ++++++--
 modelopt/torch/export/unified_export_megatron.py      |  7 ++++++-
 modelopt_recipes/ptq.md                               | 11 ++++++-----
 tests/unit/torch/export/test_get_quantization.py      |  2 ++
 .../torch/quantization/plugins/test_fused_experts.py  |  5 +++++
 6 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py
index cebb2cfe205..2806c2836b0 100644
--- a/modelopt/torch/export/convert_hf_config.py
+++ b/modelopt/torch/export/convert_hf_config.py
@@ -127,6 +127,8 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None)
                 "num_bits": spec["num_bits"],
                 "type": "int",
                 "group_size": spec["group_size"],
+                "packing": spec["packing"],
+                "block_payload_bytes": spec["block_payload_bytes"],
             }
         }
     else:
diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py
index f029fa7fadb..eeb06954a77 100644
--- a/modelopt/torch/export/moe_utils.py
+++ b/modelopt/torch/export/moe_utils.py
@@ -162,9 +162,13 @@ def _export_fused_experts(
             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
+                copy.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
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 8d49e9a8c27..b5687b3ff65 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -1160,7 +1160,12 @@ def _get_weight_scales(self, quantized_state: dict[str, Any], qformat: str):
     def _get_iq_weight_state(
         weight_key: str, weight: torch.Tensor, qformat: str
     ) -> dict[str, torch.Tensor]:
-        """Pack one final-layout weight into the IQ unified-checkpoint representation."""
+        """Pack one final-layout weight into the IQ unified-checkpoint representation.
+
+        Every caller reaches this helper through :meth:`_get_quantized_state`, which validates
+        the owning module before its final-layout tensor is packed. A new direct caller must run
+        ``_validate_iq_quantizer_config`` before calling this helper.
+        """
         if get_tensor_model_parallel_world_size() != 1:
             raise NotImplementedError(
                 "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model "
diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md
index b2a08308507..6d25ff1424c 100644
--- a/modelopt_recipes/ptq.md
+++ b/modelopt_recipes/ptq.md
@@ -138,11 +138,12 @@ activations and tensor-core math are what deliver the throughput.
   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 all linear
-  layers, with BF16 activations. No calibration data is required. Prefer weights whose final
-  dimension is divisible by 256: the fake-quant path pads other shapes with zeros, which
-  prevents reconstruction-cache reuse and can degrade the fit in the final partial block. These
-  recipes configure simulated weight quantization only; packed checkpoint export is added
-  separately.
+  layers, with BF16 activations. Unified export replaces each logical weight with a shaped
+  `uint8` tensor whose final axis contains 50 or 74 bytes per 256-value block; there is no
+  separate shape tensor. No calibration data is required. Every selected weight's final
+  dimension must be divisible by 256 for packed export. Exclude other layers: the fake-quant
+  path pads them with zeros, which prevents reconstruction-cache reuse and can degrade the fit
+  in the final partial block.
 
 ---
 
diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py
index d6e4727f8f6..382b646a027 100644
--- a/tests/unit/torch/export/test_get_quantization.py
+++ b/tests/unit/torch/export/test_get_quantization.py
@@ -94,6 +94,8 @@ def test_iq_quantization_config(num_bits, quantization_format, payload_bytes, st
         "num_bits": storage_bits,
         "type": "int",
         "group_size": 256,
+        "packing": "ggml",
+        "block_payload_bytes": payload_bytes,
     }
 
 
diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 04738134504..852f025c89b 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -572,6 +572,11 @@ def test_iq_export_does_not_create_or_warn_about_amax(self, monkeypatch):
                 "modelopt.torch.export.unified_export_hf._export_quantized_weight",
                 lambda *_args, **_kwargs: None,
             )
+
+            def reject_copy(_quantizer):
+                raise AssertionError("IQ export must reuse its source quantizer")
+
+            monkeypatch.setattr("modelopt.torch.export.moe_utils.copy.deepcopy", reject_copy)
             with warnings.catch_warnings(record=True) as caught:
                 warnings.simplefilter("always")
                 _export_fused_experts(converted, torch.float16)

From f0b42054127b28d4a02c42edd30da48d297ffa92 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Tue, 15 Sep 2026 23:31:38 -0700
Subject: [PATCH 09/32] Scope IQ export documentation to export PR

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 CHANGELOG.rst                            | 2 +-
 modelopt_recipes/general/ptq/iq1_s.yaml  | 5 +++--
 modelopt_recipes/general/ptq/iq2_xs.yaml | 5 +++--
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 62e6d19ddec..780e2812dc7 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -8,7 +8,7 @@ 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 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 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 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.
diff --git a/modelopt_recipes/general/ptq/iq1_s.yaml b/modelopt_recipes/general/ptq/iq1_s.yaml
index 0eb80fcf29d..9a402494b1b 100644
--- a/modelopt_recipes/general/ptq/iq1_s.yaml
+++ b/modelopt_recipes/general/ptq/iq1_s.yaml
@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# IQ1_S weight-only PTQ.
+# IQ1_S weight-only PTQ and unified-checkpoint export.
 
 imports:
   preset: configs/ptq/presets/model/iq1_s
@@ -10,6 +10,7 @@ metadata:
   recipe_type: ptq
   description: >-
     Applies uniform GGML-compatible IQ1_S weight-only quantization to eligible linear layers.
-    This is not a mixed per-tensor precision preset. No calibration data is required.
+    Unified export writes the packed block payload to each quantized Linear weight. This is not
+    a mixed per-tensor precision preset. No calibration data is required.
 quantize:
   $import: preset
diff --git a/modelopt_recipes/general/ptq/iq2_xs.yaml b/modelopt_recipes/general/ptq/iq2_xs.yaml
index c0edc2d5351..d5e2e36bad6 100644
--- a/modelopt_recipes/general/ptq/iq2_xs.yaml
+++ b/modelopt_recipes/general/ptq/iq2_xs.yaml
@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# IQ2_XS weight-only PTQ.
+# IQ2_XS weight-only PTQ and unified-checkpoint export.
 
 imports:
   preset: configs/ptq/presets/model/iq2_xs
@@ -10,6 +10,7 @@ metadata:
   recipe_type: ptq
   description: >-
     Applies uniform GGML-compatible IQ2_XS weight-only quantization to eligible linear layers.
-    This is not a mixed per-tensor precision preset. No calibration data is required.
+    Unified export writes the packed block payload to each quantized Linear weight. This is not
+    a mixed per-tensor precision preset. No calibration data is required.
 quantize:
   $import: preset

From c44331393a17d83d41d90e99e68ececc43579f5a Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Tue, 15 Sep 2026 23:47:39 -0700
Subject: [PATCH 10/32] Preserve export metadata compatibility

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          | 20 +++++++++----
 .../torch/export/unified_export_megatron.py   |  2 +-
 tests/unit/torch/export/test_export_weight.py | 30 +++++++++++++++++++
 3 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 555dcbb9b39..82133f622c6 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -121,12 +121,19 @@ def _pack_iq_weight(
     """Pack one IQ weight after checking any attached quantizer configuration."""
     if module is not None:
         _validate_iq_quantizer_config(module, quantization_format, weight_name)
-    if quantization_format == QUANTIZATION_IQ1_S:
-        packed_weight, _ = quantize_iq1_s(weight)
-    elif quantization_format == QUANTIZATION_IQ2_XS:
-        packed_weight, _ = quantize_iq2_xs(weight)
-    else:
-        raise ValueError(f"Unsupported IQ quantization format: {quantization_format}")
+    try:
+        if quantization_format == QUANTIZATION_IQ1_S:
+            packed_weight, _ = quantize_iq1_s(weight)
+        elif quantization_format == QUANTIZATION_IQ2_XS:
+            packed_weight, _ = quantize_iq2_xs(weight)
+        else:
+            raise ValueError(f"Unsupported IQ quantization format: {quantization_format}")
+    except ValueError as exc:
+        owner = type(module).__name__ if module is not None else "state_dict"
+        raise ValueError(
+            f"Failed to pack {quantization_format.upper()} weight "
+            f"'{owner}.{weight_name}' with shape {tuple(weight.shape)}: {exc}"
+        ) from exc
     return packed_weight
 
 
@@ -1007,6 +1014,7 @@ def from_quantized_weight(
     "_amax",
     "_bias_value",
     "input_quantizer._pre_quant_scale",
+    "weight_shape",
 )
 
 
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index b5687b3ff65..d3428df3bb3 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -1171,7 +1171,7 @@ def _get_iq_weight_state(
                 "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model "
                 "parallel size 1"
             )
-        packed_weight = _pack_iq_weight(weight, qformat)
+        packed_weight = _pack_iq_weight(weight, qformat, weight_name=weight_key)
         return {weight_key: packed_weight.detach().cpu()}
 
     def _record_layer_quant_config(self, prefix: str, qformat: str | None, block_size: int | None):
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 7d8a7d5621b..5c1b845ad64 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -127,6 +127,36 @@ def test_export_iq_payload_as_weight(num_bits, payload_bytes):
     assert "weight_shape" not in state_dict
 
 
+def test_postprocess_state_dict_drops_legacy_weight_shape():
+    state_dict = {
+        "layer.weight": torch.ones(2, 2),
+        "layer.weight_shape": torch.tensor([2, 2]),
+    }
+
+    processed = postprocess_state_dict(state_dict, maxbound=448, quantization=None)
+
+    assert set(processed) == {"layer.weight"}
+
+
+@pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"])
+def test_export_iq_divisibility_error_identifies_weight(num_bits):
+    linear = nn.Linear(192, 4, bias=False, dtype=torch.bfloat16)
+    linear.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits=num_bits,
+            block_sizes={-1: 256},
+            backend="ggml",
+            backend_extra_args={"search_impl": "auto"},
+        )
+    )
+
+    with pytest.raises(
+        ValueError,
+        match=rf"Failed to pack {num_bits.upper()} weight 'Linear.weight' with shape \(4, 192\)",
+    ):
+        _export_quantized_weight(linear, torch.bfloat16)
+
+
 @pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"])
 def test_export_iq_rejects_an_unhandled_search_impl(num_bits):
     linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)

From 8a6fb983181af00baad18739b905e8071bcc7ddd Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 00:00:02 -0700
Subject: [PATCH 11/32] Align IQ export tests and release notes

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 CHANGELOG.rst                                               | 2 +-
 modelopt/torch/export/moe_utils.py                          | 4 ++--
 tests/unit/torch/quantization/plugins/test_fused_experts.py | 4 ++--
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 780e2812dc7..6ea6bd06c84 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -9,7 +9,7 @@ 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. 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 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 ``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.
 
diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py
index eeb06954a77..1cabef8a37d 100644
--- a/modelopt/torch/export/moe_utils.py
+++ b/modelopt/torch/export/moe_utils.py
@@ -15,8 +15,8 @@
 
 """Utilities for Mixture-of-Experts (MoE) model export."""
 
-import copy
 import warnings
+from copy import deepcopy
 from pathlib import Path
 
 import torch
@@ -166,7 +166,7 @@ def _export_fused_experts(
             # 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)
+                deepcopy(w_quantizer_src)
                 if uses_first_proj_quantizers and not is_iq
                 else w_quantizer_src
             )
diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 852f025c89b..10289ac71c2 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -573,10 +573,10 @@ def test_iq_export_does_not_create_or_warn_about_amax(self, monkeypatch):
                 lambda *_args, **_kwargs: None,
             )
 
-            def reject_copy(_quantizer):
+            def reject_copy(*_args, **_kwargs):
                 raise AssertionError("IQ export must reuse its source quantizer")
 
-            monkeypatch.setattr("modelopt.torch.export.moe_utils.copy.deepcopy", reject_copy)
+            monkeypatch.setattr("modelopt.torch.export.moe_utils.deepcopy", reject_copy)
             with warnings.catch_warnings(record=True) as caught:
                 warnings.simplefilter("always")
                 _export_fused_experts(converted, torch.float16)

From a1e5c20e714961eb019429dc003fccf41353e8e3 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 00:10:27 -0700
Subject: [PATCH 12/32] Make Megatron IQ failures rank safe

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          |  9 +++++---
 .../torch/export/unified_export_megatron.py   | 21 +++++++++++++++--
 .../export/test_unified_export_megatron.py    | 23 +++++++++++++++++++
 3 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 82133f622c6..5485091158b 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -117,8 +117,9 @@ def _pack_iq_weight(
     quantization_format: str,
     module: nn.Module | None = None,
     weight_name: str = "weight",
+    describe_as: str | None = None,
 ) -> torch.Tensor:
-    """Pack one IQ weight after checking any attached quantizer configuration."""
+    """Pack one IQ weight, using ``describe_as`` only for failure attribution."""
     if module is not None:
         _validate_iq_quantizer_config(module, quantization_format, weight_name)
     try:
@@ -129,10 +130,12 @@ def _pack_iq_weight(
         else:
             raise ValueError(f"Unsupported IQ quantization format: {quantization_format}")
     except ValueError as exc:
-        owner = type(module).__name__ if module is not None else "state_dict"
+        if describe_as is None:
+            owner = type(module).__name__ if module is not None else "state_dict"
+            describe_as = f"{owner}.{weight_name}"
         raise ValueError(
             f"Failed to pack {quantization_format.upper()} weight "
-            f"'{owner}.{weight_name}' with shape {tuple(weight.shape)}: {exc}"
+            f"'{describe_as}' with shape {tuple(weight.shape)}: {exc}"
         ) from exc
     return packed_weight
 
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index d3428df3bb3..740bdcf22d2 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -292,6 +292,21 @@ def _is_sidecar_writer_rank(is_last_stage_main_rank: bool) -> bool:
             and get_expert_model_parallel_rank() == 0
         )
 
+    @staticmethod
+    def _collective_any_iq(quantization_format: str | None) -> bool:
+        """Return whether any distributed rank owns an IQ-quantized module."""
+        local_iq = quantization_format in IQ_FORMATS
+        if not torch.distributed.is_initialized():
+            return local_iq
+        device = (
+            torch.device("cuda", torch.cuda.current_device())
+            if torch.distributed.get_backend() == torch.distributed.Backend.NCCL
+            else torch.device("cpu")
+        )
+        any_iq = torch.tensor(int(local_iq), dtype=torch.int32, device=device)
+        torch.distributed.all_reduce(any_iq, op=torch.distributed.ReduceOp.MAX)
+        return bool(any_iq.item())
+
     def save_pretrained(
         self,
         save_directory: str | os.PathLike,
@@ -315,7 +330,9 @@ def save_pretrained(
         is_writer_rank = self._is_sidecar_writer_rank(is_last_stage_main_rank)
 
         quantization_format = self._get_quantization_format(self.model)
-        if quantization_format in IQ_FORMATS and get_tensor_model_parallel_world_size() != 1:
+        if get_tensor_model_parallel_world_size() != 1 and self._collective_any_iq(
+            quantization_format
+        ):
             raise NotImplementedError(
                 "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model "
                 "parallel size 1"
@@ -1171,7 +1188,7 @@ def _get_iq_weight_state(
                 "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model "
                 "parallel size 1"
             )
-        packed_weight = _pack_iq_weight(weight, qformat, weight_name=weight_key)
+        packed_weight = _pack_iq_weight(weight, qformat, describe_as=weight_key)
         return {weight_key: packed_weight.detach().cpu()}
 
     def _record_layer_quant_config(self, prefix: str, qformat: str | None, block_size: int | None):
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 8fa1d69be8e..3454b368a85 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -159,6 +159,18 @@ def test_megatron_iq_export_rejects_tensor_parallelism():
         exporter.save_pretrained("unused", "unused")
 
 
+def test_megatron_iq_tensor_parallel_guard_is_collective():
+    def report_remote_iq(flag, **_kwargs):
+        flag.fill_(1)
+
+    with (
+        patch.object(torch.distributed, "is_initialized", return_value=True),
+        patch.object(torch.distributed, "get_backend", return_value=torch.distributed.Backend.GLOO),
+        patch.object(torch.distributed, "all_reduce", side_effect=report_remote_iq),
+    ):
+        assert GPTModelExporter._collective_any_iq(None)
+
+
 def test_megatron_iq_packer_rejects_tensor_parallelism():
     """Internal packing keeps the TP=1 guard for extra-module export paths."""
     weight = torch.randn(2, 256, dtype=torch.bfloat16)
@@ -170,6 +182,17 @@ def test_megatron_iq_packer_rejects_tensor_parallelism():
         GPTModelExporter._get_iq_weight_state("weight", weight, "iq2_xs")
 
 
+def test_megatron_iq_packer_error_identifies_state_dict_key():
+    weight_key = "model.layers.0.mlp.down_proj.weight"
+    weight = torch.randn(2, 192, dtype=torch.bfloat16)
+
+    with (
+        patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1),
+        pytest.raises(ValueError, match=weight_key),
+    ):
+        GPTModelExporter._get_iq_weight_state(weight_key, weight, "iq2_xs")
+
+
 @pytest.mark.parametrize("method_name", ["_pack_name_remapping", "_pack_name_remapping_gpt_oss"])
 def test_megatron_iq_export_rejects_packed_experts(method_name):
     """Packed expert layouts cannot preserve the calibrated IQ block axis."""

From 354565f7cc0d43ae4f04316ab16c558c737f0444 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 00:23:03 -0700
Subject: [PATCH 13/32] Assert IQ packed layout invariants

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          | 18 +++++++++++++++---
 tests/unit/torch/export/test_export_weight.py | 11 ++++++++++-
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 5485091158b..bd7b298eb9b 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -122,13 +122,13 @@ def _pack_iq_weight(
     """Pack one IQ weight, using ``describe_as`` only for failure attribution."""
     if module is not None:
         _validate_iq_quantizer_config(module, quantization_format, weight_name)
+    if quantization_format not in IQ_FORMATS:
+        raise ValueError(f"Unsupported IQ quantization format: {quantization_format}")
     try:
         if quantization_format == QUANTIZATION_IQ1_S:
             packed_weight, _ = quantize_iq1_s(weight)
-        elif quantization_format == QUANTIZATION_IQ2_XS:
-            packed_weight, _ = quantize_iq2_xs(weight)
         else:
-            raise ValueError(f"Unsupported IQ quantization format: {quantization_format}")
+            packed_weight, _ = quantize_iq2_xs(weight)
     except ValueError as exc:
         if describe_as is None:
             owner = type(module).__name__ if module is not None else "state_dict"
@@ -137,6 +137,18 @@ def _pack_iq_weight(
             f"Failed to pack {quantization_format.upper()} weight "
             f"'{describe_as}' with shape {tuple(weight.shape)}: {exc}"
         ) from exc
+
+    spec = iq_format_spec(quantization_format)
+    expected_shape = (
+        *weight.shape[:-1],
+        weight.shape[-1] // spec["group_size"],
+        spec["block_payload_bytes"],
+    )
+    if packed_weight.shape != expected_shape:
+        raise RuntimeError(
+            f"{quantization_format.upper()} packing produced {tuple(packed_weight.shape)}, "
+            f"expected {expected_shape}; logical-shape recovery is no longer valid"
+        )
     return packed_weight
 
 
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 5c1b845ad64..4899d5f434f 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -20,7 +20,7 @@
 from _test_utils.torch.export.utils import ToyModel, partial_fp8_config, partial_w4a8_config
 
 import modelopt.torch.quantization as mtq
-from modelopt.torch.export.quant_utils import postprocess_state_dict
+from modelopt.torch.export.quant_utils import _pack_iq_weight, postprocess_state_dict
 from modelopt.torch.export.unified_export_hf import (
     _export_quantized_weight,
     _process_quantized_modules,
@@ -127,6 +127,15 @@ def test_export_iq_payload_as_weight(num_bits, payload_bytes):
     assert "weight_shape" not in state_dict
 
 
+@pytest.mark.parametrize(("num_bits", "payload_bytes"), [("iq1_s", 50), ("iq2_xs", 74)])
+def test_iq_packer_preserves_leading_dimensions(num_bits, payload_bytes):
+    weight = torch.randn(2, 3, 512, dtype=torch.bfloat16)
+
+    packed = _pack_iq_weight(weight, num_bits)
+
+    assert packed.shape == (2, 3, 2, payload_bytes)
+
+
 def test_postprocess_state_dict_drops_legacy_weight_shape():
     state_dict = {
         "layer.weight": torch.ones(2, 2),

From c9e5da1aff36ee0e18deab34ef3c43cee57f338e Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 00:29:38 -0700
Subject: [PATCH 14/32] Enforce final IQ export validation

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 .../torch/export/unified_export_megatron.py   | 31 +++++++------
 .../export/test_unified_export_megatron.py    | 44 +++++++++++++++++--
 2 files changed, 57 insertions(+), 18 deletions(-)

diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 740bdcf22d2..b39c2eeebcf 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -1175,20 +1175,15 @@ def _get_weight_scales(self, quantized_state: dict[str, Any], qformat: str):
 
     @staticmethod
     def _get_iq_weight_state(
-        weight_key: str, weight: torch.Tensor, qformat: str
+        weight_key: str, weight: torch.Tensor, qformat: str, module: torch.nn.Module
     ) -> dict[str, torch.Tensor]:
-        """Pack one final-layout weight into the IQ unified-checkpoint representation.
-
-        Every caller reaches this helper through :meth:`_get_quantized_state`, which validates
-        the owning module before its final-layout tensor is packed. A new direct caller must run
-        ``_validate_iq_quantizer_config`` before calling this helper.
-        """
+        """Validate and pack one final-layout weight into the IQ checkpoint representation."""
         if get_tensor_model_parallel_world_size() != 1:
             raise NotImplementedError(
                 "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model "
                 "parallel size 1"
             )
-        packed_weight = _pack_iq_weight(weight, qformat, describe_as=weight_key)
+        packed_weight = _pack_iq_weight(weight, qformat, module=module, describe_as=weight_key)
         return {weight_key: packed_weight.detach().cpu()}
 
     def _record_layer_quant_config(self, prefix: str, qformat: str | None, block_size: int | None):
@@ -1256,7 +1251,9 @@ def _name_remapping(
         weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat)
 
         if qformat in IQ_FORMATS:
-            self._state_dict.update(self._get_iq_weight_state(prefix + "weight", weight, qformat))
+            self._state_dict.update(
+                self._get_iq_weight_state(prefix + "weight", weight, qformat, module)
+            )
         elif weight_scale is None:
             self._state_dict[prefix + "weight"] = weight
         else:
@@ -1304,10 +1301,14 @@ def _gated_mlp_slicing(
 
         if qformat in IQ_FORMATS:
             self._state_dict.update(
-                self._get_iq_weight_state(gate_proj_prefix + "weight", gate_proj_weight, qformat)
+                self._get_iq_weight_state(
+                    gate_proj_prefix + "weight", gate_proj_weight, qformat, module
+                )
             )
             self._state_dict.update(
-                self._get_iq_weight_state(up_proj_prefix + "weight", up_proj_weight, qformat)
+                self._get_iq_weight_state(
+                    up_proj_prefix + "weight", up_proj_weight, qformat, module
+                )
             )
         elif weight_scale is None:
             self._state_dict[gate_proj_prefix + "weight"] = gate_proj_weight
@@ -1515,7 +1516,7 @@ def _grouped_mlp_slicing(
                     if qformat in IQ_FORMATS:
                         local_expert_state.update(
                             self._get_iq_weight_state(
-                                shard_prefix + "weight", shard_weight, qformat
+                                shard_prefix + "weight", shard_weight, qformat, module
                             )
                         )
                     elif shard_scale is None:
@@ -1683,7 +1684,7 @@ def _take(tensor, index, last_dim, with_gate=False):
 
         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))
+                self._state_dict.update(self._get_iq_weight_state(key, weight, qformat, module))
         elif weight_scale is None:
             for key, weight in zip(proj_keys, proj_weights):
                 self._state_dict[key] = weight
@@ -1805,7 +1806,9 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False):
                     self._state_dict[proj_prefix + "weight"] = proj_weight.cpu()
                 else:
                     self._state_dict.update(
-                        self._get_iq_weight_state(proj_prefix + "weight", proj_weight, qformat)
+                        self._get_iq_weight_state(
+                            proj_prefix + "weight", proj_weight, qformat, in_proj
+                        )
                     )
         elif weight_scale is None:
             for key, proj_weight in zip(proj_keys, proj_weights):
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 3454b368a85..7967a9a5762 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -173,24 +173,60 @@ def report_remote_iq(flag, **_kwargs):
 
 def test_megatron_iq_packer_rejects_tensor_parallelism():
     """Internal packing keeps the TP=1 guard for extra-module export paths."""
-    weight = torch.randn(2, 256, dtype=torch.bfloat16)
+    module = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16)
+    module.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+            backend_extra_args={"search_impl": "auto"},
+        )
+    )
 
     with (
         patch.object(uem, "get_tensor_model_parallel_world_size", return_value=2),
         pytest.raises(NotImplementedError, match="tensor model parallel size 1"),
     ):
-        GPTModelExporter._get_iq_weight_state("weight", weight, "iq2_xs")
+        GPTModelExporter._get_iq_weight_state("weight", module.weight, "iq2_xs", module)
 
 
 def test_megatron_iq_packer_error_identifies_state_dict_key():
     weight_key = "model.layers.0.mlp.down_proj.weight"
-    weight = torch.randn(2, 192, dtype=torch.bfloat16)
+    module = torch.nn.Linear(192, 2, bias=False, dtype=torch.bfloat16)
+    module.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+            backend_extra_args={"search_impl": "auto"},
+        )
+    )
 
     with (
         patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1),
         pytest.raises(ValueError, match=weight_key),
     ):
-        GPTModelExporter._get_iq_weight_state(weight_key, weight, "iq2_xs")
+        GPTModelExporter._get_iq_weight_state(weight_key, module.weight, "iq2_xs", module)
+
+
+def test_megatron_iq_packer_revalidates_final_layout_owner():
+    """The final packing boundary rejects activation state omitted from the checkpoint."""
+    module = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16)
+    module.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+            backend_extra_args={"search_impl": "auto"},
+        )
+    )
+    module.input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8))
+
+    with (
+        patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1),
+        pytest.raises(NotImplementedError, match="weight-only"),
+    ):
+        GPTModelExporter._get_iq_weight_state("weight", module.weight, "iq2_xs", module)
 
 
 @pytest.mark.parametrize("method_name", ["_pack_name_remapping", "_pack_name_remapping_gpt_oss"])

From 9a75a12a2798899f1fdc259af5c12550bd757090 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 00:37:57 -0700
Subject: [PATCH 15/32] Detect mixed IQ export collectively

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 .../torch/export/unified_export_megatron.py   | 20 +++++++++++++++----
 .../export/test_unified_export_megatron.py    | 17 +++++++++++++++-
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index b39c2eeebcf..ad4c21998e6 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -35,7 +35,10 @@
 from safetensors.torch import save_file
 
 from modelopt import __version__
-from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer
+from modelopt.torch.quantization.nn.modules.tensor_quantizer import (
+    GroupedQuantizer,
+    TensorQuantizer,
+)
 from modelopt.torch.utils import import_plugin, warn_rank_0
 
 from .convert_hf_config import convert_hf_quant_config_format
@@ -293,9 +296,18 @@ def _is_sidecar_writer_rank(is_last_stage_main_rank: bool) -> bool:
         )
 
     @staticmethod
-    def _collective_any_iq(quantization_format: str | None) -> bool:
+    def _model_has_iq_quantizer(model: torch.nn.Module) -> bool:
+        """Return whether this rank owns any enabled IQ weight quantizer."""
+        return any(
+            isinstance(module, TensorQuantizer)
+            and module.is_enabled
+            and module.num_bits in IQ_FORMATS
+            for module in model.modules()
+        )
+
+    @staticmethod
+    def _collective_any_iq(local_iq: bool) -> bool:
         """Return whether any distributed rank owns an IQ-quantized module."""
-        local_iq = quantization_format in IQ_FORMATS
         if not torch.distributed.is_initialized():
             return local_iq
         device = (
@@ -331,7 +343,7 @@ def save_pretrained(
 
         quantization_format = self._get_quantization_format(self.model)
         if get_tensor_model_parallel_world_size() != 1 and self._collective_any_iq(
-            quantization_format
+            self._model_has_iq_quantizer(self.model)
         ):
             raise NotImplementedError(
                 "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model "
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 7967a9a5762..a6e99f546a0 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -168,7 +168,22 @@ def report_remote_iq(flag, **_kwargs):
         patch.object(torch.distributed, "get_backend", return_value=torch.distributed.Backend.GLOO),
         patch.object(torch.distributed, "all_reduce", side_effect=report_remote_iq),
     ):
-        assert GPTModelExporter._collective_any_iq(None)
+        assert GPTModelExporter._collective_any_iq(False)
+
+
+def test_megatron_iq_tensor_parallel_guard_finds_later_mixed_format():
+    fp8 = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16)
+    fp8.weight_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3)))
+    iq = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16)
+    iq.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+        )
+    )
+
+    assert GPTModelExporter._model_has_iq_quantizer(torch.nn.Sequential(fp8, iq))
 
 
 def test_megatron_iq_packer_rejects_tensor_parallelism():

From 85e7300c004fc8977d5349c2ca9a1b9c84b088ac Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 00:52:27 -0700
Subject: [PATCH 16/32] Preflight IQ export weight shapes

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/layerwise_export.py     |  2 +
 modelopt/torch/export/quant_utils.py          | 34 ++++++++++++-
 modelopt/torch/export/unified_export_hf.py    |  3 ++
 tests/unit/torch/export/test_export_weight.py | 49 ++++++++++++++++++-
 4 files changed, 85 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py
index 9767d5635a4..ed35db30703 100644
--- a/modelopt/torch/export/layerwise_export.py
+++ b/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_weight_shapes,
     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_weight_shapes(model)
 
     tied = _tied_quantized_modules(model)
     if tied:
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index bd7b298eb9b..7eea868037c 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -82,6 +82,32 @@
 logger = logging.getLogger(__name__)
 
 
+def _validate_iq_export_weight_shapes(model: nn.Module) -> None:
+    """Reject all IQ weights that cannot be packed before export mutates the model."""
+    incompatible: list[str] = []
+    for module_name, module in model.named_modules():
+        for weight_name in weight_attr_names(module):
+            quantizer = representative_weight_quantizer(module, weight_name)
+            if not (
+                isinstance(quantizer, TensorQuantizer)
+                and quantizer.is_enabled
+                and quantizer.num_bits in IQ_FORMATS
+            ):
+                continue
+            weight = getattr(module, weight_name)
+            group_size = iq_format_spec(quantizer.num_bits)["group_size"]
+            if weight.dim() == 0 or weight.shape[-1] % group_size:
+                qualified_name = f"{module_name}.{weight_name}".lstrip(".")
+                incompatible.append(f"{qualified_name}: {tuple(weight.shape)}")
+
+    if incompatible:
+        details = "\n  - ".join(incompatible)
+        raise ValueError(
+            "IQ export requires every selected weight's final dimension to be divisible by "
+            f"its format group size. Incompatible weights:\n  - {details}"
+        )
+
+
 def _validate_iq_quantizer_config(
     module: nn.Module, quantization_format: str, weight_name: str = "weight"
 ) -> None:
@@ -127,8 +153,12 @@ def _pack_iq_weight(
     try:
         if quantization_format == QUANTIZATION_IQ1_S:
             packed_weight, _ = quantize_iq1_s(weight)
-        else:
+        elif quantization_format == QUANTIZATION_IQ2_XS:
             packed_weight, _ = quantize_iq2_xs(weight)
+        else:
+            raise NotImplementedError(
+                f"No checkpoint packer is registered for {quantization_format.upper()}"
+            )
     except ValueError as exc:
         if describe_as is None:
             owner = type(module).__name__ if module is not None else "state_dict"
@@ -549,7 +579,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")
             return weight_quantizer.num_bits
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index d61ac8a60cb..bafcf3ff475 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -110,6 +110,7 @@
 from .quant_utils import (
     _get_kv_cache_postprocess_config,
     _pack_iq_weight,
+    _validate_iq_export_weight_shapes,
     fuse_prequant_layernorm,
     fuse_prequant_to_linear,
     get_activation_scaling_factor,
@@ -1680,6 +1681,8 @@ def export_hf_checkpoint(
         )
         return
 
+    _validate_iq_export_weight_shapes(model)
+
     is_fsdp2_sharded = (
         torch.distributed.is_available()
         and torch.distributed.is_initialized()
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 4899d5f434f..92316b01bec 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -20,10 +20,15 @@
 from _test_utils.torch.export.utils import ToyModel, partial_fp8_config, partial_w4a8_config
 
 import modelopt.torch.quantization as mtq
-from modelopt.torch.export.quant_utils import _pack_iq_weight, postprocess_state_dict
+from modelopt.torch.export.quant_utils import (
+    _pack_iq_weight,
+    _validate_iq_export_weight_shapes,
+    postprocess_state_dict,
+)
 from modelopt.torch.export.unified_export_hf import (
     _export_quantized_weight,
     _process_quantized_modules,
+    export_hf_checkpoint,
 )
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.nn import TensorQuantizer
@@ -166,6 +171,48 @@ def test_export_iq_divisibility_error_identifies_weight(num_bits):
         _export_quantized_weight(linear, torch.bfloat16)
 
 
+def test_export_iq_shape_preflight_reports_every_incompatible_weight():
+    model = nn.Sequential(
+        nn.Linear(192, 4, bias=False, dtype=torch.bfloat16),
+        nn.Linear(256, 4, bias=False, dtype=torch.bfloat16),
+        nn.Linear(320, 4, bias=False, dtype=torch.bfloat16),
+    )
+    for linear in model:
+        linear.weight_quantizer = TensorQuantizer(
+            QuantizerAttributeConfig(
+                num_bits="iq2_xs",
+                block_sizes={-1: 256},
+                backend="ggml",
+            )
+        )
+
+    with pytest.raises(ValueError) as exc_info:
+        _validate_iq_export_weight_shapes(model)
+
+    message = str(exc_info.value)
+    assert "0.weight: (4, 192)" in message
+    assert "2.weight: (4, 320)" in message
+    assert "1.weight" not in message
+
+
+def test_export_hf_checkpoint_runs_iq_shape_preflight_before_mutation(tmp_path):
+    linear = nn.Linear(192, 4, bias=False, dtype=torch.bfloat16)
+    linear.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+        )
+    )
+    original_weight = linear.weight.detach().clone()
+
+    with pytest.raises(ValueError, match=r"weight: \(4, 192\)"):
+        export_hf_checkpoint(linear, export_dir=tmp_path)
+
+    assert linear.weight.dtype == torch.bfloat16
+    assert torch.equal(linear.weight, original_weight)
+
+
 @pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"])
 def test_export_iq_rejects_an_unhandled_search_impl(num_bits):
     linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)

From 15aaeac67d23016feec116acb60dd6d605551990 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 00:56:42 -0700
Subject: [PATCH 17/32] Harden IQ export guards

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 .../export/plugins/vllm_fakequant_megatron.py |  2 ++
 modelopt/torch/export/quant_utils.py          |  4 +--
 .../torch/export/unified_export_megatron.py   | 25 +++++++++++------
 .../export/test_unified_export_megatron.py    | 27 +++++++++++++++++++
 tests/unit/torch/export/test_export_weight.py | 11 ++++++++
 5 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py
index 832342391d0..96355cc30d5 100644
--- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py
+++ b/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)."""
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 7eea868037c..d6bb10244d5 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -159,11 +159,11 @@ def _pack_iq_weight(
             raise NotImplementedError(
                 f"No checkpoint packer is registered for {quantization_format.upper()}"
             )
-    except ValueError as exc:
+    except (TypeError, ValueError) as exc:
         if describe_as is None:
             owner = type(module).__name__ if module is not None else "state_dict"
             describe_as = f"{owner}.{weight_name}"
-        raise ValueError(
+        raise type(exc)(
             f"Failed to pack {quantization_format.upper()} weight "
             f"'{describe_as}' with shape {tuple(weight.shape)}: {exc}"
         ) from exc
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index ad4c21998e6..d0ced4be05a 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -137,6 +137,8 @@ class GPTModelExporter:
         clamp_kv_cache_scales: Whether to clamp FP8 KV cache scaling factors to at least 1.0.
     """
 
+    _packs_iq_weights = True
+
     def __init__(
         self,
         model: torch.nn.Module,
@@ -272,6 +274,7 @@ def save_pretrained_extra_modules(
         save_directory: str | os.PathLike,
     ):
         """Save a EAGLE or Medusa checkpoints which can be deployed by vLLM and TensorRT-LLM."""
+        self._validate_iq_tensor_parallelism()
         # We use the last PP rank to write the config because
         # medusa_heads and eagle_module only exist in the last stage.
         pp_rank = get_pipeline_model_parallel_rank()
@@ -319,6 +322,18 @@ def _collective_any_iq(local_iq: bool) -> bool:
         torch.distributed.all_reduce(any_iq, op=torch.distributed.ReduceOp.MAX)
         return bool(any_iq.item())
 
+    def _validate_iq_tensor_parallelism(self) -> None:
+        """Fail collectively when this exporter would pack IQ weights across TP ranks."""
+        if (
+            self._packs_iq_weights
+            and get_tensor_model_parallel_world_size() != 1
+            and self._collective_any_iq(self._model_has_iq_quantizer(self.model))
+        ):
+            raise NotImplementedError(
+                "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model "
+                "parallel size 1"
+            )
+
     def save_pretrained(
         self,
         save_directory: str | os.PathLike,
@@ -329,6 +344,8 @@ def save_pretrained(
         Args:
             save_directory: Directory to which to save. Will be created if it doesn't exist.
         """
+        self._validate_iq_tensor_parallelism()
+
         pp_rank = get_pipeline_model_parallel_rank()
         pp_size = get_pipeline_model_parallel_world_size()
         tp_rank = get_tensor_model_parallel_rank()
@@ -342,14 +359,6 @@ def save_pretrained(
         is_writer_rank = self._is_sidecar_writer_rank(is_last_stage_main_rank)
 
         quantization_format = self._get_quantization_format(self.model)
-        if get_tensor_model_parallel_world_size() != 1 and self._collective_any_iq(
-            self._model_has_iq_quantizer(self.model)
-        ):
-            raise NotImplementedError(
-                "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model "
-                "parallel size 1"
-            )
-
         # Main export process
         layer_state_dicts = self.layer_state_dicts
 
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 a6e99f546a0..4c7149ff40f 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -45,6 +45,7 @@
 import modelopt.torch.quantization as mtq
 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.plugins.vllm_fakequant_megatron import VllmFqGPTModelExporter
 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
@@ -186,6 +187,32 @@ def test_megatron_iq_tensor_parallel_guard_finds_later_mixed_format():
     assert GPTModelExporter._model_has_iq_quantizer(torch.nn.Sequential(fp8, iq))
 
 
+def test_megatron_iq_tensor_parallel_guard_covers_extra_modules():
+    exporter = object.__new__(GPTModelExporter)
+    exporter.model = torch.nn.Module()
+
+    with (
+        patch.object(uem, "get_tensor_model_parallel_world_size", return_value=2),
+        patch.object(exporter, "_model_has_iq_quantizer", return_value=True),
+        patch.object(exporter, "_collective_any_iq", return_value=True),
+        pytest.raises(NotImplementedError, match="tensor model parallel size 1"),
+    ):
+        exporter.save_pretrained_extra_modules("unused")
+
+
+def test_vllm_fakequant_export_skips_iq_packing_guard():
+    exporter = object.__new__(VllmFqGPTModelExporter)
+    exporter.model = torch.nn.Module()
+
+    with (
+        patch.object(uem, "get_tensor_model_parallel_world_size", return_value=2),
+        patch.object(exporter, "_collective_any_iq") as collective,
+    ):
+        exporter._validate_iq_tensor_parallelism()
+
+    collective.assert_not_called()
+
+
 def test_megatron_iq_packer_rejects_tensor_parallelism():
     """Internal packing keeps the TP=1 guard for extra-module export paths."""
     module = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16)
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 92316b01bec..ef9d24540f2 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -171,6 +171,17 @@ def test_export_iq_divisibility_error_identifies_weight(num_bits):
         _export_quantized_weight(linear, torch.bfloat16)
 
 
+@pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"])
+def test_export_iq_dtype_error_identifies_weight(num_bits):
+    weight = torch.ones((4, 256), dtype=torch.uint8)
+
+    with pytest.raises(
+        TypeError,
+        match=rf"Failed to pack {num_bits.upper()} weight 'model.layers.0.weight'",
+    ):
+        _pack_iq_weight(weight, num_bits, describe_as="model.layers.0.weight")
+
+
 def test_export_iq_shape_preflight_reports_every_incompatible_weight():
     model = nn.Sequential(
         nn.Linear(192, 4, bias=False, dtype=torch.bfloat16),

From b536535d6f7cb437677b5f59b7c6380c10bc091d Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 01:16:05 -0700
Subject: [PATCH 18/32] Complete IQ export preflight coverage

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          | 10 ++++--
 modelopt/torch/export/unified_export_hf.py    |  4 +++
 .../torch/export/unified_export_megatron.py   | 18 +++++++++++
 .../export/test_unified_export_megatron.py    | 31 +++++++++++++++++++
 .../torch/export/test_export_diffusers.py     | 29 +++++++++++++++++
 5 files changed, 90 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index d6bb10244d5..540cd63cea5 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -82,8 +82,8 @@
 logger = logging.getLogger(__name__)
 
 
-def _validate_iq_export_weight_shapes(model: nn.Module) -> None:
-    """Reject all IQ weights that cannot be packed before export mutates the model."""
+def _iq_export_weight_shape_errors(model: nn.Module) -> list[str]:
+    """Return selected IQ weights whose final dimension cannot be packed."""
     incompatible: list[str] = []
     for module_name, module in model.named_modules():
         for weight_name in weight_attr_names(module):
@@ -99,6 +99,12 @@ def _validate_iq_export_weight_shapes(model: nn.Module) -> None:
             if weight.dim() == 0 or weight.shape[-1] % group_size:
                 qualified_name = f"{module_name}.{weight_name}".lstrip(".")
                 incompatible.append(f"{qualified_name}: {tuple(weight.shape)}")
+    return incompatible
+
+
+def _validate_iq_export_weight_shapes(model: nn.Module) -> None:
+    """Reject all IQ weights that cannot be packed before export mutates the model."""
+    incompatible = _iq_export_weight_shape_errors(model)
 
     if incompatible:
         details = "\n  - ".join(incompatible)
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index bafcf3ff475..9d19e8c41ef 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -1280,6 +1280,10 @@ def _export_diffusers_checkpoint(
     module_components = {
         name: comp for name, comp in all_components.items() if isinstance(comp, nn.Module)
     }
+    # Validate every component before processing the first one so a later incompatible weight
+    # cannot leave earlier components partially packed.
+    for component in module_components.values():
+        _validate_iq_export_weight_shapes(component)
 
     # Best-effort diffusers pipeline check (kept for folder layout + model_index.json behavior)
     is_diffusers_pipe = False
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index d0ced4be05a..a3b94c62368 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -71,6 +71,7 @@
     iq_format_spec,
 )
 from .quant_utils import (
+    _iq_export_weight_shape_errors,
     _pack_iq_weight,
     _validate_iq_quantizer_config,
     get_activation_scaling_factor,
@@ -275,6 +276,7 @@ def save_pretrained_extra_modules(
     ):
         """Save a EAGLE or Medusa checkpoints which can be deployed by vLLM and TensorRT-LLM."""
         self._validate_iq_tensor_parallelism()
+        self._validate_iq_weight_shapes_collectively()
         # We use the last PP rank to write the config because
         # medusa_heads and eagle_module only exist in the last stage.
         pp_rank = get_pipeline_model_parallel_rank()
@@ -334,6 +336,21 @@ def _validate_iq_tensor_parallelism(self) -> None:
                 "parallel size 1"
             )
 
+    def _validate_iq_weight_shapes_collectively(self) -> None:
+        """Fail every rank before packing when any rank owns an incompatible IQ weight."""
+        if not self._packs_iq_weights:
+            return
+        local_errors = _iq_export_weight_shape_errors(self.model)
+        if not self._collective_any_iq(bool(local_errors)):
+            return
+        if local_errors:
+            details = "\n  - ".join(local_errors)
+            raise ValueError(
+                "IQ export requires every selected weight's final dimension to be divisible by "
+                f"its format group size. Incompatible weights:\n  - {details}"
+            )
+        raise ValueError("Another distributed rank owns an IQ weight that cannot be packed")
+
     def save_pretrained(
         self,
         save_directory: str | os.PathLike,
@@ -345,6 +362,7 @@ def save_pretrained(
             save_directory: Directory to which to save. Will be created if it doesn't exist.
         """
         self._validate_iq_tensor_parallelism()
+        self._validate_iq_weight_shapes_collectively()
 
         pp_rank = get_pipeline_model_parallel_rank()
         pp_size = get_pipeline_model_parallel_world_size()
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 4c7149ff40f..0504a7adc73 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -213,6 +213,37 @@ def test_vllm_fakequant_export_skips_iq_packing_guard():
     collective.assert_not_called()
 
 
+def test_megatron_iq_shape_guard_fails_all_ranks_before_packing():
+    exporter = object.__new__(GPTModelExporter)
+    exporter.model = torch.nn.Linear(192, 2, bias=False, dtype=torch.bfloat16)
+    exporter.model.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+        )
+    )
+
+    with (
+        patch.object(exporter, "_collective_any_iq", return_value=True) as collective,
+        pytest.raises(ValueError, match=r"weight: \(2, 192\)"),
+    ):
+        exporter._validate_iq_weight_shapes_collectively()
+
+    collective.assert_called_once_with(True)
+
+
+def test_megatron_iq_shape_guard_reports_another_rank():
+    exporter = object.__new__(GPTModelExporter)
+    exporter.model = torch.nn.Module()
+
+    with (
+        patch.object(exporter, "_collective_any_iq", return_value=True),
+        pytest.raises(ValueError, match="Another distributed rank"),
+    ):
+        exporter._validate_iq_weight_shapes_collectively()
+
+
 def test_megatron_iq_packer_rejects_tensor_parallelism():
     """Internal packing keeps the TP=1 guard for extra-module export paths."""
     module = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16)
diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py
index 753c81a4b0e..0b04faf98c2 100644
--- a/tests/unit/torch/export/test_export_diffusers.py
+++ b/tests/unit/torch/export/test_export_diffusers.py
@@ -39,6 +39,8 @@
     hide_quantizers_from_state_dict,
 )
 from modelopt.torch.export.unified_export_hf import _postprocess_safetensors, export_hf_checkpoint
+from modelopt.torch.quantization.config import QuantizerAttributeConfig
+from modelopt.torch.quantization.nn import TensorQuantizer
 
 
 def _load_config(config_path):
@@ -127,6 +129,33 @@ def _process_stub(*_args, **_kwargs):
     assert config_data["quantization_config"] == convert_hf_quant_config_format(dummy_quant_config)
 
 
+def test_diffusers_iq_shape_preflight_runs_before_component_mutation(tmp_path, monkeypatch):
+    component = nn.Linear(192, 4, bias=False, dtype=torch.bfloat16)
+    component.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+        )
+    )
+    original_weight = component.weight.detach().clone()
+
+    monkeypatch.setattr(
+        unified_export_hf,
+        "get_diffusion_components",
+        lambda *_args, **_kwargs: {"transformer": component},
+    )
+    monkeypatch.setattr(unified_export_hf, "has_quantized_modules", lambda *_args: True)
+
+    with pytest.raises(ValueError, match=r"weight: \(4, 192\)"):
+        unified_export_hf._export_diffusers_checkpoint(
+            object(), torch.bfloat16, tmp_path, components=None
+        )
+
+    assert component.weight.dtype == torch.bfloat16
+    assert torch.equal(component.weight, original_weight)
+
+
 def test_flux2_dummy_inputs_shape():
     """Verify Flux2-specific dummy input shapes: 4-col RoPE ids, no pooled_projections, guidance."""
     model = get_tiny_flux2()

From c57c7155c5e0ebbb932ca25045b6364c73d423d4 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 01:23:25 -0700
Subject: [PATCH 19/32] Attribute IQ export configuration errors

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          | 26 +++++++++++++------
 .../torch/export/unified_export_megatron.py   |  4 ++-
 tests/unit/torch/export/test_export_weight.py | 20 ++++++++++++++
 3 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 540cd63cea5..4453a3aba73 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -115,21 +115,27 @@ def _validate_iq_export_weight_shapes(model: nn.Module) -> None:
 
 
 def _validate_iq_quantizer_config(
-    module: nn.Module, quantization_format: str, weight_name: str = "weight"
+    module: nn.Module,
+    quantization_format: str,
+    weight_name: str = "weight",
+    describe_as: str | None = None,
 ) -> None:
     """Reject IQ export settings that the checkpoint encoder cannot reproduce."""
+    where = f" for '{describe_as}'" if describe_as else ""
     quantizer = representative_weight_quantizer(module, weight_name)
     if not isinstance(quantizer, TensorQuantizer):
-        raise ValueError(f"{quantization_format.upper()} export requires one TensorQuantizer")
+        raise ValueError(
+            f"{quantization_format.upper()} export requires one TensorQuantizer{where}"
+        )
     if quantizer.num_bits != quantization_format or quantizer.backend != "ggml":
         raise ValueError(
-            f"{quantization_format.upper()} export requires the matching ggml quantizer"
+            f"{quantization_format.upper()} export requires the matching ggml quantizer{where}"
         )
     extra_args = quantizer.backend_extra_args or {}
     search_impl = extra_args.get("search_impl", "auto")
     if search_impl != "auto":
         raise NotImplementedError(
-            f"{quantization_format.upper()} export only supports search_impl='auto'"
+            f"{quantization_format.upper()} export only supports search_impl='auto'{where}"
         )
     quantizer_names = quantizer_attr_names(weight_name)
     for quantizer_name in (quantizer_names.input_quantizer, quantizer_names.output_quantizer):
@@ -139,7 +145,8 @@ def _validate_iq_quantizer_config(
             or getattr(activation_quantizer, "pre_quant_scale", None) is not None
         ):
             raise NotImplementedError(
-                f"{quantization_format.upper()} export is weight-only; activation quantization "
+                f"{quantization_format.upper()} export is weight-only{where}; activation "
+                "quantization "
                 "and pre_quant_scale are not represented in the checkpoint"
             )
 
@@ -153,7 +160,11 @@ def _pack_iq_weight(
 ) -> torch.Tensor:
     """Pack one IQ weight, using ``describe_as`` only for failure attribution."""
     if module is not None:
-        _validate_iq_quantizer_config(module, quantization_format, weight_name)
+        if describe_as is None:
+            describe_as = f"{type(module).__name__}.{weight_name}"
+        _validate_iq_quantizer_config(
+            module, quantization_format, weight_name, describe_as=describe_as
+        )
     if quantization_format not in IQ_FORMATS:
         raise ValueError(f"Unsupported IQ quantization format: {quantization_format}")
     try:
@@ -167,8 +178,7 @@ def _pack_iq_weight(
             )
     except (TypeError, ValueError) as exc:
         if describe_as is None:
-            owner = type(module).__name__ if module is not None else "state_dict"
-            describe_as = f"{owner}.{weight_name}"
+            describe_as = f"state_dict.{weight_name}"
         raise type(exc)(
             f"Failed to pack {quantization_format.upper()} weight "
             f"'{describe_as}' with shape {tuple(weight.shape)}: {exc}"
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index a3b94c62368..a97cb4dec3c 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -1162,7 +1162,9 @@ def _get_quantized_state(
 
         is_iq = qformat in IQ_FORMATS
         if is_iq:
-            _validate_iq_quantizer_config(module, qformat)
+            _validate_iq_quantizer_config(
+                module, qformat, describe_as=f"{prefix}weight" if prefix else "weight"
+            )
         name_to_value = self._get_weight_bias(
             module, dtype, name_to_value, keep_weight_device=is_iq
         )
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index ef9d24540f2..1052e0d42ff 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -270,6 +270,26 @@ def test_export_iq_rejects_non_weight_only_config(num_bits, quantizer_name, mode
         _export_quantized_weight(linear, torch.bfloat16)
 
 
+def test_iq_config_error_identifies_weight():
+    linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)
+    linear.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+        )
+    )
+    linear.input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8, axis=None))
+
+    with pytest.raises(NotImplementedError, match=r"model\.layers\.3\.mlp\.down_proj\.weight"):
+        _pack_iq_weight(
+            linear.weight,
+            "iq2_xs",
+            module=linear,
+            describe_as="model.layers.3.mlp.down_proj.weight",
+        )
+
+
 class QuantMoELinear(nn.Module):
     def __init__(self):
         super().__init__()

From 215b240158095ed0c84e3e2b42f384e8e34f42c7 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 01:57:55 -0700
Subject: [PATCH 20/32] Refine IQ export preflight

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/hf_export_handlers.py   | 12 +++--
 modelopt/torch/export/moe_utils.py            |  6 ++-
 modelopt/torch/export/quant_utils.py          |  4 +-
 modelopt/torch/export/unified_export_hf.py    |  3 +-
 .../torch/export/unified_export_megatron.py   | 49 ++++++++++---------
 .../export/test_unified_export_megatron.py    | 46 ++++++++++++-----
 .../unit/torch/export/test_export_registry.py | 18 +++++++
 tests/unit/torch/export/test_export_weight.py |  7 +++
 .../plugins/test_fused_experts.py             | 22 +++++++--
 9 files changed, 117 insertions(+), 50 deletions(-)

diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py
index 550fb0dbd6a..6ebc02a304f 100644
--- a/modelopt/torch/export/hf_export_handlers.py
+++ b/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)
diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py
index 1cabef8a37d..892267b64ed 100644
--- a/modelopt/torch/export/moe_utils.py
+++ b/modelopt/torch/export/moe_utils.py
@@ -50,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.
 
@@ -225,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
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 4453a3aba73..c26fb0a05fe 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -159,14 +159,14 @@ def _pack_iq_weight(
     describe_as: str | None = None,
 ) -> torch.Tensor:
     """Pack one IQ weight, using ``describe_as`` only for failure attribution."""
+    if quantization_format not in IQ_FORMATS:
+        raise ValueError(f"Unsupported IQ quantization format: {quantization_format}")
     if module is not None:
         if describe_as is None:
             describe_as = f"{type(module).__name__}.{weight_name}"
         _validate_iq_quantizer_config(
             module, quantization_format, weight_name, describe_as=describe_as
         )
-    if quantization_format not in IQ_FORMATS:
-        raise ValueError(f"Unsupported IQ quantization format: {quantization_format}")
     try:
         if quantization_format == QUANTIZATION_IQ1_S:
             packed_weight, _ = quantize_iq1_s(weight)
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 9d19e8c41ef..5e95c2e9f42 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -597,6 +597,7 @@ def _export_quantized_weight(
     sub_module: nn.Module,
     dtype: torch.dtype,
     weight_name: str = "weight",
+    describe_as: str | None = None,
 ):
     """For the given weight attr of the sub_module, export the quantization info of it.
 
@@ -631,7 +632,7 @@ def _export_quantized_weight(
                 f"attribute, got {weight_name!r} on {type(sub_module).__name__}"
             )
         packed_weight = _pack_iq_weight(
-            weight.to(dtype), quantization_format, sub_module, weight_name
+            weight.to(dtype), quantization_format, sub_module, weight_name, describe_as
         )
         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 a97cb4dec3c..be5c067e04c 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -275,8 +275,7 @@ def save_pretrained_extra_modules(
         save_directory: str | os.PathLike,
     ):
         """Save a EAGLE or Medusa checkpoints which can be deployed by vLLM and TensorRT-LLM."""
-        self._validate_iq_tensor_parallelism()
-        self._validate_iq_weight_shapes_collectively()
+        self._validate_iq_export()
         # We use the last PP rank to write the config because
         # medusa_heads and eagle_module only exist in the last stage.
         pp_rank = get_pipeline_model_parallel_rank()
@@ -311,37 +310,40 @@ def _model_has_iq_quantizer(model: torch.nn.Module) -> bool:
         )
 
     @staticmethod
-    def _collective_any_iq(local_iq: bool) -> bool:
-        """Return whether any distributed rank owns an IQ-quantized module."""
+    def _collective_iq_export_flags(local_iq: bool, local_bad_shape: bool) -> tuple[bool, bool]:
+        """Combine IQ presence and incompatible-shape flags across all ranks."""
         if not torch.distributed.is_initialized():
-            return local_iq
+            return local_iq, local_bad_shape
         device = (
             torch.device("cuda", torch.cuda.current_device())
             if torch.distributed.get_backend() == torch.distributed.Backend.NCCL
             else torch.device("cpu")
         )
-        any_iq = torch.tensor(int(local_iq), dtype=torch.int32, device=device)
-        torch.distributed.all_reduce(any_iq, op=torch.distributed.ReduceOp.MAX)
-        return bool(any_iq.item())
+        flags = torch.tensor(
+            [int(local_iq), int(local_bad_shape)], dtype=torch.int32, device=device
+        )
+        torch.distributed.all_reduce(flags, op=torch.distributed.ReduceOp.MAX)
+        return bool(flags[0].item()), bool(flags[1].item())
 
-    def _validate_iq_tensor_parallelism(self) -> None:
-        """Fail collectively when this exporter would pack IQ weights across TP ranks."""
-        if (
-            self._packs_iq_weights
-            and get_tensor_model_parallel_world_size() != 1
-            and self._collective_any_iq(self._model_has_iq_quantizer(self.model))
-        ):
+    def _validate_iq_export(self) -> None:
+        """Reject unsupported IQ export configurations in one collective preflight."""
+        if not self._packs_iq_weights:
+            return
+
+        local_iq = self._model_has_iq_quantizer(self.model)
+        tp_size = get_tensor_model_parallel_world_size()
+        local_errors = (
+            _iq_export_weight_shape_errors(self.model) if local_iq and tp_size == 1 else []
+        )
+        any_iq, any_bad_shape = self._collective_iq_export_flags(local_iq, bool(local_errors))
+        if not any_iq:
+            return
+        if tp_size != 1:
             raise NotImplementedError(
                 "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model "
                 "parallel size 1"
             )
-
-    def _validate_iq_weight_shapes_collectively(self) -> None:
-        """Fail every rank before packing when any rank owns an incompatible IQ weight."""
-        if not self._packs_iq_weights:
-            return
-        local_errors = _iq_export_weight_shape_errors(self.model)
-        if not self._collective_any_iq(bool(local_errors)):
+        if not any_bad_shape:
             return
         if local_errors:
             details = "\n  - ".join(local_errors)
@@ -361,8 +363,7 @@ def save_pretrained(
         Args:
             save_directory: Directory to which to save. Will be created if it doesn't exist.
         """
-        self._validate_iq_tensor_parallelism()
-        self._validate_iq_weight_shapes_collectively()
+        self._validate_iq_export()
 
         pp_rank = get_pipeline_model_parallel_rank()
         pp_size = get_pipeline_model_parallel_world_size()
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 0504a7adc73..b6da351e6d8 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -160,16 +160,16 @@ def test_megatron_iq_export_rejects_tensor_parallelism():
         exporter.save_pretrained("unused", "unused")
 
 
-def test_megatron_iq_tensor_parallel_guard_is_collective():
-    def report_remote_iq(flag, **_kwargs):
-        flag.fill_(1)
+def test_megatron_iq_export_guard_uses_one_collective_for_both_flags():
+    def report_remote_iq(flags, **_kwargs):
+        flags.copy_(torch.tensor([1, 0], dtype=flags.dtype, device=flags.device))
 
     with (
         patch.object(torch.distributed, "is_initialized", return_value=True),
         patch.object(torch.distributed, "get_backend", return_value=torch.distributed.Backend.GLOO),
         patch.object(torch.distributed, "all_reduce", side_effect=report_remote_iq),
     ):
-        assert GPTModelExporter._collective_any_iq(False)
+        assert GPTModelExporter._collective_iq_export_flags(False, False) == (True, False)
 
 
 def test_megatron_iq_tensor_parallel_guard_finds_later_mixed_format():
@@ -194,7 +194,7 @@ def test_megatron_iq_tensor_parallel_guard_covers_extra_modules():
     with (
         patch.object(uem, "get_tensor_model_parallel_world_size", return_value=2),
         patch.object(exporter, "_model_has_iq_quantizer", return_value=True),
-        patch.object(exporter, "_collective_any_iq", return_value=True),
+        patch.object(exporter, "_collective_iq_export_flags", return_value=(True, False)),
         pytest.raises(NotImplementedError, match="tensor model parallel size 1"),
     ):
         exporter.save_pretrained_extra_modules("unused")
@@ -206,13 +206,31 @@ def test_vllm_fakequant_export_skips_iq_packing_guard():
 
     with (
         patch.object(uem, "get_tensor_model_parallel_world_size", return_value=2),
-        patch.object(exporter, "_collective_any_iq") as collective,
+        patch.object(exporter, "_collective_iq_export_flags") as collective,
     ):
-        exporter._validate_iq_tensor_parallelism()
+        exporter._validate_iq_export()
 
     collective.assert_not_called()
 
 
+def test_megatron_non_iq_export_skips_shape_scan():
+    exporter = object.__new__(GPTModelExporter)
+    exporter.model = torch.nn.Module()
+
+    with (
+        patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1),
+        patch.object(exporter, "_model_has_iq_quantizer", return_value=False),
+        patch.object(
+            exporter, "_collective_iq_export_flags", return_value=(False, False)
+        ) as collective,
+        patch.object(uem, "_iq_export_weight_shape_errors") as shape_errors,
+    ):
+        exporter._validate_iq_export()
+
+    collective.assert_called_once_with(False, False)
+    shape_errors.assert_not_called()
+
+
 def test_megatron_iq_shape_guard_fails_all_ranks_before_packing():
     exporter = object.__new__(GPTModelExporter)
     exporter.model = torch.nn.Linear(192, 2, bias=False, dtype=torch.bfloat16)
@@ -225,12 +243,15 @@ def test_megatron_iq_shape_guard_fails_all_ranks_before_packing():
     )
 
     with (
-        patch.object(exporter, "_collective_any_iq", return_value=True) as collective,
+        patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1),
+        patch.object(
+            exporter, "_collective_iq_export_flags", return_value=(True, True)
+        ) as collective,
         pytest.raises(ValueError, match=r"weight: \(2, 192\)"),
     ):
-        exporter._validate_iq_weight_shapes_collectively()
+        exporter._validate_iq_export()
 
-    collective.assert_called_once_with(True)
+    collective.assert_called_once_with(True, True)
 
 
 def test_megatron_iq_shape_guard_reports_another_rank():
@@ -238,10 +259,11 @@ def test_megatron_iq_shape_guard_reports_another_rank():
     exporter.model = torch.nn.Module()
 
     with (
-        patch.object(exporter, "_collective_any_iq", return_value=True),
+        patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1),
+        patch.object(exporter, "_collective_iq_export_flags", return_value=(True, True)),
         pytest.raises(ValueError, match="Another distributed rank"),
     ):
-        exporter._validate_iq_weight_shapes_collectively()
+        exporter._validate_iq_export()
 
 
 def test_megatron_iq_packer_rejects_tensor_parallelism():
diff --git a/tests/unit/torch/export/test_export_registry.py b/tests/unit/torch/export/test_export_registry.py
index 4c93388ca08..5a57dbdc79f 100644
--- a/tests/unit/torch/export/test_export_registry.py
+++ b/tests/unit/torch/export/test_export_registry.py
@@ -40,6 +40,8 @@
     _ExportHandlerRegistryCls,
 )
 from modelopt.torch.export.unified_export_hf import _process_quantized_modules
+from modelopt.torch.quantization.config import QuantizerAttributeConfig
+from modelopt.torch.quantization.nn import TensorQuantizer
 
 
 class _Experts(nn.Module):
@@ -303,6 +305,22 @@ def test_process_quantized_modules_exports_via_registry():
         assert weight.dtype == torch.float8_e4m3fn
 
 
+def test_iq_handler_attributes_configuration_error_to_qualified_weight():
+    linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)
+    linear.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+        )
+    )
+    linear.input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8, axis=None))
+    ctx = ExportContext(model=linear, dtype=torch.bfloat16)
+
+    with pytest.raises(NotImplementedError, match=r"model\.layers\.3\.mlp\.down_proj\.weight"):
+        _export_quant_linear("model.layers.3.mlp.down_proj", linear, ctx)
+
+
 def test_export_context_carries_no_resolver():
     # Tied-weight dedup is driven by the driver's resolver (fed to sync_tied_input_amax /
     # postprocess_state_dict), not by handlers, so ExportContext no longer builds or holds
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 1052e0d42ff..04efa2d90e0 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -182,6 +182,13 @@ def test_export_iq_dtype_error_identifies_weight(num_bits):
         _pack_iq_weight(weight, num_bits, describe_as="model.layers.0.weight")
 
 
+def test_iq_packer_rejects_non_iq_format_before_inspecting_module():
+    linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)
+
+    with pytest.raises(ValueError, match="Unsupported IQ quantization format: fp8"):
+        _pack_iq_weight(linear.weight, "fp8", module=linear)
+
+
 def test_export_iq_shape_preflight_reports_every_incompatible_weight():
     model = nn.Sequential(
         nn.Linear(192, 4, bias=False, dtype=torch.bfloat16),
diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index 10289ac71c2..d2b91a4ceb3 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -525,14 +525,21 @@ def test_export_preserves_packed_iq_weights(self, monkeypatch):
             for quantizer in quantizers:
                 quantizer.set_from_attribute_config(config)
 
-            monkeypatch.setattr(
-                "modelopt.torch.export.unified_export_hf._pack_iq_weight",
-                lambda weight, *_args: torch.zeros(
+            labels = []
+
+            def fake_pack(weight, *_args):
+                labels.append(_args[-1])
+                return torch.zeros(
                     (weight.numel() // 256, 74), dtype=torch.uint8, device=weight.device
-                ),
+                )
+
+            monkeypatch.setattr(
+                "modelopt.torch.export.unified_export_hf._pack_iq_weight", fake_pack
             )
 
-            _export_fused_experts(converted, torch.float16)
+            _export_fused_experts(
+                converted, torch.float16, module_name="model.layers.0.mlp.experts"
+            )
 
             state_dict = converted.state_dict()
             for idx in range(NUM_EXPERTS):
@@ -543,6 +550,11 @@ def test_export_preserves_packed_iq_weights(self, monkeypatch):
                     assert isinstance(
                         getattr(getattr(converted, str(idx)), projection).weight, nn.Parameter
                     )
+            assert labels == [
+                f"model.layers.0.mlp.experts.{idx}.{projection}.weight"
+                for idx in range(NUM_EXPERTS)
+                for projection in ("gate_proj", "up_proj", "down_proj")
+            ]
         finally:
             self._cleanup_registry(expert_type)
 

From 60179e46762c78af48cd6582e8d4e17d5c64a8f6 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 02:24:14 -0700
Subject: [PATCH 21/32] Validate grouped IQ export weights

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          | 34 +++++++++++++++----
 tests/unit/torch/export/test_export_weight.py | 30 +++++++++++++++-
 2 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index c26fb0a05fe..22b827a36e3 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -48,7 +48,12 @@
 )
 from modelopt.torch.utils import clear_cuda_cache
 
-from ..quantization.nn import NVFP4StaticQuantizer, SequentialQuantizer, TensorQuantizer
+from ..quantization.nn import (
+    GroupedQuantizer,
+    NVFP4StaticQuantizer,
+    SequentialQuantizer,
+    TensorQuantizer,
+)
 from .model_utils import TiedWeightMap, get_language_model_from_vl
 from .quant_format import (
     IQ_FORMATS,
@@ -86,19 +91,36 @@ def _iq_export_weight_shape_errors(model: nn.Module) -> list[str]:
     """Return selected IQ weights whose final dimension cannot be packed."""
     incompatible: list[str] = []
     for module_name, module in model.named_modules():
-        for weight_name in weight_attr_names(module):
-            quantizer = representative_weight_quantizer(module, weight_name)
+        inspected: set[str] = set()
+
+        def inspect_weight(weight_name: str, quantizer: nn.Module | None) -> None:
+            qualified_name = f"{module_name}.{weight_name}".lstrip(".")
+            if qualified_name in inspected:
+                return
+            inspected.add(qualified_name)
             if not (
                 isinstance(quantizer, TensorQuantizer)
                 and quantizer.is_enabled
                 and quantizer.num_bits in IQ_FORMATS
             ):
-                continue
-            weight = getattr(module, weight_name)
+                return
+            weight = getattr(module, weight_name, None)
+            if not isinstance(weight, torch.Tensor):
+                return
             group_size = iq_format_spec(quantizer.num_bits)["group_size"]
             if weight.dim() == 0 or weight.shape[-1] % group_size:
-                qualified_name = f"{module_name}.{weight_name}".lstrip(".")
                 incompatible.append(f"{qualified_name}: {tuple(weight.shape)}")
+
+        for weight_name in weight_attr_names(module):
+            quantizer = representative_weight_quantizer(module, weight_name)
+            inspect_weight(weight_name, quantizer)
+
+        grouped_quantizer = getattr(module, "weight_quantizer", None)
+        if isinstance(grouped_quantizer, GroupedQuantizer) and len(grouped_quantizer) > 0:
+            num_gemms = int(getattr(module, "num_gemms", len(grouped_quantizer)))
+            for index in range(num_gemms):
+                quantizer = grouped_quantizer[min(index, len(grouped_quantizer) - 1)]
+                inspect_weight(f"weight{index}", quantizer)
     return incompatible
 
 
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 04efa2d90e0..45a02299507 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -31,7 +31,7 @@
     export_hf_checkpoint,
 )
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
-from modelopt.torch.quantization.nn import TensorQuantizer
+from modelopt.torch.quantization.nn import GroupedQuantizer, TensorQuantizer
 from modelopt.torch.quantization.utils import quantizer_attr_names
 
 
@@ -213,6 +213,34 @@ def test_export_iq_shape_preflight_reports_every_incompatible_weight():
     assert "1.weight" not in message
 
 
+def test_export_iq_shape_preflight_reports_grouped_weights():
+    class GroupedLinear(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.num_gemms = 2
+            self.weight0 = nn.Parameter(torch.ones(4, 256, dtype=torch.bfloat16))
+            self.weight1 = nn.Parameter(torch.ones(4, 192, dtype=torch.bfloat16))
+            quantizer_config = QuantizerAttributeConfig(
+                num_bits="iq2_xs",
+                block_sizes={-1: 256},
+                backend="ggml",
+            )
+            self.weight_quantizer = GroupedQuantizer(
+                TensorQuantizer(quantizer_config),
+                TensorQuantizer(quantizer_config),
+            )
+
+    model = nn.Module()
+    model.experts = GroupedLinear()
+
+    with pytest.raises(ValueError) as exc_info:
+        _validate_iq_export_weight_shapes(model)
+
+    message = str(exc_info.value)
+    assert "experts.weight1: (4, 192)" in message
+    assert "experts.weight0" not in message
+
+
 def test_export_hf_checkpoint_runs_iq_shape_preflight_before_mutation(tmp_path):
     linear = nn.Linear(192, 4, bias=False, dtype=torch.bfloat16)
     linear.weight_quantizer = TensorQuantizer(

From 6c1d18ee32b1097280d29acf1205d3cd7238c1e1 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 03:22:55 -0700
Subject: [PATCH 22/32] Preflight unsupported IQ export weights

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          | 30 ++++++++----
 modelopt/torch/export/unified_export_hf.py    |  3 +-
 tests/unit/torch/export/test_export_weight.py | 48 +++++++++++++++++++
 3 files changed, 72 insertions(+), 9 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 22b827a36e3..34197ea20c8 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -88,12 +88,17 @@
 
 
 def _iq_export_weight_shape_errors(model: nn.Module) -> list[str]:
-    """Return selected IQ weights whose final dimension cannot be packed."""
+    """Return selected IQ weights that cannot be handled by unified export."""
     incompatible: list[str] = []
     for module_name, module in model.named_modules():
         inspected: set[str] = set()
 
-        def inspect_weight(weight_name: str, quantizer: nn.Module | None) -> None:
+        def inspect_weight(
+            weight_name: str,
+            quantizer: nn.Module | None,
+            *,
+            allow_nonstandard_name: bool = False,
+        ) -> None:
             qualified_name = f"{module_name}.{weight_name}".lstrip(".")
             if qualified_name in inspected:
                 return
@@ -104,6 +109,11 @@ def inspect_weight(weight_name: str, quantizer: nn.Module | None) -> None:
                 and quantizer.num_bits in IQ_FORMATS
             ):
                 return
+            if weight_name != "weight" and not allow_nonstandard_name:
+                incompatible.append(
+                    f"{qualified_name}: nonstandard weight attributes are not supported"
+                )
+                return
             weight = getattr(module, weight_name, None)
             if not isinstance(weight, torch.Tensor):
                 return
@@ -120,19 +130,21 @@ def inspect_weight(weight_name: str, quantizer: nn.Module | None) -> None:
             num_gemms = int(getattr(module, "num_gemms", len(grouped_quantizer)))
             for index in range(num_gemms):
                 quantizer = grouped_quantizer[min(index, len(grouped_quantizer) - 1)]
-                inspect_weight(f"weight{index}", quantizer)
+                # Grouped linears are rebuilt as standard per-expert modules before packing.
+                inspect_weight(f"weight{index}", quantizer, allow_nonstandard_name=True)
     return incompatible
 
 
 def _validate_iq_export_weight_shapes(model: nn.Module) -> None:
-    """Reject all IQ weights that cannot be packed before export mutates the model."""
+    """Reject all unsupported IQ weights before export mutates the model."""
     incompatible = _iq_export_weight_shape_errors(model)
 
     if incompatible:
         details = "\n  - ".join(incompatible)
         raise ValueError(
-            "IQ export requires every selected weight's final dimension to be divisible by "
-            f"its format group size. Incompatible weights:\n  - {details}"
+            "IQ export cannot handle the following selected weights. Each weight must use the "
+            "standard 'weight' attribute and have a final dimension divisible by its format "
+            f"group size:\n  - {details}"
         )
 
 
@@ -213,9 +225,11 @@ def _pack_iq_weight(
         spec["block_payload_bytes"],
     )
     if packed_weight.shape != expected_shape:
+        where = describe_as or f"state_dict.{weight_name}"
         raise RuntimeError(
-            f"{quantization_format.upper()} packing produced {tuple(packed_weight.shape)}, "
-            f"expected {expected_shape}; logical-shape recovery is no longer valid"
+            f"{quantization_format.upper()} packing for weight '{where}' produced "
+            f"{tuple(packed_weight.shape)}, expected {expected_shape}; logical-shape recovery "
+            "is no longer valid"
         )
     return packed_weight
 
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 5e95c2e9f42..6d75a09a83a 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -627,9 +627,10 @@ def _export_quantized_weight(
 
     if quantization_format in IQ_FORMATS:
         if weight_name != "weight":
+            where = describe_as or f"{type(sub_module).__name__}.{weight_name}"
             raise NotImplementedError(
                 "IQ unified export currently supports modules with a standard 'weight' "
-                f"attribute, got {weight_name!r} on {type(sub_module).__name__}"
+                f"attribute, got '{where}'"
             )
         packed_weight = _pack_iq_weight(
             weight.to(dtype), quantization_format, sub_module, weight_name, describe_as
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 45a02299507..15b65588254 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -19,6 +19,7 @@
 import torch.nn as nn
 from _test_utils.torch.export.utils import ToyModel, partial_fp8_config, partial_w4a8_config
 
+import modelopt.torch.export.quant_utils as quant_utils
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.quant_utils import (
     _pack_iq_weight,
@@ -182,6 +183,21 @@ def test_export_iq_dtype_error_identifies_weight(num_bits):
         _pack_iq_weight(weight, num_bits, describe_as="model.layers.0.weight")
 
 
+def test_export_iq_payload_shape_error_identifies_weight(monkeypatch):
+    weight = torch.ones((4, 256), dtype=torch.bfloat16)
+
+    def pack_with_wrong_shape(weight):
+        return torch.zeros((4, 2, 74), dtype=torch.uint8), None
+
+    monkeypatch.setattr(quant_utils, "quantize_iq2_xs", pack_with_wrong_shape)
+
+    with pytest.raises(
+        RuntimeError,
+        match=r"packing for weight 'model\.layers\.0\.weight'.*expected \(4, 1, 74\)",
+    ):
+        _pack_iq_weight(weight, "iq2_xs", describe_as="model.layers.0.weight")
+
+
 def test_iq_packer_rejects_non_iq_format_before_inspecting_module():
     linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)
 
@@ -241,6 +257,38 @@ def __init__(self):
     assert "experts.weight0" not in message
 
 
+def test_export_iq_preflight_rejects_nonstandard_weight_before_mutation(tmp_path):
+    class CustomWeightModule(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.proj = nn.Parameter(torch.ones(4, 256, dtype=torch.bfloat16))
+            self.proj_weight_quantizer = TensorQuantizer(
+                QuantizerAttributeConfig(
+                    num_bits="iq2_xs",
+                    block_sizes={-1: 256},
+                    backend="ggml",
+                )
+            )
+
+    model = nn.Module()
+    model.block = CustomWeightModule()
+    original_weight = model.block.proj.detach().clone()
+
+    with pytest.raises(ValueError, match=r"block\.proj: nonstandard weight"):
+        export_hf_checkpoint(model, export_dir=tmp_path)
+
+    assert model.block.proj.dtype == torch.bfloat16
+    assert torch.equal(model.block.proj, original_weight)
+
+    with pytest.raises(NotImplementedError, match=r"got 'model\.block\.proj'"):
+        _export_quantized_weight(
+            model.block,
+            torch.bfloat16,
+            weight_name="proj",
+            describe_as="model.block.proj",
+        )
+
+
 def test_export_hf_checkpoint_runs_iq_shape_preflight_before_mutation(tmp_path):
     linear = nn.Linear(192, 4, bias=False, dtype=torch.bfloat16)
     linear.weight_quantizer = TensorQuantizer(

From 0a397aac93b8f859e16217d6add77c7c0b3eb907 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 03:56:58 -0700
Subject: [PATCH 23/32] Preserve fused IQ export preflight

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          | 19 ++++++-
 .../plugins/test_fused_experts.py             | 49 +++++++++++++++++++
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 34197ea20c8..b166ac5fd4a 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -93,6 +93,18 @@ def _iq_export_weight_shape_errors(model: nn.Module) -> list[str]:
     for module_name, module in model.named_modules():
         inspected: set[str] = set()
 
+        def uses_plural_quantizers(weight_name: str) -> bool:
+            quantizer_name = quantizer_attr_names(weight_name).weight_quantizer
+            if isinstance(getattr(module, quantizer_name + "s", None), nn.ModuleList):
+                return True
+            return (
+                weight_name == getattr(module, "_first_proj_attr", None)
+                and weight_name != "gate_up_proj"
+                and isinstance(
+                    getattr(module, "gate_up_proj_weight_quantizers", None), nn.ModuleList
+                )
+            )
+
         def inspect_weight(
             weight_name: str,
             quantizer: nn.Module | None,
@@ -123,7 +135,12 @@ def inspect_weight(
 
         for weight_name in weight_attr_names(module):
             quantizer = representative_weight_quantizer(module, weight_name)
-            inspect_weight(weight_name, quantizer)
+            # Fused-expert plural quantizer lists are rebuilt as standard per-expert modules.
+            inspect_weight(
+                weight_name,
+                quantizer,
+                allow_nonstandard_name=uses_plural_quantizers(weight_name),
+            )
 
         grouped_quantizer = getattr(module, "weight_quantizer", None)
         if isinstance(grouped_quantizer, GroupedQuantizer) and len(grouped_quantizer) > 0:
diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py
index d2b91a4ceb3..402130942fc 100644
--- a/tests/unit/torch/quantization/plugins/test_fused_experts.py
+++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py
@@ -30,6 +30,7 @@
 import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module
 from modelopt.torch.export.moe_utils import _export_fused_experts
 from modelopt.torch.export.quant_utils import get_quant_config, get_quantization_format
+from modelopt.torch.export.unified_export_hf import export_hf_checkpoint
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.conversion import _normalize_fused_experts_quantizer_name
 from modelopt.torch.quantization.model_calib import local_hessian_calibrate
@@ -558,6 +559,54 @@ def fake_pack(weight, *_args):
         finally:
             self._cleanup_registry(expert_type)
 
+    def test_export_hf_checkpoint_preflight_allows_fused_iq_weights(self, monkeypatch, tmp_path):
+        """The model-wide preflight must preserve the fused-expert split path."""
+        model = _TinyMoEModel()
+        expert_type = type(model.moe.experts)
+        self._cleanup_registry(expert_type)
+        register_fused_experts_on_the_fly(model)
+
+        try:
+            converted = QuantModuleRegistry.convert(model.moe.experts)
+            converted.gate_up_proj = nn.Parameter(torch.ones(NUM_EXPERTS, 512, 256))
+            converted.down_proj = nn.Parameter(torch.ones(NUM_EXPERTS, 256, 256))
+            config = QuantizerAttributeConfig(
+                num_bits="iq2_xs",
+                block_sizes={-1: 256},
+                backend="ggml",
+                backend_extra_args={"search_impl": "auto"},
+            )
+            quantizers = list(converted.gate_up_proj_weight_quantizers) + list(
+                converted.down_proj_weight_quantizers
+            )
+            for quantizer in quantizers:
+                quantizer.set_from_attribute_config(config)
+
+            export_called = False
+
+            def fake_export(*_args, **_kwargs):
+                nonlocal export_called
+                export_called = True
+                return {}, {}
+
+            monkeypatch.setattr(
+                "modelopt.torch.export.unified_export_hf._export_transformers_checkpoint",
+                fake_export,
+            )
+            monkeypatch.setattr(
+                model, "save_pretrained", lambda *_args, **_kwargs: None, raising=False
+            )
+            monkeypatch.setattr(
+                "modelopt.torch.export.unified_export_hf._write_hf_export_config",
+                lambda *_args, **_kwargs: None,
+            )
+
+            export_hf_checkpoint(model, export_dir=tmp_path)
+
+            assert export_called
+        finally:
+            self._cleanup_registry(expert_type)
+
     def test_iq_export_does_not_create_or_warn_about_amax(self, monkeypatch):
         """IQ expert packing is weight-only and must not enter amax fallback paths."""
         model = _TinyMoEModel()

From 5033f55e9e518458e35921d56392e3f1ca96bc0b Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 04:18:16 -0700
Subject: [PATCH 24/32] Cover grouped export preflight unions

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          |  2 +-
 tests/unit/torch/export/test_export_weight.py | 24 +++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index b166ac5fd4a..ab462b80e59 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -145,7 +145,7 @@ def inspect_weight(
         grouped_quantizer = getattr(module, "weight_quantizer", None)
         if isinstance(grouped_quantizer, GroupedQuantizer) and len(grouped_quantizer) > 0:
             num_gemms = int(getattr(module, "num_gemms", len(grouped_quantizer)))
-            for index in range(num_gemms):
+            for index in range(max(num_gemms, len(grouped_quantizer))):
                 quantizer = grouped_quantizer[min(index, len(grouped_quantizer) - 1)]
                 # Grouped linears are rebuilt as standard per-expert modules before packing.
                 inspect_weight(f"weight{index}", quantizer, allow_nonstandard_name=True)
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 15b65588254..674fcebf40e 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -257,6 +257,30 @@ def __init__(self):
     assert "experts.weight0" not in message
 
 
+def test_export_iq_shape_preflight_reports_grouped_quantizers_beyond_num_gemms():
+    class GroupedLinear(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.num_gemms = 1
+            self.weight0 = nn.Parameter(torch.ones(4, 256, dtype=torch.bfloat16))
+            self.weight1 = nn.Parameter(torch.ones(4, 192, dtype=torch.bfloat16))
+            quantizer_config = QuantizerAttributeConfig(
+                num_bits="iq2_xs",
+                block_sizes={-1: 256},
+                backend="ggml",
+            )
+            self.weight_quantizer = GroupedQuantizer(
+                TensorQuantizer(quantizer_config),
+                TensorQuantizer(quantizer_config),
+            )
+
+    model = nn.Module()
+    model.experts = GroupedLinear()
+
+    with pytest.raises(ValueError, match=r"experts\.weight1: \(4, 192\)"):
+        _validate_iq_export_weight_shapes(model)
+
+
 def test_export_iq_preflight_rejects_nonstandard_weight_before_mutation(tmp_path):
     class CustomWeightModule(nn.Module):
         def __init__(self):

From 8a6ed077288405dfcdbffdf7773a7c5967711aeb Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 04:39:06 -0700
Subject: [PATCH 25/32] Preflight direct IQ checkpoint exports

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py    |  1 +
 tests/unit/torch/export/test_export_weight.py | 27 +++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 6d75a09a83a..fa12577b731 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -988,6 +988,7 @@ def _prepare_model_for_export(model, dtype, is_modelopt_qlora):
     the quant config.
     """
     dtype = _resolve_export_dtype(model, dtype)
+    _validate_iq_export_weight_shapes(model)
     # One tied-weight map for the whole export (amax sync + final dedup in postprocess_state_dict).
     # Sourced from HF's name-based all_tied_weights_keys, so it is correct even under FSDP/offload.
     tied_map = TiedWeightMap(model)
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 674fcebf40e..a4ceb88b8b4 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -14,6 +14,8 @@
 # limitations under the License.
 
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 import torch.nn as nn
@@ -28,6 +30,7 @@
 )
 from modelopt.torch.export.unified_export_hf import (
     _export_quantized_weight,
+    _export_transformers_checkpoint,
     _process_quantized_modules,
     export_hf_checkpoint,
 )
@@ -331,6 +334,30 @@ def test_export_hf_checkpoint_runs_iq_shape_preflight_before_mutation(tmp_path):
     assert torch.equal(linear.weight, original_weight)
 
 
+def test_export_transformers_checkpoint_runs_iq_shape_preflight_before_mutation():
+    model = nn.Sequential(
+        nn.Linear(256, 4, bias=False, dtype=torch.bfloat16),
+        nn.Linear(192, 4, bias=False, dtype=torch.bfloat16),
+    )
+    model.config = SimpleNamespace(torch_dtype=torch.bfloat16)
+    for linear in model:
+        linear.weight_quantizer = TensorQuantizer(
+            QuantizerAttributeConfig(
+                num_bits="iq2_xs",
+                block_sizes={-1: 256},
+                backend="ggml",
+            )
+        )
+    original_weights = [linear.weight.detach().clone() for linear in model]
+
+    with pytest.raises(ValueError, match=r"1\.weight: \(4, 192\)"):
+        _export_transformers_checkpoint(model)
+
+    for linear, original_weight in zip(model, original_weights):
+        assert linear.weight.dtype == torch.bfloat16
+        assert torch.equal(linear.weight, original_weight)
+
+
 @pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"])
 def test_export_iq_rejects_an_unhandled_search_impl(num_bits):
     linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)

From 802682bc91bd9c269781a2c62982f30af96a6a43 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 05:00:09 -0700
Subject: [PATCH 26/32] Tighten IQ export diagnostics

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py        |  2 --
 modelopt/torch/export/unified_export_megatron.py  |  2 +-
 .../torch/export/test_unified_export_megatron.py  | 15 +++++++++++++++
 tests/unit/torch/export/test_export_weight.py     |  2 ++
 4 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index fa12577b731..eb7f8d0854c 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -1688,8 +1688,6 @@ def export_hf_checkpoint(
         )
         return
 
-    _validate_iq_export_weight_shapes(model)
-
     is_fsdp2_sharded = (
         torch.distributed.is_available()
         and torch.distributed.is_initialized()
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index be5c067e04c..f7b5d7ce7a8 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -1516,7 +1516,7 @@ def _grouped_mlp_slicing(
                         temp_amax_wqs.append(_wq)
 
                 name_to_value, qformat, block_size = self._get_quantized_state(
-                    module, self.dtype, prefix=prefix
+                    module, self.dtype, prefix=expert_prefix
                 )
                 weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat)
                 name_to_value.pop("weight", None)
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 b6da351e6d8..bb06fd7826f 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -981,6 +981,21 @@ def test_grouped_mlp_slicing_maps_local_to_global_expert_ids():
     assert "experts.1.gate_up_proj.weight" not in exporter._state_dict
 
 
+def test_grouped_mlp_slicing_formats_quantized_state_prefix_per_expert():
+    exporter = _make_exporter_for_grouped_mlp()
+    prefixes = []
+    exporter._get_quantized_state = lambda *a, **k: (
+        prefixes.append(k["prefix"]) or {},
+        None,
+        0,
+    )
+    module = _FakeTEGroupedMLP(num_gemms=2, local_expert_indices=[4, 5])
+
+    exporter._grouped_mlp_slicing(module, "experts.{}.gate_up_proj")
+
+    assert prefixes == ["experts.4.gate_up_proj.", "experts.5.gate_up_proj."]
+
+
 def test_grouped_mlp_slicing_normalizes_tensor_local_expert_indices():
     """local_expert_indices may arrive as a torch.Tensor (Megatron path). It must be
     normalized to list[int] -- a naive `bool(tensor)` on a multi-element tensor raises.
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index a4ceb88b8b4..45820378424 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -299,6 +299,7 @@ def __init__(self):
 
     model = nn.Module()
     model.block = CustomWeightModule()
+    model.config = SimpleNamespace(torch_dtype=torch.bfloat16)
     original_weight = model.block.proj.detach().clone()
 
     with pytest.raises(ValueError, match=r"block\.proj: nonstandard weight"):
@@ -325,6 +326,7 @@ def test_export_hf_checkpoint_runs_iq_shape_preflight_before_mutation(tmp_path):
             backend="ggml",
         )
     )
+    linear.config = SimpleNamespace(torch_dtype=torch.bfloat16)
     original_weight = linear.weight.detach().clone()
 
     with pytest.raises(ValueError, match=r"weight: \(4, 192\)"):

From 2f11f1e44849c03b801509464f93d87272854085 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 05:49:26 -0700
Subject: [PATCH 27/32] Preflight IQ weights in offload export

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 .../export/unified_export_hf_streaming.py     |  2 ++
 tests/unit/torch/export/test_export_weight.py | 25 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py
index 32477173e51..7e11d85a853 100644
--- a/modelopt/torch/export/unified_export_hf_streaming.py
+++ b/modelopt/torch/export/unified_export_hf_streaming.py
@@ -48,6 +48,7 @@
 from .quant_utils import (
     _get_kv_cache_postprocess_config,
     _postprocess_single_tensor,
+    _validate_iq_export_weight_shapes,
     get_quant_config,
 )
 from .registry import ExportContext
@@ -366,6 +367,7 @@ def _export_transformers_checkpoint_streaming(
 
     # --- Same model-level setup as _export_transformers_checkpoint ---
     dtype = _resolve_export_dtype(model, dtype)
+    _validate_iq_export_weight_shapes(model)
     _prepare_moe_inputs(model, dtype, is_modelopt_qlora)
 
     requantize_resmooth_fused_llm_layers(model)
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 45820378424..fe72d35dc2b 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -34,6 +34,9 @@
     _process_quantized_modules,
     export_hf_checkpoint,
 )
+from modelopt.torch.export.unified_export_hf_streaming import (
+    _export_transformers_checkpoint_streaming,
+)
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.nn import GroupedQuantizer, TensorQuantizer
 from modelopt.torch.quantization.utils import quantizer_attr_names
@@ -360,6 +363,28 @@ def test_export_transformers_checkpoint_runs_iq_shape_preflight_before_mutation(
         assert torch.equal(linear.weight, original_weight)
 
 
+def test_export_transformers_checkpoint_streaming_runs_iq_shape_preflight_before_mutation(
+    tmp_path,
+):
+    linear = nn.Linear(192, 4, bias=False, dtype=torch.bfloat16)
+    linear.config = SimpleNamespace(torch_dtype=torch.bfloat16)
+    linear.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+        )
+    )
+    original_weight = linear.weight.detach().clone()
+
+    with pytest.raises(ValueError, match=r"weight: \(4, 192\)"):
+        _export_transformers_checkpoint_streaming(linear, export_dir=tmp_path)
+
+    assert linear.weight.dtype == torch.bfloat16
+    assert torch.equal(linear.weight, original_weight)
+    assert not any(tmp_path.iterdir())
+
+
 @pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"])
 def test_export_iq_rejects_an_unhandled_search_impl(num_bits):
     linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)

From fe064eba39e42421797476e20ad802f8c06fea58 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 06:10:12 -0700
Subject: [PATCH 28/32] Preserve export labels and fast paths

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/quant_utils.py                 |  5 +++++
 modelopt/torch/export/unified_export_megatron.py     | 12 ++++++++----
 .../torch/export/test_unified_export_megatron.py     | 11 +++++++----
 tests/unit/torch/export/test_export_weight.py        | 11 +++++++++++
 4 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index ab462b80e59..818c68570f1 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -154,6 +154,11 @@ def inspect_weight(
 
 def _validate_iq_export_weight_shapes(model: nn.Module) -> None:
     """Reject all unsupported IQ weights before export mutates the model."""
+    if not any(
+        isinstance(module, TensorQuantizer) and module.is_enabled and module.num_bits in IQ_FORMATS
+        for module in model.modules()
+    ):
+        return
     incompatible = _iq_export_weight_shape_errors(model)
 
     if incompatible:
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index f7b5d7ce7a8..6300e574476 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -1144,6 +1144,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.
 
@@ -1151,6 +1152,7 @@ def _get_quantized_state(
             module: The target module to perform real quantization.
             dtype: The default data type.
             prefix: The prefix of the layer.
+            describe_as: Optional logical weight label used only in IQ diagnostics.
 
         Returns:
             Tuple: state_dict, quantization format, and block_size of the module.
@@ -1163,9 +1165,8 @@ def _get_quantized_state(
 
         is_iq = qformat in IQ_FORMATS
         if is_iq:
-            _validate_iq_quantizer_config(
-                module, qformat, describe_as=f"{prefix}weight" if prefix else "weight"
-            )
+            weight_label = describe_as or (f"{prefix}weight" if prefix else "weight")
+            _validate_iq_quantizer_config(module, qformat, describe_as=weight_label)
         name_to_value = self._get_weight_bias(
             module, dtype, name_to_value, keep_weight_device=is_iq
         )
@@ -1516,7 +1517,10 @@ def _grouped_mlp_slicing(
                         temp_amax_wqs.append(_wq)
 
                 name_to_value, qformat, block_size = self._get_quantized_state(
-                    module, self.dtype, prefix=expert_prefix
+                    module,
+                    self.dtype,
+                    prefix=prefix,
+                    describe_as=expert_prefix + "weight",
                 )
                 weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat)
                 name_to_value.pop("weight", None)
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 bb06fd7826f..dc9de60b3df 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -981,11 +981,11 @@ def test_grouped_mlp_slicing_maps_local_to_global_expert_ids():
     assert "experts.1.gate_up_proj.weight" not in exporter._state_dict
 
 
-def test_grouped_mlp_slicing_formats_quantized_state_prefix_per_expert():
+def test_grouped_mlp_slicing_preserves_exclusion_prefix_and_formats_iq_label():
     exporter = _make_exporter_for_grouped_mlp()
-    prefixes = []
+    calls = []
     exporter._get_quantized_state = lambda *a, **k: (
-        prefixes.append(k["prefix"]) or {},
+        calls.append((k["prefix"], k["describe_as"])) or {},
         None,
         0,
     )
@@ -993,7 +993,10 @@ def test_grouped_mlp_slicing_formats_quantized_state_prefix_per_expert():
 
     exporter._grouped_mlp_slicing(module, "experts.{}.gate_up_proj")
 
-    assert prefixes == ["experts.4.gate_up_proj.", "experts.5.gate_up_proj."]
+    assert calls == [
+        ("experts.{}.gate_up_proj", "experts.4.gate_up_proj.weight"),
+        ("experts.{}.gate_up_proj", "experts.5.gate_up_proj.weight"),
+    ]
 
 
 def test_grouped_mlp_slicing_normalizes_tensor_local_expert_indices():
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index fe72d35dc2b..40d9ca923cc 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -235,6 +235,17 @@ def test_export_iq_shape_preflight_reports_every_incompatible_weight():
     assert "1.weight" not in message
 
 
+def test_export_iq_shape_preflight_skips_detailed_scan_without_iq(monkeypatch):
+    model = nn.Sequential(nn.Linear(32, 16), nn.ReLU())
+    monkeypatch.setattr(
+        quant_utils,
+        "_iq_export_weight_shape_errors",
+        lambda _model: pytest.fail("non-IQ export must skip the detailed weight scan"),
+    )
+
+    _validate_iq_export_weight_shapes(model)
+
+
 def test_export_iq_shape_preflight_reports_grouped_weights():
     class GroupedLinear(nn.Module):
         def __init__(self):

From 1862b921a36f37d9eb9d1f18248167e624ae2bdf Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 07:31:29 -0700
Subject: [PATCH 29/32] Preserve grouped fakequant export compatibility

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 .../export/plugins/vllm_fakequant_megatron.py     |  3 +++
 .../torch/export/test_unified_export_megatron.py  | 15 +++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py
index 96355cc30d5..858b8bfee09 100644
--- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py
+++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py
@@ -116,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.
 
@@ -126,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.
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 dc9de60b3df..11ebfd86108 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -999,6 +999,21 @@ def test_grouped_mlp_slicing_preserves_exclusion_prefix_and_formats_iq_label():
     ]
 
 
+def test_vllm_fakequant_grouped_mlp_accepts_diagnostic_label():
+    exporter = object.__new__(VllmFqGPTModelExporter)
+    exporter.dtype = torch.bfloat16
+    exporter._state_dict = {}
+    exporter.exclude_modules = []
+    exporter._get_weight_bias = lambda module, _dtype, _values: {"weight": module.weight}
+    exporter._get_weight_scales = lambda *_args: (None, None)
+    exporter._record_layer_quant_config = lambda *_args: None
+    module = _FakeTEGroupedMLP(num_gemms=1, local_expert_indices=[3])
+
+    exporter._grouped_mlp_slicing(module, "experts.{}.down_proj")
+
+    assert "experts.3.down_proj.weight" in exporter._state_dict
+
+
 def test_grouped_mlp_slicing_normalizes_tensor_local_expert_indices():
     """local_expert_indices may arrive as a torch.Tensor (Megatron path). It must be
     normalized to list[int] -- a naive `bool(tensor)` on a multi-element tensor raises.

From 84df0eece1b412e271c1615014e41dbb509b9ac9 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 07:52:24 -0700
Subject: [PATCH 30/32] Harden distributed IQ export preflight

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/convert_hf_config.py    |   2 +
 modelopt/torch/export/quant_utils.py          | 114 +++++++++++-------
 .../torch/export/unified_export_megatron.py   |  56 +++++++--
 .../export/test_unified_export_megatron.py    |  64 ++++++++--
 tests/unit/torch/export/test_export_weight.py |  25 ++++
 5 files changed, 200 insertions(+), 61 deletions(-)

diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py
index 2806c2836b0..0dff17f660e 100644
--- a/modelopt/torch/export/convert_hf_config.py
+++ b/modelopt/torch/export/convert_hf_config.py
@@ -121,6 +121,8 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None)
         }
     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,
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 818c68570f1..9721a2b4b34 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -87,60 +87,75 @@
 logger = logging.getLogger(__name__)
 
 
+def _iq_export_weight_quantizer(
+    module: nn.Module, weight_name: str
+) -> tuple[nn.Module | None, bool]:
+    """Return a representative quantizer and whether ``weight_name`` uses a plural layout."""
+    quantizer = representative_weight_quantizer(module, weight_name)
+    quantizer_name = quantizer_attr_names(weight_name).weight_quantizer
+    plural_quantizers = getattr(module, quantizer_name + "s", None)
+    if quantizer is not None:
+        return quantizer, isinstance(plural_quantizers, nn.ModuleList)
+
+    if weight_name == getattr(module, "_first_proj_attr", None) and weight_name != "gate_up_proj":
+        # Older non-gated fused-expert wrappers kept the first-projection quantizers under
+        # the gate_up_proj sentinel name rather than the projection's actual attribute name.
+        plural_quantizers = getattr(module, "gate_up_proj_weight_quantizers", None)
+        if isinstance(plural_quantizers, nn.ModuleList) and len(plural_quantizers) > 0:
+            return plural_quantizers[0], True
+    return None, False
+
+
+def _iq_export_weight_shape_error(
+    module: nn.Module,
+    module_name: str,
+    weight_name: str,
+    quantizer: nn.Module | None,
+    *,
+    allow_nonstandard_name: bool = False,
+) -> str | None:
+    """Return one IQ export compatibility error, if any, for ``weight_name``."""
+    if not (
+        isinstance(quantizer, TensorQuantizer)
+        and quantizer.is_enabled
+        and quantizer.num_bits in IQ_FORMATS
+    ):
+        return None
+
+    qualified_name = f"{module_name}.{weight_name}".lstrip(".")
+    if weight_name != "weight" and not allow_nonstandard_name:
+        return f"{qualified_name}: nonstandard weight attributes are not supported"
+
+    weight = getattr(module, weight_name, None)
+    if not isinstance(weight, torch.Tensor):
+        return None
+    group_size = iq_format_spec(quantizer.num_bits)["group_size"]
+    if weight.dim() == 0 or weight.shape[-1] % group_size:
+        return f"{qualified_name}: {tuple(weight.shape)}"
+    return None
+
+
 def _iq_export_weight_shape_errors(model: nn.Module) -> list[str]:
     """Return selected IQ weights that cannot be handled by unified export."""
     incompatible: list[str] = []
     for module_name, module in model.named_modules():
         inspected: set[str] = set()
-
-        def uses_plural_quantizers(weight_name: str) -> bool:
-            quantizer_name = quantizer_attr_names(weight_name).weight_quantizer
-            if isinstance(getattr(module, quantizer_name + "s", None), nn.ModuleList):
-                return True
-            return (
-                weight_name == getattr(module, "_first_proj_attr", None)
-                and weight_name != "gate_up_proj"
-                and isinstance(
-                    getattr(module, "gate_up_proj_weight_quantizers", None), nn.ModuleList
-                )
-            )
-
-        def inspect_weight(
-            weight_name: str,
-            quantizer: nn.Module | None,
-            *,
-            allow_nonstandard_name: bool = False,
-        ) -> None:
+        for weight_name in weight_attr_names(module):
             qualified_name = f"{module_name}.{weight_name}".lstrip(".")
             if qualified_name in inspected:
-                return
+                continue
             inspected.add(qualified_name)
-            if not (
-                isinstance(quantizer, TensorQuantizer)
-                and quantizer.is_enabled
-                and quantizer.num_bits in IQ_FORMATS
-            ):
-                return
-            if weight_name != "weight" and not allow_nonstandard_name:
-                incompatible.append(
-                    f"{qualified_name}: nonstandard weight attributes are not supported"
-                )
-                return
-            weight = getattr(module, weight_name, None)
-            if not isinstance(weight, torch.Tensor):
-                return
-            group_size = iq_format_spec(quantizer.num_bits)["group_size"]
-            if weight.dim() == 0 or weight.shape[-1] % group_size:
-                incompatible.append(f"{qualified_name}: {tuple(weight.shape)}")
-
-        for weight_name in weight_attr_names(module):
-            quantizer = representative_weight_quantizer(module, weight_name)
+            quantizer, uses_plural_quantizers = _iq_export_weight_quantizer(module, weight_name)
             # Fused-expert plural quantizer lists are rebuilt as standard per-expert modules.
-            inspect_weight(
+            error = _iq_export_weight_shape_error(
+                module,
+                module_name,
                 weight_name,
                 quantizer,
-                allow_nonstandard_name=uses_plural_quantizers(weight_name),
+                allow_nonstandard_name=uses_plural_quantizers,
             )
+            if error is not None:
+                incompatible.append(error)
 
         grouped_quantizer = getattr(module, "weight_quantizer", None)
         if isinstance(grouped_quantizer, GroupedQuantizer) and len(grouped_quantizer) > 0:
@@ -148,7 +163,20 @@ def inspect_weight(
             for index in range(max(num_gemms, len(grouped_quantizer))):
                 quantizer = grouped_quantizer[min(index, len(grouped_quantizer) - 1)]
                 # Grouped linears are rebuilt as standard per-expert modules before packing.
-                inspect_weight(f"weight{index}", quantizer, allow_nonstandard_name=True)
+                weight_name = f"weight{index}"
+                qualified_name = f"{module_name}.{weight_name}".lstrip(".")
+                if qualified_name in inspected:
+                    continue
+                inspected.add(qualified_name)
+                error = _iq_export_weight_shape_error(
+                    module,
+                    module_name,
+                    weight_name,
+                    quantizer,
+                    allow_nonstandard_name=True,
+                )
+                if error is not None:
+                    incompatible.append(error)
     return incompatible
 
 
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 6300e574476..0302bcc709e 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -309,21 +309,46 @@ def _model_has_iq_quantizer(model: torch.nn.Module) -> bool:
             for module in model.modules()
         )
 
+    def _model_has_packed_expert_iq_quantizer(self) -> bool:
+        """Return whether this rank would route an IQ expert through a packed-expert rule."""
+        packed_projection_names = {
+            rule_name.rsplit(".", 1)[-1]
+            for rule_name, rule in getattr(self, "rules", {}).items()
+            if getattr(rule, "_modelopt_export_func_name", None)
+            in {"pack_name_remapping", "pack_name_remapping_gpt_oss"}
+        }
+        if not packed_projection_names:
+            return False
+
+        for module_name, module in self.model.named_modules():
+            path = module_name.split(".")
+            if (
+                "local_experts" in path
+                and path[-1] in packed_projection_names
+                and self._model_has_iq_quantizer(module)
+            ):
+                return True
+        return False
+
     @staticmethod
-    def _collective_iq_export_flags(local_iq: bool, local_bad_shape: bool) -> tuple[bool, bool]:
-        """Combine IQ presence and incompatible-shape flags across all ranks."""
+    def _collective_iq_export_flags(
+        local_iq: bool, local_bad_shape: bool, local_packed_expert_iq: bool
+    ) -> tuple[bool, bool, bool]:
+        """Combine IQ presence, incompatible-shape, and packed-expert flags across ranks."""
         if not torch.distributed.is_initialized():
-            return local_iq, local_bad_shape
+            return local_iq, local_bad_shape, local_packed_expert_iq
         device = (
             torch.device("cuda", torch.cuda.current_device())
             if torch.distributed.get_backend() == torch.distributed.Backend.NCCL
             else torch.device("cpu")
         )
         flags = torch.tensor(
-            [int(local_iq), int(local_bad_shape)], dtype=torch.int32, device=device
+            [int(local_iq), int(local_bad_shape), int(local_packed_expert_iq)],
+            dtype=torch.int32,
+            device=device,
         )
         torch.distributed.all_reduce(flags, op=torch.distributed.ReduceOp.MAX)
-        return bool(flags[0].item()), bool(flags[1].item())
+        return bool(flags[0].item()), bool(flags[1].item()), bool(flags[2].item())
 
     def _validate_iq_export(self) -> None:
         """Reject unsupported IQ export configurations in one collective preflight."""
@@ -335,7 +360,10 @@ def _validate_iq_export(self) -> None:
         local_errors = (
             _iq_export_weight_shape_errors(self.model) if local_iq and tp_size == 1 else []
         )
-        any_iq, any_bad_shape = self._collective_iq_export_flags(local_iq, bool(local_errors))
+        local_packed_expert_iq = local_iq and self._model_has_packed_expert_iq_quantizer()
+        any_iq, any_bad_shape, any_packed_expert_iq = self._collective_iq_export_flags(
+            local_iq, bool(local_errors), local_packed_expert_iq
+        )
         if not any_iq:
             return
         if tp_size != 1:
@@ -343,6 +371,11 @@ def _validate_iq_export(self) -> None:
                 "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model "
                 "parallel size 1"
             )
+        if any_packed_expert_iq:
+            raise NotImplementedError(
+                "Megatron IQ1_S/IQ2_XS export does not yet support packed expert weights "
+                "because the HF transpose changes the 256-value block axis"
+            )
         if not any_bad_shape:
             return
         if local_errors:
@@ -1084,9 +1117,14 @@ def _custom_mapping_to_lambda(mapping):
             func = method_map[mapping.func_name]
             prefix = mapping.target_name_or_prefix
             func_kwargs = mapping.func_kwargs
-            return lambda m, *args, **kwargs: func(
-                m, prefix.format(*args), **{**func_kwargs, **kwargs}
-            )
+
+            def apply_mapping(m, *args, **kwargs):
+                return func(m, prefix.format(*args), **{**func_kwargs, **kwargs})
+
+            # Preflight uses this marker to identify packed-expert rules before any rank enters
+            # per-layer conversion, where a local exception would strand other pipeline ranks.
+            setattr(apply_mapping, "_modelopt_export_func_name", mapping.func_name)
+            return apply_mapping
 
         for arch, mappings in all_mcore_hf_export_mapping.items():
             all_rules[arch] = {
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 11ebfd86108..9fc1a83159a 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -160,16 +160,20 @@ def test_megatron_iq_export_rejects_tensor_parallelism():
         exporter.save_pretrained("unused", "unused")
 
 
-def test_megatron_iq_export_guard_uses_one_collective_for_both_flags():
+def test_megatron_iq_export_guard_uses_one_collective_for_all_flags():
     def report_remote_iq(flags, **_kwargs):
-        flags.copy_(torch.tensor([1, 0], dtype=flags.dtype, device=flags.device))
+        flags.copy_(torch.tensor([1, 0, 1], dtype=flags.dtype, device=flags.device))
 
     with (
         patch.object(torch.distributed, "is_initialized", return_value=True),
         patch.object(torch.distributed, "get_backend", return_value=torch.distributed.Backend.GLOO),
         patch.object(torch.distributed, "all_reduce", side_effect=report_remote_iq),
     ):
-        assert GPTModelExporter._collective_iq_export_flags(False, False) == (True, False)
+        assert GPTModelExporter._collective_iq_export_flags(False, False, False) == (
+            True,
+            False,
+            True,
+        )
 
 
 def test_megatron_iq_tensor_parallel_guard_finds_later_mixed_format():
@@ -187,6 +191,30 @@ def test_megatron_iq_tensor_parallel_guard_finds_later_mixed_format():
     assert GPTModelExporter._model_has_iq_quantizer(torch.nn.Sequential(fp8, iq))
 
 
+def test_megatron_iq_packed_expert_guard_finds_selected_projection():
+    expert = torch.nn.Module()
+    expert.linear_fc1 = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16)
+    expert.linear_fc1.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+        )
+    )
+    model = torch.nn.Module()
+    model.local_experts = torch.nn.ModuleList([expert])
+
+    def packed_rule(*_args, **_kwargs):
+        pass
+
+    setattr(packed_rule, "_modelopt_export_func_name", "pack_name_remapping")
+    exporter = object.__new__(GPTModelExporter)
+    exporter.model = model
+    exporter.rules = {"local_experts.linear_fc1": packed_rule}
+
+    assert exporter._model_has_packed_expert_iq_quantizer()
+
+
 def test_megatron_iq_tensor_parallel_guard_covers_extra_modules():
     exporter = object.__new__(GPTModelExporter)
     exporter.model = torch.nn.Module()
@@ -194,7 +222,7 @@ def test_megatron_iq_tensor_parallel_guard_covers_extra_modules():
     with (
         patch.object(uem, "get_tensor_model_parallel_world_size", return_value=2),
         patch.object(exporter, "_model_has_iq_quantizer", return_value=True),
-        patch.object(exporter, "_collective_iq_export_flags", return_value=(True, False)),
+        patch.object(exporter, "_collective_iq_export_flags", return_value=(True, False, False)),
         pytest.raises(NotImplementedError, match="tensor model parallel size 1"),
     ):
         exporter.save_pretrained_extra_modules("unused")
@@ -221,13 +249,13 @@ def test_megatron_non_iq_export_skips_shape_scan():
         patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1),
         patch.object(exporter, "_model_has_iq_quantizer", return_value=False),
         patch.object(
-            exporter, "_collective_iq_export_flags", return_value=(False, False)
+            exporter, "_collective_iq_export_flags", return_value=(False, False, False)
         ) as collective,
         patch.object(uem, "_iq_export_weight_shape_errors") as shape_errors,
     ):
         exporter._validate_iq_export()
 
-    collective.assert_called_once_with(False, False)
+    collective.assert_called_once_with(False, False, False)
     shape_errors.assert_not_called()
 
 
@@ -245,13 +273,13 @@ def test_megatron_iq_shape_guard_fails_all_ranks_before_packing():
     with (
         patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1),
         patch.object(
-            exporter, "_collective_iq_export_flags", return_value=(True, True)
+            exporter, "_collective_iq_export_flags", return_value=(True, True, False)
         ) as collective,
         pytest.raises(ValueError, match=r"weight: \(2, 192\)"),
     ):
         exporter._validate_iq_export()
 
-    collective.assert_called_once_with(True, True)
+    collective.assert_called_once_with(True, True, False)
 
 
 def test_megatron_iq_shape_guard_reports_another_rank():
@@ -260,12 +288,30 @@ def test_megatron_iq_shape_guard_reports_another_rank():
 
     with (
         patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1),
-        patch.object(exporter, "_collective_iq_export_flags", return_value=(True, True)),
+        patch.object(exporter, "_collective_iq_export_flags", return_value=(True, True, False)),
         pytest.raises(ValueError, match="Another distributed rank"),
     ):
         exporter._validate_iq_export()
 
 
+def test_megatron_iq_packed_expert_guard_fails_all_ranks_before_conversion():
+    exporter = object.__new__(GPTModelExporter)
+    exporter.model = torch.nn.Module()
+
+    with (
+        patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1),
+        patch.object(exporter, "_model_has_iq_quantizer", return_value=False),
+        patch.object(exporter, "_model_has_packed_expert_iq_quantizer", return_value=False),
+        patch.object(
+            exporter, "_collective_iq_export_flags", return_value=(True, False, True)
+        ) as collective,
+        pytest.raises(NotImplementedError, match="packed expert weights"),
+    ):
+        exporter._validate_iq_export()
+
+    collective.assert_called_once_with(False, False, False)
+
+
 def test_megatron_iq_packer_rejects_tensor_parallelism():
     """Internal packing keeps the TP=1 guard for extra-module export paths."""
     module = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16)
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 40d9ca923cc..6673d51b34d 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -298,6 +298,31 @@ def __init__(self):
         _validate_iq_export_weight_shapes(model)
 
 
+def test_export_iq_shape_preflight_reports_legacy_fused_expert_quantizers():
+    class LegacyFusedExperts(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self._first_proj_attr = "up_proj"
+            self.up_proj = nn.Parameter(torch.ones(4, 192, dtype=torch.bfloat16))
+            self.gate_up_proj_weight_quantizers = nn.ModuleList(
+                [
+                    TensorQuantizer(
+                        QuantizerAttributeConfig(
+                            num_bits="iq2_xs",
+                            block_sizes={-1: 256},
+                            backend="ggml",
+                        )
+                    )
+                ]
+            )
+
+    model = nn.Module()
+    model.experts = LegacyFusedExperts()
+
+    with pytest.raises(ValueError, match=r"experts\.up_proj: \(4, 192\)"):
+        _validate_iq_export_weight_shapes(model)
+
+
 def test_export_iq_preflight_rejects_nonstandard_weight_before_mutation(tmp_path):
     class CustomWeightModule(nn.Module):
         def __init__(self):

From d9be3bb045f168e754c1109afc740a6948c719f3 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 08:01:57 -0700
Subject: [PATCH 31/32] Collect all IQ export preflight failures

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/quant_utils.py          | 153 +++++++++++-------
 .../torch/export/unified_export_megatron.py   |  31 ++--
 .../export/test_unified_export_megatron.py    |  25 +++
 tests/unit/torch/export/test_export_weight.py |  21 +++
 4 files changed, 161 insertions(+), 69 deletions(-)

diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 9721a2b4b34..86d65cd062d 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -106,22 +106,53 @@ def _iq_export_weight_quantizer(
     return None, False
 
 
+def _iq_export_weights(
+    model: nn.Module,
+) -> Generator[tuple[nn.Module, str, str, TensorQuantizer, bool], None, None]:
+    """Yield each selected IQ weight and the metadata needed by export preflight."""
+    for module_name, module in model.named_modules():
+        inspected: set[str] = set()
+        for weight_name in weight_attr_names(module):
+            qualified_name = f"{module_name}.{weight_name}".lstrip(".")
+            if qualified_name in inspected:
+                continue
+            inspected.add(qualified_name)
+            quantizer, uses_plural_quantizers = _iq_export_weight_quantizer(module, weight_name)
+            if (
+                isinstance(quantizer, TensorQuantizer)
+                and quantizer.is_enabled
+                and quantizer.num_bits in IQ_FORMATS
+            ):
+                yield module, module_name, weight_name, quantizer, uses_plural_quantizers
+
+        grouped_quantizer = getattr(module, "weight_quantizer", None)
+        if isinstance(grouped_quantizer, GroupedQuantizer) and len(grouped_quantizer) > 0:
+            num_gemms = int(getattr(module, "num_gemms", len(grouped_quantizer)))
+            for index in range(max(num_gemms, len(grouped_quantizer))):
+                quantizer = grouped_quantizer[min(index, len(grouped_quantizer) - 1)]
+                weight_name = f"weight{index}"
+                qualified_name = f"{module_name}.{weight_name}".lstrip(".")
+                if qualified_name in inspected:
+                    continue
+                inspected.add(qualified_name)
+                if (
+                    isinstance(quantizer, TensorQuantizer)
+                    and quantizer.is_enabled
+                    and quantizer.num_bits in IQ_FORMATS
+                ):
+                    # Grouped linears are rebuilt as standard per-expert modules before packing.
+                    yield module, module_name, weight_name, quantizer, True
+
+
 def _iq_export_weight_shape_error(
     module: nn.Module,
     module_name: str,
     weight_name: str,
-    quantizer: nn.Module | None,
+    quantizer: TensorQuantizer,
     *,
     allow_nonstandard_name: bool = False,
 ) -> str | None:
     """Return one IQ export compatibility error, if any, for ``weight_name``."""
-    if not (
-        isinstance(quantizer, TensorQuantizer)
-        and quantizer.is_enabled
-        and quantizer.num_bits in IQ_FORMATS
-    ):
-        return None
-
     qualified_name = f"{module_name}.{weight_name}".lstrip(".")
     if weight_name != "weight" and not allow_nonstandard_name:
         return f"{qualified_name}: nonstandard weight attributes are not supported"
@@ -138,45 +169,19 @@ def _iq_export_weight_shape_error(
 def _iq_export_weight_shape_errors(model: nn.Module) -> list[str]:
     """Return selected IQ weights that cannot be handled by unified export."""
     incompatible: list[str] = []
-    for module_name, module in model.named_modules():
-        inspected: set[str] = set()
-        for weight_name in weight_attr_names(module):
-            qualified_name = f"{module_name}.{weight_name}".lstrip(".")
-            if qualified_name in inspected:
-                continue
-            inspected.add(qualified_name)
-            quantizer, uses_plural_quantizers = _iq_export_weight_quantizer(module, weight_name)
-            # Fused-expert plural quantizer lists are rebuilt as standard per-expert modules.
-            error = _iq_export_weight_shape_error(
-                module,
-                module_name,
-                weight_name,
-                quantizer,
-                allow_nonstandard_name=uses_plural_quantizers,
-            )
-            if error is not None:
-                incompatible.append(error)
-
-        grouped_quantizer = getattr(module, "weight_quantizer", None)
-        if isinstance(grouped_quantizer, GroupedQuantizer) and len(grouped_quantizer) > 0:
-            num_gemms = int(getattr(module, "num_gemms", len(grouped_quantizer)))
-            for index in range(max(num_gemms, len(grouped_quantizer))):
-                quantizer = grouped_quantizer[min(index, len(grouped_quantizer) - 1)]
-                # Grouped linears are rebuilt as standard per-expert modules before packing.
-                weight_name = f"weight{index}"
-                qualified_name = f"{module_name}.{weight_name}".lstrip(".")
-                if qualified_name in inspected:
-                    continue
-                inspected.add(qualified_name)
-                error = _iq_export_weight_shape_error(
-                    module,
-                    module_name,
-                    weight_name,
-                    quantizer,
-                    allow_nonstandard_name=True,
-                )
-                if error is not None:
-                    incompatible.append(error)
+    for module, module_name, weight_name, quantizer, uses_plural_quantizers in _iq_export_weights(
+        model
+    ):
+        # Fused-expert plural and grouped quantizers are rebuilt as standard per-expert modules.
+        error = _iq_export_weight_shape_error(
+            module,
+            module_name,
+            weight_name,
+            quantizer,
+            allow_nonstandard_name=uses_plural_quantizers,
+        )
+        if error is not None:
+            incompatible.append(error)
     return incompatible
 
 
@@ -198,27 +203,35 @@ def _validate_iq_export_weight_shapes(model: nn.Module) -> None:
         )
 
 
-def _validate_iq_quantizer_config(
+def _iq_quantizer_config_error(
     module: nn.Module,
     quantization_format: str,
     weight_name: str = "weight",
     describe_as: str | None = None,
-) -> None:
-    """Reject IQ export settings that the checkpoint encoder cannot reproduce."""
+    quantizer: nn.Module | None = None,
+) -> Exception | None:
+    """Return an IQ export configuration error without mutating the module."""
     where = f" for '{describe_as}'" if describe_as else ""
-    quantizer = representative_weight_quantizer(module, weight_name)
+    if quantizer is None:
+        quantizer = representative_weight_quantizer(module, weight_name)
     if not isinstance(quantizer, TensorQuantizer):
-        raise ValueError(
+        return ValueError(
             f"{quantization_format.upper()} export requires one TensorQuantizer{where}"
         )
     if quantizer.num_bits != quantization_format or quantizer.backend != "ggml":
-        raise ValueError(
+        return ValueError(
             f"{quantization_format.upper()} export requires the matching ggml quantizer{where}"
         )
     extra_args = quantizer.backend_extra_args or {}
+    unknown_keys = set(extra_args) - {"search_impl"}
+    if unknown_keys:
+        return ValueError(
+            f"{quantization_format.upper()} export received unsupported backend_extra_args keys "
+            f"{sorted(unknown_keys)}{where}"
+        )
     search_impl = extra_args.get("search_impl", "auto")
     if search_impl != "auto":
-        raise NotImplementedError(
+        return NotImplementedError(
             f"{quantization_format.upper()} export only supports search_impl='auto'{where}"
         )
     quantizer_names = quantizer_attr_names(weight_name)
@@ -228,11 +241,41 @@ def _validate_iq_quantizer_config(
             getattr(activation_quantizer, "is_enabled", False)
             or getattr(activation_quantizer, "pre_quant_scale", None) is not None
         ):
-            raise NotImplementedError(
+            return NotImplementedError(
                 f"{quantization_format.upper()} export is weight-only{where}; activation "
                 "quantization "
                 "and pre_quant_scale are not represented in the checkpoint"
             )
+    return None
+
+
+def _validate_iq_quantizer_config(
+    module: nn.Module,
+    quantization_format: str,
+    weight_name: str = "weight",
+    describe_as: str | None = None,
+) -> None:
+    """Reject IQ export settings that the checkpoint encoder cannot reproduce."""
+    error = _iq_quantizer_config_error(module, quantization_format, weight_name, describe_as)
+    if error is not None:
+        raise error
+
+
+def _iq_export_quantizer_config_errors(model: nn.Module) -> list[str]:
+    """Return selected IQ quantizer configurations unsupported by checkpoint export."""
+    incompatible: list[str] = []
+    for module, module_name, weight_name, quantizer, _ in _iq_export_weights(model):
+        qualified_name = f"{module_name}.{weight_name}".lstrip(".")
+        error = _iq_quantizer_config_error(
+            module,
+            quantizer.num_bits,
+            weight_name,
+            qualified_name,
+            quantizer=quantizer,
+        )
+        if error is not None:
+            incompatible.append(str(error))
+    return incompatible
 
 
 def _pack_iq_weight(
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 0302bcc709e..721c08cefaf 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -71,6 +71,7 @@
     iq_format_spec,
 )
 from .quant_utils import (
+    _iq_export_quantizer_config_errors,
     _iq_export_weight_shape_errors,
     _pack_iq_weight,
     _validate_iq_quantizer_config,
@@ -332,18 +333,18 @@ def _model_has_packed_expert_iq_quantizer(self) -> bool:
 
     @staticmethod
     def _collective_iq_export_flags(
-        local_iq: bool, local_bad_shape: bool, local_packed_expert_iq: bool
+        local_iq: bool, local_error: bool, local_packed_expert_iq: bool
     ) -> tuple[bool, bool, bool]:
-        """Combine IQ presence, incompatible-shape, and packed-expert flags across ranks."""
+        """Combine IQ presence, preflight-error, and packed-expert flags across ranks."""
         if not torch.distributed.is_initialized():
-            return local_iq, local_bad_shape, local_packed_expert_iq
+            return local_iq, local_error, local_packed_expert_iq
         device = (
             torch.device("cuda", torch.cuda.current_device())
             if torch.distributed.get_backend() == torch.distributed.Backend.NCCL
             else torch.device("cpu")
         )
         flags = torch.tensor(
-            [int(local_iq), int(local_bad_shape), int(local_packed_expert_iq)],
+            [int(local_iq), int(local_error), int(local_packed_expert_iq)],
             dtype=torch.int32,
             device=device,
         )
@@ -357,11 +358,12 @@ def _validate_iq_export(self) -> None:
 
         local_iq = self._model_has_iq_quantizer(self.model)
         tp_size = get_tensor_model_parallel_world_size()
-        local_errors = (
-            _iq_export_weight_shape_errors(self.model) if local_iq and tp_size == 1 else []
-        )
+        local_errors = []
+        if local_iq and tp_size == 1:
+            local_errors.extend(_iq_export_weight_shape_errors(self.model))
+            local_errors.extend(_iq_export_quantizer_config_errors(self.model))
         local_packed_expert_iq = local_iq and self._model_has_packed_expert_iq_quantizer()
-        any_iq, any_bad_shape, any_packed_expert_iq = self._collective_iq_export_flags(
+        any_iq, any_error, any_packed_expert_iq = self._collective_iq_export_flags(
             local_iq, bool(local_errors), local_packed_expert_iq
         )
         if not any_iq:
@@ -376,15 +378,15 @@ def _validate_iq_export(self) -> None:
                 "Megatron IQ1_S/IQ2_XS export does not yet support packed expert weights "
                 "because the HF transpose changes the 256-value block axis"
             )
-        if not any_bad_shape:
+        if not any_error:
             return
         if local_errors:
             details = "\n  - ".join(local_errors)
             raise ValueError(
-                "IQ export requires every selected weight's final dimension to be divisible by "
-                f"its format group size. Incompatible weights:\n  - {details}"
+                "IQ export rejected unsupported selected weights or quantizer settings:"
+                f"\n  - {details}"
             )
-        raise ValueError("Another distributed rank owns an IQ weight that cannot be packed")
+        raise ValueError("Another distributed rank owns an IQ weight that failed export preflight")
 
     def save_pretrained(
         self,
@@ -1215,8 +1217,9 @@ def _get_quantized_state(
         if qformat == QUANTIZATION_NONE:
             return name_to_value, qformat, block_size
         # IQ formats derive all block metadata directly from the weight and do not use amax or
-        # separately exported scaling tensors. Keep the weight on-device until its final HF layout
-        # has been produced, so the CUDA packer can be used.
+        # separately exported scaling tensors. Keep the weight on its current device until its
+        # final HF layout has been produced: this avoids a CPU round-trip in slicing paths and
+        # leaves this seam ready for a device-side packer.
         if is_iq:
             return name_to_value, qformat, block_size
         # Getting the weight scales
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 9fc1a83159a..a08372a4d15 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -294,6 +294,31 @@ def test_megatron_iq_shape_guard_reports_another_rank():
         exporter._validate_iq_export()
 
 
+def test_megatron_iq_config_guard_fails_collectively_before_conversion():
+    exporter = object.__new__(GPTModelExporter)
+    exporter.model = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16)
+    exporter.model.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+            backend_extra_args={"search_impl": "auto"},
+        )
+    )
+    exporter.model.input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8))
+
+    with (
+        patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1),
+        patch.object(
+            exporter, "_collective_iq_export_flags", return_value=(True, True, False)
+        ) as collective,
+        pytest.raises(ValueError, match="weight-only"),
+    ):
+        exporter._validate_iq_export()
+
+    collective.assert_called_once_with(True, True, False)
+
+
 def test_megatron_iq_packed_expert_guard_fails_all_ranks_before_conversion():
     exporter = object.__new__(GPTModelExporter)
     exporter.model = torch.nn.Module()
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 6673d51b34d..63f274b466a 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -467,6 +467,27 @@ def test_export_iq_rejects_non_weight_only_config(num_bits, quantizer_name, mode
         _export_quantized_weight(linear, torch.bfloat16)
 
 
+def test_export_iq_config_preflight_reports_weight_owner():
+    linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)
+    linear.weight_quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(
+            num_bits="iq2_xs",
+            block_sizes={-1: 256},
+            backend="ggml",
+            backend_extra_args={"search_impl": "auto"},
+        )
+    )
+    linear.input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8))
+    model = nn.Module()
+    model.block = linear
+
+    errors = quant_utils._iq_export_quantizer_config_errors(model)
+
+    assert len(errors) == 1
+    assert "block.weight" in errors[0]
+    assert "weight-only" in errors[0]
+
+
 def test_iq_config_error_identifies_weight():
     linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)
     linear.weight_quantizer = TensorQuantizer(

From ea612a87cdddb6891ec5097719380f1f68eb757c Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Wed, 16 Sep 2026 08:24:41 -0700
Subject: [PATCH 32/32] Preflight every IQ export entry point

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 modelopt/torch/export/layerwise_export.py     |  4 +-
 modelopt/torch/export/quant_utils.py          | 45 +++++++++----
 modelopt/torch/export/unified_export_hf.py    |  6 +-
 .../export/unified_export_hf_streaming.py     |  4 +-
 tests/unit/torch/export/test_export_weight.py | 66 +++++++++++++++++--
 5 files changed, 99 insertions(+), 26 deletions(-)

diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py
index ed35db30703..4aa14b6eced 100644
--- a/modelopt/torch/export/layerwise_export.py
+++ b/modelopt/torch/export/layerwise_export.py
@@ -42,7 +42,7 @@
 from .quant_utils import (
     _get_kv_cache_postprocess_config,
     _postprocess_single_tensor,
-    _validate_iq_export_weight_shapes,
+    _validate_iq_export_support,
     get_quant_config,
     get_quantization_format,
 )
@@ -138,7 +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_weight_shapes(model)
+    _validate_iq_export_support(model)
 
     tied = _tied_quantized_modules(model)
     if tied:
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 86d65cd062d..b3eec78480a 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -108,7 +108,7 @@ def _iq_export_weight_quantizer(
 
 def _iq_export_weights(
     model: nn.Module,
-) -> Generator[tuple[nn.Module, str, str, TensorQuantizer, bool], None, None]:
+) -> Generator[tuple[nn.Module, str, str, str, TensorQuantizer, bool], None, None]:
     """Yield each selected IQ weight and the metadata needed by export preflight."""
     for module_name, module in model.named_modules():
         inspected: set[str] = set()
@@ -123,7 +123,14 @@ def _iq_export_weights(
                 and quantizer.is_enabled
                 and quantizer.num_bits in IQ_FORMATS
             ):
-                yield module, module_name, weight_name, quantizer, uses_plural_quantizers
+                yield (
+                    module,
+                    module_name,
+                    weight_name,
+                    weight_name,
+                    quantizer,
+                    uses_plural_quantizers,
+                )
 
         grouped_quantizer = getattr(module, "weight_quantizer", None)
         if isinstance(grouped_quantizer, GroupedQuantizer) and len(grouped_quantizer) > 0:
@@ -141,7 +148,7 @@ def _iq_export_weights(
                     and quantizer.num_bits in IQ_FORMATS
                 ):
                     # Grouped linears are rebuilt as standard per-expert modules before packing.
-                    yield module, module_name, weight_name, quantizer, True
+                    yield module, module_name, weight_name, "weight", quantizer, True
 
 
 def _iq_export_weight_shape_error(
@@ -169,9 +176,14 @@ def _iq_export_weight_shape_error(
 def _iq_export_weight_shape_errors(model: nn.Module) -> list[str]:
     """Return selected IQ weights that cannot be handled by unified export."""
     incompatible: list[str] = []
-    for module, module_name, weight_name, quantizer, uses_plural_quantizers in _iq_export_weights(
-        model
-    ):
+    for (
+        module,
+        module_name,
+        weight_name,
+        _,
+        quantizer,
+        uses_plural_quantizers,
+    ) in _iq_export_weights(model):
         # Fused-expert plural and grouped quantizers are rebuilt as standard per-expert modules.
         error = _iq_export_weight_shape_error(
             module,
@@ -185,21 +197,21 @@ def _iq_export_weight_shape_errors(model: nn.Module) -> list[str]:
     return incompatible
 
 
-def _validate_iq_export_weight_shapes(model: nn.Module) -> None:
-    """Reject all unsupported IQ weights before export mutates the model."""
+def _validate_iq_export_support(model: nn.Module) -> None:
+    """Reject unsupported IQ weights and quantizer settings before export mutates the model."""
     if not any(
         isinstance(module, TensorQuantizer) and module.is_enabled and module.num_bits in IQ_FORMATS
         for module in model.modules()
     ):
         return
     incompatible = _iq_export_weight_shape_errors(model)
+    incompatible.extend(_iq_export_quantizer_config_errors(model))
 
     if incompatible:
         details = "\n  - ".join(incompatible)
         raise ValueError(
-            "IQ export cannot handle the following selected weights. Each weight must use the "
-            "standard 'weight' attribute and have a final dimension divisible by its format "
-            f"group size:\n  - {details}"
+            "IQ export cannot handle the following selected weights or quantizer settings:"
+            f"\n  - {details}"
         )
 
 
@@ -264,12 +276,19 @@ def _validate_iq_quantizer_config(
 def _iq_export_quantizer_config_errors(model: nn.Module) -> list[str]:
     """Return selected IQ quantizer configurations unsupported by checkpoint export."""
     incompatible: list[str] = []
-    for module, module_name, weight_name, quantizer, _ in _iq_export_weights(model):
+    for (
+        module,
+        module_name,
+        weight_name,
+        quantizer_weight_name,
+        quantizer,
+        _,
+    ) in _iq_export_weights(model):
         qualified_name = f"{module_name}.{weight_name}".lstrip(".")
         error = _iq_quantizer_config_error(
             module,
             quantizer.num_bits,
-            weight_name,
+            quantizer_weight_name,
             qualified_name,
             quantizer=quantizer,
         )
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index eb7f8d0854c..50d3a903af9 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -110,7 +110,7 @@
 from .quant_utils import (
     _get_kv_cache_postprocess_config,
     _pack_iq_weight,
-    _validate_iq_export_weight_shapes,
+    _validate_iq_export_support,
     fuse_prequant_layernorm,
     fuse_prequant_to_linear,
     get_activation_scaling_factor,
@@ -988,7 +988,7 @@ def _prepare_model_for_export(model, dtype, is_modelopt_qlora):
     the quant config.
     """
     dtype = _resolve_export_dtype(model, dtype)
-    _validate_iq_export_weight_shapes(model)
+    _validate_iq_export_support(model)
     # One tied-weight map for the whole export (amax sync + final dedup in postprocess_state_dict).
     # Sourced from HF's name-based all_tied_weights_keys, so it is correct even under FSDP/offload.
     tied_map = TiedWeightMap(model)
@@ -1286,7 +1286,7 @@ def _export_diffusers_checkpoint(
     # Validate every component before processing the first one so a later incompatible weight
     # cannot leave earlier components partially packed.
     for component in module_components.values():
-        _validate_iq_export_weight_shapes(component)
+        _validate_iq_export_support(component)
 
     # Best-effort diffusers pipeline check (kept for folder layout + model_index.json behavior)
     is_diffusers_pipe = False
diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py
index 7e11d85a853..3b5c6ac7e81 100644
--- a/modelopt/torch/export/unified_export_hf_streaming.py
+++ b/modelopt/torch/export/unified_export_hf_streaming.py
@@ -48,7 +48,7 @@
 from .quant_utils import (
     _get_kv_cache_postprocess_config,
     _postprocess_single_tensor,
-    _validate_iq_export_weight_shapes,
+    _validate_iq_export_support,
     get_quant_config,
 )
 from .registry import ExportContext
@@ -367,7 +367,7 @@ def _export_transformers_checkpoint_streaming(
 
     # --- Same model-level setup as _export_transformers_checkpoint ---
     dtype = _resolve_export_dtype(model, dtype)
-    _validate_iq_export_weight_shapes(model)
+    _validate_iq_export_support(model)
     _prepare_moe_inputs(model, dtype, is_modelopt_qlora)
 
     requantize_resmooth_fused_llm_layers(model)
diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py
index 63f274b466a..925cdeb5b1e 100644
--- a/tests/unit/torch/export/test_export_weight.py
+++ b/tests/unit/torch/export/test_export_weight.py
@@ -25,7 +25,7 @@
 import modelopt.torch.quantization as mtq
 from modelopt.torch.export.quant_utils import (
     _pack_iq_weight,
-    _validate_iq_export_weight_shapes,
+    _validate_iq_export_support,
     postprocess_state_dict,
 )
 from modelopt.torch.export.unified_export_hf import (
@@ -227,7 +227,7 @@ def test_export_iq_shape_preflight_reports_every_incompatible_weight():
         )
 
     with pytest.raises(ValueError) as exc_info:
-        _validate_iq_export_weight_shapes(model)
+        _validate_iq_export_support(model)
 
     message = str(exc_info.value)
     assert "0.weight: (4, 192)" in message
@@ -243,7 +243,7 @@ def test_export_iq_shape_preflight_skips_detailed_scan_without_iq(monkeypatch):
         lambda _model: pytest.fail("non-IQ export must skip the detailed weight scan"),
     )
 
-    _validate_iq_export_weight_shapes(model)
+    _validate_iq_export_support(model)
 
 
 def test_export_iq_shape_preflight_reports_grouped_weights():
@@ -267,7 +267,7 @@ def __init__(self):
     model.experts = GroupedLinear()
 
     with pytest.raises(ValueError) as exc_info:
-        _validate_iq_export_weight_shapes(model)
+        _validate_iq_export_support(model)
 
     message = str(exc_info.value)
     assert "experts.weight1: (4, 192)" in message
@@ -295,7 +295,7 @@ def __init__(self):
     model.experts = GroupedLinear()
 
     with pytest.raises(ValueError, match=r"experts\.weight1: \(4, 192\)"):
-        _validate_iq_export_weight_shapes(model)
+        _validate_iq_export_support(model)
 
 
 def test_export_iq_shape_preflight_reports_legacy_fused_expert_quantizers():
@@ -320,7 +320,7 @@ def __init__(self):
     model.experts = LegacyFusedExperts()
 
     with pytest.raises(ValueError, match=r"experts\.up_proj: \(4, 192\)"):
-        _validate_iq_export_weight_shapes(model)
+        _validate_iq_export_support(model)
 
 
 def test_export_iq_preflight_rejects_nonstandard_weight_before_mutation(tmp_path):
@@ -399,6 +399,31 @@ def test_export_transformers_checkpoint_runs_iq_shape_preflight_before_mutation(
         assert torch.equal(linear.weight, original_weight)
 
 
+def test_export_transformers_checkpoint_runs_iq_config_preflight_before_mutation():
+    model = nn.Sequential(
+        nn.Linear(256, 4, bias=False, dtype=torch.bfloat16),
+        nn.Linear(256, 4, bias=False, dtype=torch.bfloat16),
+    )
+    model.config = SimpleNamespace(torch_dtype=torch.bfloat16)
+    for linear in model:
+        linear.weight_quantizer = TensorQuantizer(
+            QuantizerAttributeConfig(
+                num_bits="iq2_xs",
+                block_sizes={-1: 256},
+                backend="ggml",
+            )
+        )
+    model[1].input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8))
+    original_weights = [linear.weight.detach().clone() for linear in model]
+
+    with pytest.raises(ValueError, match=r"weight-only.*1\.weight"):
+        _export_transformers_checkpoint(model)
+
+    for linear, original_weight in zip(model, original_weights):
+        assert linear.weight.dtype == torch.bfloat16
+        assert torch.equal(linear.weight, original_weight)
+
+
 def test_export_transformers_checkpoint_streaming_runs_iq_shape_preflight_before_mutation(
     tmp_path,
 ):
@@ -488,6 +513,35 @@ def test_export_iq_config_preflight_reports_weight_owner():
     assert "weight-only" in errors[0]
 
 
+def test_export_iq_config_preflight_uses_unprefixed_grouped_activation_quantizer():
+    class GroupedLinear(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.num_gemms = 2
+            self.weight0 = nn.Parameter(torch.ones(4, 256, dtype=torch.bfloat16))
+            self.weight1 = nn.Parameter(torch.ones(4, 256, dtype=torch.bfloat16))
+            quantizer_config = QuantizerAttributeConfig(
+                num_bits="iq2_xs",
+                block_sizes={-1: 256},
+                backend="ggml",
+            )
+            self.weight_quantizer = GroupedQuantizer(
+                TensorQuantizer(quantizer_config),
+                TensorQuantizer(quantizer_config),
+            )
+            self.input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8))
+
+    model = nn.Module()
+    model.experts = GroupedLinear()
+
+    errors = quant_utils._iq_export_quantizer_config_errors(model)
+
+    assert len(errors) == 2
+    assert all("weight-only" in error for error in errors)
+    assert any("experts.weight0" in error for error in errors)
+    assert any("experts.weight1" in error for error in errors)
+
+
 def test_iq_config_error_identifies_weight():
     linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16)
     linear.weight_quantizer = TensorQuantizer(