From d789b450d6ce8e77d6917b67a1d0fba4dbf4c103 Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Wed, 16 Sep 2026 09:27:57 -0700 Subject: [PATCH 01/11] Export IQ checkpoints from HF and Megatron Signed-off-by: Hung-Yueh Chiang --- 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 ++++++++++++++++--- 7 files changed, 175 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 ``.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 7ac06e73e17..9179f473e18 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 @@ -312,10 +316,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, @@ -328,6 +341,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: @@ -1031,6 +1046,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. @@ -1039,6 +1055,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. @@ -1049,7 +1066,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: @@ -1086,13 +1105,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) @@ -1128,6 +1155,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): @@ -1192,7 +1228,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( @@ -1237,7 +1275,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: @@ -1403,7 +1448,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 ) @@ -1434,7 +1481,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( @@ -1597,7 +1650,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: @@ -1712,7 +1768,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: @@ -1846,7 +1910,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( @@ -1957,7 +2023,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: From 9f9519be53c10f40289dfd966008b0e6bb33e441 Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Wed, 16 Sep 2026 14:52:46 -0700 Subject: [PATCH 02/11] Test IQ checkpoint export Signed-off-by: Hung-Yueh Chiang --- .../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 ++++++++ 4 files changed, 168 insertions(+) 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 614b5d96e2a..c4a881051d4 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 13cdcecd6d96318ebcc8d21540f8e02cc879beef Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Wed, 16 Sep 2026 16:30:31 -0700 Subject: [PATCH 03/11] Remove unused IQ search config from export tests Signed-off-by: Hung-Yueh Chiang --- tests/gpu_megatron/torch/export/test_unified_export_megatron.py | 2 -- tests/unit/torch/export/test_export_weight.py | 1 - tests/unit/torch/export/test_get_quantization.py | 1 - 3 files changed, 4 deletions(-) diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py index c4a881051d4..0d684e0f82c 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -101,7 +101,6 @@ def test_megatron_name_remapping_exports_iq_payload(qformat, payload_bytes, dequ num_bits=qformat, block_sizes={-1: 256}, backend="ggml", - backend_extra_args={"search_impl": "auto"}, ) ) exporter = object.__new__(GPTModelExporter) @@ -141,7 +140,6 @@ def test_megatron_iq_export_rejects_tensor_parallelism(): num_bits="iq2_xs", block_sizes={-1: 256}, backend="ggml", - backend_extra_args={"search_impl": "auto"}, ) ) exporter = object.__new__(GPTModelExporter) diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index f15377dc7a4..121fe564b30 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -113,7 +113,6 @@ def test_export_iq_payload_as_weight(num_bits, payload_bytes): num_bits=num_bits, block_sizes={-1: 256}, backend="ggml", - backend_extra_args={"search_impl": "auto"}, ) ) diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index 15f9657dffa..caa38033157 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -75,7 +75,6 @@ def test_iq_quantization_config(num_bits, quantization_format, payload_bytes, ef "num_bits": num_bits, "block_sizes": {-1: 256}, "backend": "ggml", - "backend_extra_args": {"search_impl": "auto"}, }, }, ], From 542b601782e4ff9eea3f930645c4d00c5acbba0e Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Wed, 16 Sep 2026 16:43:17 -0700 Subject: [PATCH 04/11] Fix IQ checkpoint export review findings Signed-off-by: Hung-Yueh Chiang --- docs/source/deployment/3_unified_hf.rst | 9 + modelopt/torch/export/convert_hf_config.py | 44 ++-- modelopt/torch/export/moe_utils.py | 5 +- modelopt/torch/export/quant_utils.py | 25 +- modelopt/torch/export/unified_export_hf.py | 3 +- .../torch/export/unified_export_megatron.py | 46 ++-- .../export/test_unified_export_megatron.py | 215 +++++++++++++++++- tests/unit/torch/export/test_export_weight.py | 1 + .../torch/export/test_get_quantization.py | 45 +++- .../plugins/test_fused_experts.py | 31 --- 10 files changed, 350 insertions(+), 74 deletions(-) diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst index 2923c227a79..5756cb85c26 100644 --- a/docs/source/deployment/3_unified_hf.rst +++ b/docs/source/deployment/3_unified_hf.rst @@ -69,6 +69,15 @@ 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. +Megatron fused-MoE export packs each expert's logical ``[out_features, in_features]`` matrix before +stacking the payloads. Such tensors therefore have shape +``[num_experts, out_features, in_features // 256, payload_bytes]``; packed blocks are never +transposed across the contraction axis. + +The generated configuration records ``quant_method: modelopt``, ``packing: ggml``, the 256-value +block size, and the payload byte count. IQ payloads are not represented as compressed-tensors +integer ``weights`` groups because all scales and indices are embedded in each packed block. + Each 74-byte IQ2_XS block represents 256 logical weights: * bytes 0--1 are the little-endian FP16 super-block scale ``d``; diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index 29a9822c5db..d2c975ac72c 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -19,6 +19,15 @@ from collections import defaultdict from typing import Any +from modelopt.torch.quantization.ggml import ( + IQ1_S_BLOCK_BYTES, + IQ1_S_BLOCK_SIZE, + IQ1_S_EFFECTIVE_BITS, + IQ2_XS_BLOCK_BYTES, + IQ2_XS_BLOCK_SIZE, + IQ2_XS_EFFECTIVE_BITS, +) + def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) -> dict[str, Any]: """Map a per-layer quant_algo string to compressed-tensors config group details. @@ -29,7 +38,8 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) Returns: Dictionary with ``input_activations`` and ``weights`` entries suitable for - a compressed-tensors ``config_groups`` entry. + a compressed-tensors ``config_groups`` entry, or ModelOpt-owned metadata for + self-contained IQ payloads. """ if quant_algo == "FP8": return { @@ -118,17 +128,24 @@ 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) + if quant_algo == "IQ1_S": + block_size = IQ1_S_BLOCK_SIZE + payload_bytes = IQ1_S_BLOCK_BYTES + effective_bits = IQ1_S_EFFECTIVE_BITS + else: + block_size = IQ2_XS_BLOCK_SIZE + payload_bytes = IQ2_XS_BLOCK_BYTES + effective_bits = IQ2_XS_EFFECTIVE_BITS + if group_size not in (None, block_size): + raise ValueError(f"{quant_algo} requires group size {block_size}, got {group_size}") + # IQ payloads are self-contained blocks, not compressed-tensors integer groups. + # Keep their format marker outside a ``weights`` quantization scheme. return { - "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, - } + "quant_algo": quant_algo, + "effective_bits": effective_bits, + "group_size": block_size, + "packing": "ggml", + "block_payload_bytes": payload_bytes, } else: warnings.warn( @@ -223,9 +240,8 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An } 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} + iq_metadata = _quant_algo_to_group_config(quant_algo_value) + new_config.update(iq_metadata) elif quant_algo_value == "NVFP4_SVD": # NVFP4 + SVDQuant: NVFP4 weights/activations plus an AWQ-style # pre_quant_scale and a low-rank residual (svdquant_lora_a/b) stored as diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index e8ae02090d2..734302690f1 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -218,10 +218,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/quant_utils.py b/modelopt/torch/export/quant_utils.py index 4d8a0e47f19..0641cad7988 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -26,6 +26,14 @@ from modelopt import __version__ from modelopt.torch.models import get_spec, list_all_possible +from modelopt.torch.quantization.ggml import ( + IQ1_S_BLOCK_BYTES, + IQ1_S_BLOCK_SIZE, + IQ1_S_EFFECTIVE_BITS, + IQ2_XS_BLOCK_BYTES, + IQ2_XS_BLOCK_SIZE, + IQ2_XS_EFFECTIVE_BITS, +) from modelopt.torch.quantization.model_calib import ( enable_stats_collection, finish_stats_collection, @@ -730,10 +738,22 @@ def process_layer_quant_config(layer_config_dict): "group_size": block_size_value, } elif v in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): - payload_bytes = 50 if v == QUANTIZATION_IQ1_S else 74 + if v == QUANTIZATION_IQ1_S: + block_size = IQ1_S_BLOCK_SIZE + payload_bytes = IQ1_S_BLOCK_BYTES + effective_bits = IQ1_S_EFFECTIVE_BITS + else: + block_size = IQ2_XS_BLOCK_SIZE + payload_bytes = IQ2_XS_BLOCK_BYTES + effective_bits = IQ2_XS_EFFECTIVE_BITS + if block_size_value != block_size: + raise ValueError( + f"{v.upper()} requires block size {block_size}, got {block_size_value}" + ) layer_config = { "quant_algo": v.upper(), - "group_size": 256, + "group_size": block_size, + "effective_bits": effective_bits, "block_payload_bytes": payload_bytes, "packing": "ggml", } @@ -1167,7 +1187,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..4aeddac0770 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -634,8 +634,7 @@ def _export_quantized_weight( 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) + setattr(sub_module, weight_name, nn.Parameter(packed_weight, requires_grad=False)) maybe_clear_cuda_cache() return diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 9179f473e18..d8fc661e840 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -1116,8 +1116,8 @@ 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-device until it can be packed + # along its contraction axis, so the CUDA packer can be used. if is_iq: return name_to_value, qformat, block_size # Getting the weight scales @@ -1156,13 +1156,18 @@ 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.""" + def _pack_iq_weight(weight: torch.Tensor, qformat: str) -> torch.Tensor: + """Pack one ``[out, in]`` weight and return its CPU payload.""" quantize_iq = quantize_iq1_s if qformat == QUANTIZATION_IQ1_S else quantize_iq2_xs packed_weight, _ = quantize_iq(weight) - return {weight_key: packed_weight.detach().cpu()} + return packed_weight.detach().cpu() + + @classmethod + def _get_iq_weight_state( + cls, weight_key: str, weight: torch.Tensor, qformat: str + ) -> dict[str, torch.Tensor]: + """Pack one ``[out, in]`` weight into the IQ checkpoint representation.""" + return {weight_key: cls._pack_iq_weight(weight, qformat)} 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.""" @@ -1878,16 +1883,19 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False, tr name_to_value.pop("input_scale") if "input_scale" in name_to_value else None ) + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + weight = self._pack_iq_weight(weight, qformat) weight_list.append(weight) weight_scale_list.append(weight_scale) weight_scale_2_list.append(weight_scale_2) input_scale_list.append(input_scale) self._record_layer_quant_config(prefix, qformat, block_size) + is_iq = qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) merged_weight = torch.stack(weight_list, dim=0) # Megatron is [num_experts, out, in]; most HF layouts want [num_experts, in, out]. - if transpose: + if transpose and not is_iq: merged_weight = merged_weight.transpose(-2, -1).contiguous() if weight_scale_2_list[0] is None: @@ -1910,9 +1918,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 is_iq or merged_weight_scale is None: self._state_dict[prefix] = merged_weight else: self._state_dict[prefix] = to_quantized_weight( @@ -1950,6 +1956,14 @@ def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None, is_mtp=F name_to_value.pop("input_scale") if "input_scale" in name_to_value else None ) + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if layer_type == "linear_fc1": + half_out = weight.shape[0] // 2 + interleaved_weight = torch.empty_like(weight) + interleaved_weight[::2] = weight[:half_out] + interleaved_weight[1::2] = weight[half_out:] + weight = interleaved_weight + weight = self._pack_iq_weight(weight, qformat) weight_list.append(weight) weight_scale_list.append(weight_scale) weight_scale_2_list.append(weight_scale_2) @@ -1957,6 +1971,7 @@ def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None, is_mtp=F bias_list.append(bias) self._record_layer_quant_config(prefix, qformat, block_size) + is_iq = qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) merged_weight = torch.stack(weight_list, dim=0) # Transpose the last two dimensions to match HuggingFace format (except for GptOssForCausalLM) @@ -1964,10 +1979,11 @@ def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None, is_mtp=F # HF format: [num_experts, in_features, out_features] # TODO: Need to decide if we want to transpose the weight or not. - merged_weight = merged_weight.transpose(-2, -1).contiguous() + if not is_iq: + merged_weight = merged_weight.transpose(-2, -1).contiguous() # Apply interleaving for GptOssForCausalLM linear_fc1 to match HF format - if layer_type == "linear_fc1": + if layer_type == "linear_fc1" and not is_iq: # Megatron has de-interleaved format, need to interleave for HF # Pattern: first half -> even indices, second half -> odd indices num_experts, in_features, out_features = merged_weight.shape @@ -2023,8 +2039,8 @@ 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)) + if is_iq: + self._state_dict[prefix] = merged_weight elif merged_weight_scale is None: # TODO: May need to modify the key name later. self._state_dict[prefix] = merged_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 0d684e0f82c..43aac0fc74e 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -47,7 +47,12 @@ 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.ggml import ( + dequantize_iq1_s, + dequantize_iq2_xs, + quantize_iq1_s, + quantize_iq2_xs, +) from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.speculative.eagle.default_config import default_eagle_config from modelopt.torch.speculative.plugins.megatron_eagle import _DynamicEagleGPTModel @@ -132,6 +137,214 @@ def test_megatron_name_remapping_exports_iq_payload(qformat, payload_bytes, dequ } +def _make_iq_experts(qformat, layer_type, *, bias=False): + experts = torch.nn.ModuleList() + generator = torch.Generator().manual_seed(1234) + for _ in range(2): + expert = torch.nn.Module() + linear = torch.nn.Linear(256, 4, bias=bias, dtype=torch.bfloat16) + with torch.no_grad(): + linear.weight.copy_(torch.randn(linear.weight.shape, generator=generator)) + if linear.bias is not None: + linear.bias.copy_(torch.randn(linear.bias.shape, generator=generator)) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=qformat, + block_sizes={-1: 256}, + backend="ggml", + ) + ) + expert.add_module(layer_type, linear) + experts.append(expert) + return experts + + +def _make_iq_exporter(): + exporter = object.__new__(GPTModelExporter) + exporter.dtype = torch.bfloat16 + exporter._state_dict = {} + exporter.exclude_modules = [] + exporter.layer_config_dict = {} + return exporter + + +def _make_iq_weight(rows): + return torch.linspace(-1, 1, rows * 256, dtype=torch.float32).reshape(rows, 256).bfloat16() + + +def _assert_iq2_payload_matches(packed, logical_weight): + expected, _ = quantize_iq2_xs(logical_weight) + torch.testing.assert_close(packed, expected.cpu(), rtol=0, atol=0) + + +def test_megatron_gated_mlp_slicing_exports_iq_payloads(): + weight = _make_iq_weight(8) + module = SimpleNamespace(config=SimpleNamespace(ffn_hidden_size=4)) + exporter = _make_iq_exporter() + exporter._get_quantized_state = lambda *a, **k: ({"weight": weight}, "iq2_xs", 256) + + exporter._gated_mlp_slicing(module, "model.layers.0.mlp.") + + _assert_iq2_payload_matches( + exporter._state_dict["model.layers.0.mlp.gate_proj.weight"], weight[:4] + ) + _assert_iq2_payload_matches( + exporter._state_dict["model.layers.0.mlp.up_proj.weight"], weight[4:] + ) + + +def test_megatron_grouped_mlp_slicing_exports_iq_payloads(): + weight = _make_iq_weight(8) + module = SimpleNamespace( + num_gemms=1, + weight0=weight, + local_expert_indices=[0], + state_dict=lambda: {"weight0": weight}, + ) + exporter = _make_iq_exporter() + exporter._get_quantized_state = lambda *a, **k: ( + {"weight": module.weight}, + "iq2_xs", + 256, + ) + + exporter._grouped_mlp_slicing( + module, + "model.layers.0.mlp.experts.{}", + gate_proj_name="gate_proj", + up_proj_name="up_proj", + ) + + _assert_iq2_payload_matches( + exporter._state_dict["model.layers.0.mlp.experts.0.gate_proj.weight"], weight[:4] + ) + _assert_iq2_payload_matches( + exporter._state_dict["model.layers.0.mlp.experts.0.up_proj.weight"], weight[4:] + ) + + +def test_megatron_qkv_slicing_exports_iq_payloads(): + weight = _make_iq_weight(8) + module = SimpleNamespace( + config=SimpleNamespace( + hidden_size=256, + num_query_groups=1, + num_attention_heads=2, + kv_channels=2, + attention_output_gate=False, + ) + ) + exporter = _make_iq_exporter() + exporter._get_quantized_state = lambda *a, **k: ({"weight": weight}, "iq2_xs", 256) + + exporter._qkv_slicing(module, "model.layers.0.self_attn.") + + reshaped = weight.reshape(4, 2, 256) + expected = { + "q_proj": reshaped[:2].reshape(4, 256), + "k_proj": reshaped[2].reshape(2, 256), + "v_proj": reshaped[3].reshape(2, 256), + } + for projection, logical_weight in expected.items(): + _assert_iq2_payload_matches( + exporter._state_dict[f"model.layers.0.self_attn.{projection}.weight"], + logical_weight, + ) + + +def test_megatron_gated_delta_net_slicing_exports_iq_payloads(): + weight = _make_iq_weight(12) + module = SimpleNamespace( + in_proj=object(), + in_proj_split_names=("query", "key", "value", "z", "beta", "alpha"), + in_proj_split_sections=(2, 2, 2, 2, 2, 2), + ) + exporter = _make_iq_exporter() + exporter._get_quantized_state = lambda *a, **k: ({"weight": weight}, "iq2_xs", 256) + + exporter._gated_delta_net_slicing(module, "model.layers.0.mixer.") + + _assert_iq2_payload_matches( + exporter._state_dict["model.layers.0.mixer.in_proj_qkv.weight"], weight[:6] + ) + _assert_iq2_payload_matches( + exporter._state_dict["model.layers.0.mixer.in_proj_z.weight"], weight[6:8] + ) + torch.testing.assert_close( + exporter._state_dict["model.layers.0.mixer.in_proj_b.weight"], weight[8:10] + ) + torch.testing.assert_close( + exporter._state_dict["model.layers.0.mixer.in_proj_a.weight"], weight[10:] + ) + + +@pytest.mark.parametrize( + ("qformat", "payload_bytes", "quantize", "dequantize"), + [ + ("iq1_s", 50, quantize_iq1_s, dequantize_iq1_s), + ("iq2_xs", 74, quantize_iq2_xs, dequantize_iq2_xs), + ], +) +def test_megatron_packed_experts_keep_iq_blocks_on_input_axis( + qformat, payload_bytes, quantize, dequantize +): + experts = _make_iq_experts(qformat, "linear_fc2") + expected_packed = torch.stack( + [quantize(expert.linear_fc2.weight)[0].cpu() for expert in experts] + ) + expected = torch.stack( + [expert.linear_fc2.weight_quantizer(expert.linear_fc2.weight) for expert in experts] + ) + exporter = _make_iq_exporter() + + exporter._pack_name_remapping( + experts, + "model.layers.0.mlp.experts.down_proj", + layer_type="linear_fc2", + ) + + packed = exporter._state_dict["model.layers.0.mlp.experts.down_proj"] + assert packed.shape == (2, 4, 1, payload_bytes) + assert packed.device.type == "cpu" + torch.testing.assert_close(packed, expected_packed, rtol=0, atol=0) + reconstructed = dequantize(packed, torch.tensor([2, 4, 256]), dtype=torch.bfloat16) + torch.testing.assert_close(reconstructed, expected, rtol=0.02, atol=0.005) + + +def test_megatron_gpt_oss_packed_experts_interleave_before_iq_packing(): + experts = _make_iq_experts("iq2_xs", "linear_fc1", bias=True) + interleave = torch.tensor([0, 2, 1, 3]) + expected_packed = torch.stack( + [quantize_iq2_xs(expert.linear_fc1.weight[interleave])[0].cpu() for expert in experts] + ) + expected_weight = torch.stack( + [ + expert.linear_fc1.weight_quantizer(expert.linear_fc1.weight)[interleave] + for expert in experts + ] + ) + expected_bias = torch.stack([expert.linear_fc1.bias[interleave] for expert in experts]) + exporter = _make_iq_exporter() + + exporter._pack_name_remapping_gpt_oss( + experts, + "model.layers.0.mlp.experts.gate_up_proj", + layer_type="linear_fc1", + ) + + prefix = "model.layers.0.mlp.experts.gate_up_proj" + packed = exporter._state_dict[prefix] + assert packed.shape == (2, 4, 1, 74) + torch.testing.assert_close(packed, expected_packed, rtol=0, atol=0) + reconstructed = dequantize_iq2_xs( + packed, + torch.tensor([2, 4, 256]), + dtype=torch.bfloat16, + ) + torch.testing.assert_close(reconstructed, expected_weight, rtol=0.02, atol=0.005) + torch.testing.assert_close(exporter._state_dict[prefix + "_bias"], expected_bias) + + 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) diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index 121fe564b30..94790846330 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -119,6 +119,7 @@ def test_export_iq_payload_as_weight(num_bits, payload_bytes): _export_quantized_weight(linear, torch.bfloat16) state_dict = postprocess_state_dict(linear.state_dict(), maxbound=448, quantization=None) + assert isinstance(linear.weight, nn.Parameter) assert state_dict["weight"].shape == (4, 1, payload_bytes) assert state_dict["weight"].dtype == torch.uint8 assert "packed_weights" not in state_dict diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index caa38033157..6e3d5b8a45b 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -44,6 +44,7 @@ get_quant_config, get_quantization_format, postprocess_state_dict, + process_layer_quant_config, ) from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer @@ -86,11 +87,47 @@ def test_iq_quantization_config(num_bits, quantization_format, payload_bytes, ef config = get_quant_config(model) assert config["quantization"]["quant_algo"] == num_bits.upper() assert config["quantization"]["block_payload_bytes"] == payload_bytes + assert config["quantization"]["effective_bits"] == effective_bits 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 "config_groups" not in hf_config + assert hf_config["group_size"] == 256 + assert hf_config["effective_bits"] == effective_bits + assert hf_config["packing"] == "ggml" + assert hf_config["block_payload_bytes"] == payload_bytes + + +def test_mixed_iq_config_group_does_not_claim_integer_weight_schema(): + converted = convert_hf_quant_config_format( + { + "quantization": { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "model.layers.0.mlp.down_proj": { + "quant_algo": "IQ2_XS", + "group_size": 256, + "effective_bits": 2.3125, + "packing": "ggml", + "block_payload_bytes": 74, + } + }, + } + } + ) + + group = converted["config_groups"]["group_0"] + assert "weights" not in group + assert group["quant_algo"] == "IQ2_XS" + assert group["packing"] == "ggml" + + +def test_iq_quantization_config_rejects_mismatched_block_size(): + with pytest.raises(ValueError, match="IQ2_XS requires block size 256, got 128"): + process_layer_quant_config( + { + "model.layers.0.mlp.down_proj.quantization": "iq2_xs", + "model.layers.0.mlp.down_proj.awq_block_size": 128, + } + ) 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..49cd999a205 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -502,37 +502,6 @@ 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 11cd58d907465933f5a552bc1a8065f84c9ba3b1 Mon Sep 17 00:00:00 2001 From: Hung-Yueh Chiang Date: Wed, 16 Sep 2026 17:54:42 -0700 Subject: [PATCH 05/11] Gate unsupported fused-MoE IQ export Signed-off-by: Hung-Yueh Chiang --- docs/source/deployment/3_unified_hf.rst | 11 ++- .../torch/export/unified_export_megatron.py | 36 ++++----- .../export/test_unified_export_megatron.py | 81 +++++-------------- 3 files changed, 41 insertions(+), 87 deletions(-) diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst index 5756cb85c26..92bca00c4ab 100644 --- a/docs/source/deployment/3_unified_hf.rst +++ b/docs/source/deployment/3_unified_hf.rst @@ -69,10 +69,13 @@ 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. -Megatron fused-MoE export packs each expert's logical ``[out_features, in_features]`` matrix before -stacking the payloads. Such tensors therefore have shape -``[num_experts, out_features, in_features // 256, payload_bytes]``; packed blocks are never -transposed across the contraction axis. +.. warning:: + Megatron fused-MoE IQ export is not currently supported. Its packed tensor would require the + deployment consumer to understand + ``[num_experts, out_features, in_features // 256, payload_bytes]`` rather than the ordinary HF + fused-expert order. The exporter raises ``NotImplementedError`` until a deployment loader owns + this layout and is covered by an integration test. Dense and individually named expert weights + continue to use the representation above. The generated configuration records ``quant_method: modelopt``, ``packing: ggml``, the 256-value block size, and the payload byte count. IQ payloads are not represented as compressed-tensors diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index d8fc661e840..8dba49693ab 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -1169,6 +1169,15 @@ def _get_iq_weight_state( """Pack one ``[out, in]`` weight into the IQ checkpoint representation.""" return {weight_key: cls._pack_iq_weight(weight, qformat)} + @staticmethod + def _reject_unsupported_fused_iq_export(qformat: str) -> None: + """Reject fused-expert IQ payloads until a deployment loader owns their layout.""" + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + raise NotImplementedError( + "Fused-MoE IQ export requires a deployment loader that supports " + "[num_experts, out_features, in_features // 256, payload_bytes]" + ) + 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): @@ -1877,25 +1886,23 @@ 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 ) + self._reject_unsupported_fused_iq_export(qformat) weight = name_to_value.pop("weight") weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) input_scale = ( name_to_value.pop("input_scale") if "input_scale" in name_to_value else None ) - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): - weight = self._pack_iq_weight(weight, qformat) weight_list.append(weight) weight_scale_list.append(weight_scale) weight_scale_2_list.append(weight_scale_2) input_scale_list.append(input_scale) self._record_layer_quant_config(prefix, qformat, block_size) - is_iq = qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) merged_weight = torch.stack(weight_list, dim=0) # Megatron is [num_experts, out, in]; most HF layouts want [num_experts, in, out]. - if transpose and not is_iq: + if transpose: merged_weight = merged_weight.transpose(-2, -1).contiguous() if weight_scale_2_list[0] is None: @@ -1918,7 +1925,7 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False, tr merged_input_scale = None # Save the merged weights - if is_iq or 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( @@ -1949,6 +1956,7 @@ 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 ) + self._reject_unsupported_fused_iq_export(qformat) 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) @@ -1956,14 +1964,6 @@ def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None, is_mtp=F name_to_value.pop("input_scale") if "input_scale" in name_to_value else None ) - if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): - if layer_type == "linear_fc1": - half_out = weight.shape[0] // 2 - interleaved_weight = torch.empty_like(weight) - interleaved_weight[::2] = weight[:half_out] - interleaved_weight[1::2] = weight[half_out:] - weight = interleaved_weight - weight = self._pack_iq_weight(weight, qformat) weight_list.append(weight) weight_scale_list.append(weight_scale) weight_scale_2_list.append(weight_scale_2) @@ -1971,7 +1971,6 @@ def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None, is_mtp=F bias_list.append(bias) self._record_layer_quant_config(prefix, qformat, block_size) - is_iq = qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) merged_weight = torch.stack(weight_list, dim=0) # Transpose the last two dimensions to match HuggingFace format (except for GptOssForCausalLM) @@ -1979,11 +1978,10 @@ def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None, is_mtp=F # HF format: [num_experts, in_features, out_features] # TODO: Need to decide if we want to transpose the weight or not. - if not is_iq: - merged_weight = merged_weight.transpose(-2, -1).contiguous() + merged_weight = merged_weight.transpose(-2, -1).contiguous() # Apply interleaving for GptOssForCausalLM linear_fc1 to match HF format - if layer_type == "linear_fc1" and not is_iq: + if layer_type == "linear_fc1": # Megatron has de-interleaved format, need to interleave for HF # Pattern: first half -> even indices, second half -> odd indices num_experts, in_features, out_features = merged_weight.shape @@ -2039,9 +2037,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 is_iq: - self._state_dict[prefix] = merged_weight - 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 43aac0fc74e..7894910bf5b 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -47,12 +47,7 @@ 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, - quantize_iq1_s, - quantize_iq2_xs, -) +from modelopt.torch.quantization.ggml import dequantize_iq1_s, dequantize_iq2_xs, quantize_iq2_xs from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.speculative.eagle.default_config import default_eagle_config from modelopt.torch.speculative.plugins.megatron_eagle import _DynamicEagleGPTModel @@ -278,71 +273,31 @@ def test_megatron_gated_delta_net_slicing_exports_iq_payloads(): ) -@pytest.mark.parametrize( - ("qformat", "payload_bytes", "quantize", "dequantize"), - [ - ("iq1_s", 50, quantize_iq1_s, dequantize_iq1_s), - ("iq2_xs", 74, quantize_iq2_xs, dequantize_iq2_xs), - ], -) -def test_megatron_packed_experts_keep_iq_blocks_on_input_axis( - qformat, payload_bytes, quantize, dequantize -): +@pytest.mark.parametrize("qformat", ["iq1_s", "iq2_xs"]) +def test_megatron_packed_experts_reject_iq_without_deployment_loader(qformat): experts = _make_iq_experts(qformat, "linear_fc2") - expected_packed = torch.stack( - [quantize(expert.linear_fc2.weight)[0].cpu() for expert in experts] - ) - expected = torch.stack( - [expert.linear_fc2.weight_quantizer(expert.linear_fc2.weight) for expert in experts] - ) exporter = _make_iq_exporter() - exporter._pack_name_remapping( - experts, - "model.layers.0.mlp.experts.down_proj", - layer_type="linear_fc2", - ) - - packed = exporter._state_dict["model.layers.0.mlp.experts.down_proj"] - assert packed.shape == (2, 4, 1, payload_bytes) - assert packed.device.type == "cpu" - torch.testing.assert_close(packed, expected_packed, rtol=0, atol=0) - reconstructed = dequantize(packed, torch.tensor([2, 4, 256]), dtype=torch.bfloat16) - torch.testing.assert_close(reconstructed, expected, rtol=0.02, atol=0.005) + with pytest.raises(NotImplementedError, match="Fused-MoE IQ export requires"): + exporter._pack_name_remapping( + experts, + "model.layers.0.mlp.experts.down_proj", + layer_type="linear_fc2", + ) + assert exporter._state_dict == {} -def test_megatron_gpt_oss_packed_experts_interleave_before_iq_packing(): +def test_megatron_gpt_oss_packed_experts_reject_iq_without_deployment_loader(): experts = _make_iq_experts("iq2_xs", "linear_fc1", bias=True) - interleave = torch.tensor([0, 2, 1, 3]) - expected_packed = torch.stack( - [quantize_iq2_xs(expert.linear_fc1.weight[interleave])[0].cpu() for expert in experts] - ) - expected_weight = torch.stack( - [ - expert.linear_fc1.weight_quantizer(expert.linear_fc1.weight)[interleave] - for expert in experts - ] - ) - expected_bias = torch.stack([expert.linear_fc1.bias[interleave] for expert in experts]) exporter = _make_iq_exporter() - exporter._pack_name_remapping_gpt_oss( - experts, - "model.layers.0.mlp.experts.gate_up_proj", - layer_type="linear_fc1", - ) - - prefix = "model.layers.0.mlp.experts.gate_up_proj" - packed = exporter._state_dict[prefix] - assert packed.shape == (2, 4, 1, 74) - torch.testing.assert_close(packed, expected_packed, rtol=0, atol=0) - reconstructed = dequantize_iq2_xs( - packed, - torch.tensor([2, 4, 256]), - dtype=torch.bfloat16, - ) - torch.testing.assert_close(reconstructed, expected_weight, rtol=0.02, atol=0.005) - torch.testing.assert_close(exporter._state_dict[prefix + "_bias"], expected_bias) + with pytest.raises(NotImplementedError, match="Fused-MoE IQ export requires"): + exporter._pack_name_remapping_gpt_oss( + experts, + "model.layers.0.mlp.experts.gate_up_proj", + layer_type="linear_fc1", + ) + assert exporter._state_dict == {} def test_megatron_iq_export_rejects_tensor_parallelism(): From 106686d77a4399b9e338d64a39e7c2d7c1be22f9 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 18 Sep 2026 07:02:17 +0000 Subject: [PATCH 06/11] Close the four open IQ export review findings Two of these let a wrong checkpoint out the door silently, which is the worst failure shape for a format nothing downstream validates. **TP>1 guard could be bypassed.** The Megatron guard keyed on ``_get_quantization_format(self.model)``, which by its own docstring returns only the first non-NONE format in the tree. A mixed-format model whose IQ layers follow, say, an FP8 one slipped past and packed TP-sharded weights as if they were whole. Add ``uses_iq_quantization``, which scans every layer, and guard on that. It reads ``num_bits`` directly rather than resolving each layer's full format, so an unrelated unsupported quantizer elsewhere in the model cannot turn the guard into an error. **W-IQ + A-FP8 exported as weight-only.** The IQ branch of ``_get_quantization_from_layer`` returned as soon as it validated the backend, while the neighbouring INT8 branch consults ``input_quantizer`` to tell SQ from WO. Both exporters then return before ``input_scale`` collection and before the ``pre_quant_scale`` handling, so an enabled activation quantizer vanished without a trace. Reject it instead; the GGML block payload has nowhere to put an activation scale. **group_size was not forwarded.** ``_quant_algo_to_group_config`` already validates it, and a test already covered that, but the top-level IQ branch in ``convert_hf_config`` called it without one, so a mismatched value was rewritten to the block size rather than rejected. The MIXED_PRECISION branch beside it was already correct. **Docs overstated the encoders.** "importance-aware" implies imatrix support these encoders do not have -- they neither refine the scale iteratively nor weight by importance, unlike upstream ``quantize_row_iq1_s_impl``. Call it codebook quantization. Four tests added. The two behavioural ones were mutation-checked: each fails with its fix reverted. That caught a first attempt at the group_size test which passed either way, because it exercised the MIXED_PRECISION path instead of the branch actually changed. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- docs/source/deployment/3_unified_hf.rst | 4 +- modelopt/torch/export/convert_hf_config.py | 6 +- modelopt/torch/export/quant_utils.py | 34 +++++++++ .../torch/export/unified_export_megatron.py | 9 ++- .../torch/export/test_get_quantization.py | 73 +++++++++++++++++++ 5 files changed, 119 insertions(+), 7 deletions(-) diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst index 92bca00c4ab..822097b5f85 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 - 1-bit codebook quantization using the GGML block layout +8. IQ2_XS - 2-bit codebook quantization using the GGML block layout .. note:: GGML has no equivalent for ModelOpt's per-tensor FP8 weight-and-activation format. In particular, diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index d2c975ac72c..24931b05137 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -240,7 +240,11 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An } new_config["config_groups"] = {"group_0": config_group_details} elif quant_algo_value in ("IQ1_S", "IQ2_XS"): - iq_metadata = _quant_algo_to_group_config(quant_algo_value) + # Forward the caller's group size so a mismatched one is rejected rather than rewritten + # to the format's block size. + iq_metadata = _quant_algo_to_group_config( + quant_algo_value, original_quantization_details.get("group_size") + ) new_config.update(iq_metadata) elif quant_algo_value == "NVFP4_SVD": # NVFP4 + SVDQuant: NVFP4 weights/activations plus an AWQ-style diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 0641cad7988..951844e1172 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -450,6 +450,27 @@ def get_weight_block_size(module: nn.Module, weight_name: str = "weight") -> int return 0 +def uses_iq_quantization(module) -> bool: + """Whether any weight quantizer in ``module`` or its children targets an IQ format. + + ``get_quantization_format`` returns the *first* non-``NONE`` format it finds, so in a + mixed-format model IQ layers sitting behind, say, an FP8 layer are invisible to it. Callers + that must reject IQ specifically need to see every layer. + + This reads ``num_bits`` directly rather than resolving each layer's full format, so an + unrelated unsupported quantizer elsewhere in the model cannot turn the check into an error. + """ + for weight_name in weight_attr_names(module): + weight_quantizer = representative_weight_quantizer(module, weight_name) + if ( + weight_quantizer is not None + and weight_quantizer.is_enabled + and weight_quantizer.num_bits in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) + ): + return True + return any(uses_iq_quantization(child) for _, child in module.named_children()) + + def get_quantization_format(module) -> str | None: """Gets the quantization string. @@ -487,6 +508,19 @@ def _get_quantization_from_layer(layer, quantizer_attr_names: QuantizerAttrNames if weight_quantizer.num_bits in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): if weight_quantizer.backend != "ggml": raise ValueError("IQ formats require the built-in 'ggml' quantization backend") + # Both exporters return before collecting input_scale and before the pre_quant_scale + # handling below, so an enabled activation quantizer would be dropped without a trace + # and the checkpoint would load as weight-only. Refuse instead. + if input_quantizer is not None and input_quantizer.is_enabled: + raise NotImplementedError( + "IQ1_S/IQ2_XS export is weight-only, but this layer has an enabled input " + "quantizer. The GGML block payload carries no activation scale, so the " + "activation quantization would be silently lost." + ) + if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"): + raise NotImplementedError( + "IQ1_S/IQ2_XS export does not support an AWQ-style pre_quant_scale." + ) return weight_quantizer.num_bits if weight_quantizer.num_bits == 4: diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 8dba49693ab..6a3b1e82a65 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -78,6 +78,7 @@ get_weight_scaling_factor_2, process_layer_quant_config, to_quantized_weight, + uses_iq_quantization, ) with import_plugin("transformers", verbose=False): @@ -317,10 +318,10 @@ 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 - ): + # Scan every layer rather than trusting quantization_format, which is only the first + # non-NONE format in the tree: a mixed-format model whose IQ layers follow, say, an FP8 + # one would otherwise slip past this guard and pack TP-sharded weights as whole ones. + if uses_iq_quantization(self.model) and get_tensor_model_parallel_world_size() != 1: raise NotImplementedError( "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model " "parallel size 1" diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index 6e3d5b8a45b..6c322d7205c 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -45,6 +45,7 @@ get_quantization_format, postprocess_state_dict, process_layer_quant_config, + uses_iq_quantization, ) from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer @@ -96,6 +97,78 @@ def test_iq_quantization_config(num_bits, quantization_format, payload_bytes, ef assert hf_config["block_payload_bytes"] == payload_bytes +def _quantize_sequential(layer_cfgs): + """Quantize a two-Linear model, one quantizer config per layer.""" + model = torch.nn.Sequential( + torch.nn.Linear(256, 256, bias=False), torch.nn.Linear(256, 256, bias=False) + ) + mtq.quantize( + model, + { + "quant_cfg": [{"quantizer_name": "*", "enable": False}, *layer_cfgs], + "algorithm": None, + }, + ) + return model + + +_IQ_WEIGHT_CFG = {"num_bits": "iq1_s", "block_sizes": {-1: 256}, "backend": "ggml"} + + +def test_uses_iq_quantization_sees_iq_behind_another_format(): + """get_quantization_format stops at the first format, so the TP guard cannot rely on it.""" + model = _quantize_sequential( + [ + {"quantizer_name": "0.weight_quantizer", "cfg": {"num_bits": (4, 3)}}, + {"quantizer_name": "1.weight_quantizer", "cfg": _IQ_WEIGHT_CFG}, + ] + ) + + assert get_quantization_format(model) == QUANTIZATION_FP8 + assert uses_iq_quantization(model) + + +def test_uses_iq_quantization_false_without_iq_layers(): + model = _quantize_sequential( + [{"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": (4, 3)}}] + ) + + assert not uses_iq_quantization(model) + + +def test_iq_export_rejects_enabled_input_quantizer(): + """IQ payloads carry no activation scale, so W-IQ + A-FP8 must not export as weight-only.""" + model = _quantize_sequential( + [ + {"quantizer_name": "*weight_quantizer", "cfg": _IQ_WEIGHT_CFG}, + {"quantizer_name": "*input_quantizer", "cfg": {"num_bits": (4, 3)}}, + ] + ) + + with pytest.raises(NotImplementedError, match="weight-only"): + get_quantization_format(model) + + +def test_iq_hf_config_rejects_mismatched_group_size(): + """A uniformly-IQ config must validate group_size, not silently rewrite it to the block size. + + The MIXED_PRECISION branch already forwards the per-layer group size; this covers the + top-level branch, which did not. + """ + with pytest.raises(ValueError, match="IQ2_XS requires group size 256, got 128"): + convert_hf_quant_config_format( + { + "quantization": { + "quant_algo": "IQ2_XS", + "group_size": 128, + "effective_bits": 2.3125, + "packing": "ggml", + "block_payload_bytes": 74, + } + } + ) + + def test_mixed_iq_config_group_does_not_claim_integer_weight_schema(): converted = convert_hf_quant_config_format( { From 070baa548fc6561fea8d1d74475e8b46c558d1a9 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 18 Sep 2026 16:38:59 +0000 Subject: [PATCH 07/11] Fix uses_iq_quantization on SequentialQuantizer and make its guard rank-uniform The helper added in the previous commit read ``weight_quantizer.num_bits`` directly. A ``SequentialQuantizer`` has ``is_enabled`` but no ``num_bits``, so this raised ``AttributeError`` on any W4A8_AWQ model -- and since ``save_pretrained`` calls the guard on every Megatron export, it raised before any format dispatch. A regression introduced by 106686d77; read the attribute defensively instead. A SequentialQuantizer is never IQ, which is a single quantizer with backend="ggml", so ``None`` is the right answer for it. The guard was also rank-local, as it was before 106686d77: under pipeline parallelism a stage holding no IQ layer skipped the raise and then blocked in ``_gather_exclude_modules`` while its peers exited -- a hang rather than a clean error, the same failure shape ``_check_weight_quantization_took_effect`` already warns about. Agree across ranks first via ``all_gather_object``, mirroring ``_gather_exclude_modules``, including its ``is_initialized`` guard for single-process export. Tests: a SequentialQuantizer case, mutation-checked -- with the defensive read reverted it fails with the exact AttributeError above. 30 tests in test_get_quantization.py, 246 across the export and ggml unit suites. The rank-agreement path is not covered here: it needs Megatron and more than one rank, neither available in this environment. ``_reject_unsupported_fused_iq_export`` still raises from rank-local inspection and has the same hazard; left alone as it is outside this round's findings. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/export/quant_utils.py | 5 ++++- .../torch/export/unified_export_megatron.py | 22 +++++++++++++++---- .../torch/export/test_get_quantization.py | 19 +++++++++++++++- 3 files changed, 40 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 951844e1172..ebe4f6bb27b 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -462,10 +462,13 @@ def uses_iq_quantization(module) -> bool: """ for weight_name in weight_attr_names(module): weight_quantizer = representative_weight_quantizer(module, weight_name) + # getattr: a SequentialQuantizer has is_enabled but no num_bits, and is never IQ -- + # IQ is a single quantizer with backend="ggml". if ( weight_quantizer is not None and weight_quantizer.is_enabled - and weight_quantizer.num_bits in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) + and getattr(weight_quantizer, "num_bits", None) + in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) ): return True return any(uses_iq_quantization(child) for _, child in module.named_children()) diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 6a3b1e82a65..fa6dd07c02e 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -318,10 +318,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) - # Scan every layer rather than trusting quantization_format, which is only the first - # non-NONE format in the tree: a mixed-format model whose IQ layers follow, say, an FP8 - # one would otherwise slip past this guard and pack TP-sharded weights as whole ones. - if uses_iq_quantization(self.model) and get_tensor_model_parallel_world_size() != 1: + if self._any_rank_uses_iq_quantization() and get_tensor_model_parallel_world_size() != 1: raise NotImplementedError( "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model " "parallel size 1" @@ -1140,6 +1137,23 @@ def _get_quantized_state( return name_to_value, qformat, block_size + def _any_rank_uses_iq_quantization(self) -> bool: + """Whether any rank's local stage holds an IQ layer. + + Two reasons this is not ``self._get_quantization_format(self.model) in (...)``. That + returns only the first non-NONE format in the tree, so a mixed-format model whose IQ + layers follow, say, an FP8 one would slip past the caller's guard and pack TP-sharded + weights as whole ones. And the scan is rank-local: under pipeline parallelism a stage + holding no IQ layer would skip the raise and then block in the next collective while its + peers exit. Agree across ranks first, mirroring ``_gather_exclude_modules``. + """ + local_uses_iq = uses_iq_quantization(self.model) + if not torch.distributed.is_initialized(): + return local_uses_iq + per_rank = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(per_rank, local_uses_iq) + return any(per_rank) + def _get_quantization_format(self, module: torch.nn.Module): return get_quantization_format(module) diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index 6c322d7205c..911567f75e5 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -47,7 +47,11 @@ process_layer_quant_config, uses_iq_quantization, ) -from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer +from modelopt.torch.quantization.nn import ( + NVFP4StaticQuantizer, + SequentialQuantizer, + TensorQuantizer, +) class _FakeAttention(torch.nn.Module): @@ -136,6 +140,19 @@ def test_uses_iq_quantization_false_without_iq_layers(): assert not uses_iq_quantization(model) +def test_uses_iq_quantization_tolerates_sequential_quantizer(): + """A SequentialQuantizer has is_enabled but no num_bits, and is never IQ. + + save_pretrained calls this on every Megatron export, so reading num_bits directly would + raise AttributeError on a W4A8_AWQ model before any format dispatch. + """ + layer = torch.nn.Linear(256, 256, bias=False) + layer.weight_quantizer = SequentialQuantizer(TensorQuantizer(), TensorQuantizer()) + assert not hasattr(layer.weight_quantizer, "num_bits") + + assert not uses_iq_quantization(torch.nn.Sequential(layer)) + + def test_iq_export_rejects_enabled_input_quantizer(): """IQ payloads carry no activation scale, so W-IQ + A-FP8 must not export as weight-only.""" model = _quantize_sequential( From a232c1a745a133fe65874f1a85b4d5bdb746f815 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 18 Sep 2026 17:05:04 +0000 Subject: [PATCH 08/11] Reject fused-MoE IQ export where every rank agrees _reject_unsupported_fused_iq_export raises from inside the two per-expert loops, so a pipeline stage or EP rank owning no fused expert skipped it and then blocked in the collectives in save_pretrained while its peers exited -- the same divergence just fixed for the TP>1 guard, one layer down. Decide it up front instead. Whether experts are packed into one fused tensor is a property of the architecture's export rule table, not of which layers a rank happens to own, so it is identical everywhere: pair it with the already rank-agreed IQ flag and raise beside the TP check. The in-loop rejection stays as a backstop for paths that do not enter save_pretrained. The predicate is a module-level pure function over a mapping dict rather than a method, so it is testable without Megatron or a second rank -- four parametrized cases plus a non-CustomModuleMapping entry. Also confirmed, no change needed: save_pretrained has no early return before the all_gather_object added in 070baa548, and already contained _gather_exclude_modules, _gather_layer_config_dict, three barriers and an all_gather_object, one of which carries a comment about stranding peers. It already required every rank to enter; this adds no new constraint, it only moves the first collective earlier. 35 tests in test_get_quantization.py, 251 across the export and ggml unit suites. The rank-agreement path itself remains uncovered here: it needs Megatron and more than one rank. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- .../torch/export/unified_export_megatron.py | 36 ++++++++++++++++--- .../torch/export/test_get_quantization.py | 29 +++++++++++++++ 2 files changed, 60 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index fa6dd07c02e..7204c4f9292 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -113,6 +113,23 @@ ] +_FUSED_EXPERT_MAPPING_FUNCS = ("pack_name_remapping", "pack_name_remapping_gpt_oss") + + +def _mappings_pack_fused_experts(mappings: dict) -> bool: + """Whether an architecture's export mapping packs experts into one fused tensor. + + This is a property of the architecture's rule table, which is identical on every rank -- + unlike the per-expert loops that apply it, which only the ranks owning those experts run. + That makes it safe to reject a fused IQ export up front, where every rank agrees. + """ + return any( + isinstance(mapping, CustomModuleMapping) + and mapping.func_name in _FUSED_EXPERT_MAPPING_FUNCS + for mapping in mappings.values() + ) + + class GPTModelExporter: """Megatron Core GPTModel Exporter. @@ -318,11 +335,20 @@ 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 self._any_rank_uses_iq_quantization() and get_tensor_model_parallel_world_size() != 1: - raise NotImplementedError( - "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model " - "parallel size 1" - ) + if self._any_rank_uses_iq_quantization(): + if get_tensor_model_parallel_world_size() != 1: + raise NotImplementedError( + "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model " + "parallel size 1" + ) + # _reject_unsupported_fused_iq_export raises from inside the per-expert loops, so a + # stage owning no fused expert would skip it and block in the collectives below while + # its peers exit. The architecture's rule table is the same everywhere, so decide here. + if _mappings_pack_fused_experts(all_mcore_hf_export_mapping.get(self.arch, {})): + raise NotImplementedError( + "Fused-MoE IQ export requires a deployment loader that supports " + "[num_experts, out_features, in_features // 256, payload_bytes]" + ) # Main export process layer_state_dicts = self.layer_state_dicts diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index 911567f75e5..dedcd82084c 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -28,6 +28,7 @@ import modelopt.torch.export.unified_export_megatron as unified_export_megatron import modelopt.torch.quantization as mtq from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format +from modelopt.torch.export.plugins.mcore_custom import CustomModuleMapping from modelopt.torch.export.quant_format import ( KV_CACHE_FP8, KV_CACHE_FP8_K_NVFP4_V, @@ -773,3 +774,31 @@ def test_moe_router_names_handle_root_module(): block = _FakeMoEBlock(hidden=16) # name == "" for the root module; the router must be "gate", not ".gate". assert _get_unquantized_moe_router_names(block) == ["gate"] + + +@pytest.mark.parametrize( + ("func_name", "expected"), + [ + ("pack_name_remapping", True), + ("pack_name_remapping_gpt_oss", True), + ("name_remapping", False), + ("gated_mlp_slicing", False), + ], +) +def test_mappings_pack_fused_experts(func_name, expected): + """The fused-expert decision must come from the architecture's rule table. + + That table is identical on every rank, unlike the per-expert loops that apply it, so it is + what lets the IQ rejection be raised where all ranks agree instead of only on the ranks + that own an expert. + """ + mappings = { + "word_embeddings": CustomModuleMapping("name_remapping", "model.embed_tokens."), + "experts": CustomModuleMapping(func_name, "model.layers.{}.mlp.experts."), + } + + assert unified_export_megatron._mappings_pack_fused_experts(mappings) is expected + + +def test_mappings_pack_fused_experts_ignores_non_mapping_entries(): + assert not unified_export_megatron._mappings_pack_fused_experts({"skip_output_scale": True}) From 0ba91aae8bf604719447f2f24db5f81587768c3a Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 18 Sep 2026 17:15:09 +0000 Subject: [PATCH 09/11] Revert the architecture-keyed fused-MoE IQ rejection a232c1a74 hoisted the fused-MoE rejection beside the TP guard to make it rank-uniform, keying it off the architecture's export rule table. That is the wrong key: the table says whether an architecture *has* fused experts, not whether *those experts* are IQ. So an IQ dense-layer export whose experts are BF16 or FP8 was rejected outright on every fused architecture -- measured, that is GptOssForCausalLM and Llama4ForConditionalGeneration, exactly the two the review named. Blocking a working configuration is worse than the hang it was meant to prevent, which only affects a configuration that is unsupported anyway. Restore the per-expert rejection as the sole mechanism. It is precise: it fires on the expert's own qformat. It is still rank-local, so the divergence the previous round asked about remains open, and the docstring now records both the hazard and why the architecture rule table cannot close it -- a correct fix needs a rank-uniform way to know the experts themselves are IQ, which means resolving the fused rule keys against the local model rather than reading a static table. I do not have a way to test that here: it needs Megatron and more than one rank, and my two attempts at closing it from static information have each traded one failure mode for another. Worth the author's judgement on whether mixed IQ + fused-expert models can occur at all, which would settle it without any detection. 30 tests in test_get_quantization.py, 246 across the export and ggml unit suites. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- .../torch/export/unified_export_megatron.py | 46 ++++++------------- .../torch/export/test_get_quantization.py | 29 ------------ 2 files changed, 14 insertions(+), 61 deletions(-) diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 7204c4f9292..06e32dfdeee 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -113,23 +113,6 @@ ] -_FUSED_EXPERT_MAPPING_FUNCS = ("pack_name_remapping", "pack_name_remapping_gpt_oss") - - -def _mappings_pack_fused_experts(mappings: dict) -> bool: - """Whether an architecture's export mapping packs experts into one fused tensor. - - This is a property of the architecture's rule table, which is identical on every rank -- - unlike the per-expert loops that apply it, which only the ranks owning those experts run. - That makes it safe to reject a fused IQ export up front, where every rank agrees. - """ - return any( - isinstance(mapping, CustomModuleMapping) - and mapping.func_name in _FUSED_EXPERT_MAPPING_FUNCS - for mapping in mappings.values() - ) - - class GPTModelExporter: """Megatron Core GPTModel Exporter. @@ -335,20 +318,11 @@ 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 self._any_rank_uses_iq_quantization(): - if get_tensor_model_parallel_world_size() != 1: - raise NotImplementedError( - "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model " - "parallel size 1" - ) - # _reject_unsupported_fused_iq_export raises from inside the per-expert loops, so a - # stage owning no fused expert would skip it and block in the collectives below while - # its peers exit. The architecture's rule table is the same everywhere, so decide here. - if _mappings_pack_fused_experts(all_mcore_hf_export_mapping.get(self.arch, {})): - raise NotImplementedError( - "Fused-MoE IQ export requires a deployment loader that supports " - "[num_experts, out_features, in_features // 256, payload_bytes]" - ) + if self._any_rank_uses_iq_quantization() 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 @@ -1212,7 +1186,15 @@ def _get_iq_weight_state( @staticmethod def _reject_unsupported_fused_iq_export(qformat: str) -> None: - """Reject fused-expert IQ payloads until a deployment loader owns their layout.""" + """Reject fused-expert IQ payloads until a deployment loader owns their layout. + + Raised from inside the per-expert loops, so it is rank-local: a stage owning no fused + expert skips it and blocks in ``save_pretrained``'s collectives while its peers exit. + Hoisting it beside the TP guard needs a rank-uniform way to know the *experts* are IQ; + keying off the architecture's rule table is not it, since that cannot distinguish an + architecture that has fused experts from one whose experts are actually IQ, and would + reject IQ dense-layer exports on GPT-OSS and Llama4. + """ if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): raise NotImplementedError( "Fused-MoE IQ export requires a deployment loader that supports " diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index dedcd82084c..911567f75e5 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -28,7 +28,6 @@ import modelopt.torch.export.unified_export_megatron as unified_export_megatron import modelopt.torch.quantization as mtq from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format -from modelopt.torch.export.plugins.mcore_custom import CustomModuleMapping from modelopt.torch.export.quant_format import ( KV_CACHE_FP8, KV_CACHE_FP8_K_NVFP4_V, @@ -774,31 +773,3 @@ def test_moe_router_names_handle_root_module(): block = _FakeMoEBlock(hidden=16) # name == "" for the root module; the router must be "gate", not ".gate". assert _get_unquantized_moe_router_names(block) == ["gate"] - - -@pytest.mark.parametrize( - ("func_name", "expected"), - [ - ("pack_name_remapping", True), - ("pack_name_remapping_gpt_oss", True), - ("name_remapping", False), - ("gated_mlp_slicing", False), - ], -) -def test_mappings_pack_fused_experts(func_name, expected): - """The fused-expert decision must come from the architecture's rule table. - - That table is identical on every rank, unlike the per-expert loops that apply it, so it is - what lets the IQ rejection be raised where all ranks agree instead of only on the ranks - that own an expert. - """ - mappings = { - "word_embeddings": CustomModuleMapping("name_remapping", "model.embed_tokens."), - "experts": CustomModuleMapping(func_name, "model.layers.{}.mlp.experts."), - } - - assert unified_export_megatron._mappings_pack_fused_experts(mappings) is expected - - -def test_mappings_pack_fused_experts_ignores_non_mapping_entries(): - assert not unified_export_megatron._mappings_pack_fused_experts({"skip_output_scale": True}) From 5d0bbabd14f053c9d77cfcb7de2afe518c13f8a0 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 18 Sep 2026 17:30:02 +0000 Subject: [PATCH 10/11] Require PP=1 for Megatron IQ export so the fused rejection is rank-uniform The fused-MoE rejection raises from inside the per-expert loops, so it only runs on ranks that own an expert -- a stage without one skipped it and blocked in save_pretrained's collectives while its peers exited. My two attempts at detecting the condition up front each traded that hang for a worse problem, the second rejecting IQ dense exports on every fused architecture. Constrain the topology instead. IQ export already required TP=1; require PP=1 as well. Both sizes are identical on every rank and the IQ flag is already agreed across them, so these raise everywhere or nowhere. With PP=1 every rank holds the same layers and reaches the same loops, and expert parallelism shards a set of experts quantized alike, so the per-expert rejection fires on all ranks together. No detection needed, and nothing over-rejected. The residual gap is a rank holding no local expert at all, which needs expert-parallel size to exceed the expert count; recorded in the docstring rather than guarded, since no supported topology reaches it. TP=1 was already load-bearing for correctness, not just uniformity: packing happens during export, so a tensor-parallel shard would be packed as if it were a whole weight. The deployment doc now states both constraints. 246 tests across the export and ggml unit suites. The guards themselves stay uncovered here -- they need Megatron and more than one rank. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- docs/source/deployment/3_unified_hf.rst | 6 +++ .../torch/export/unified_export_megatron.py | 37 +++++++++++++------ 2 files changed, 32 insertions(+), 11 deletions(-) diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst index 822097b5f85..1955506a86f 100644 --- a/docs/source/deployment/3_unified_hf.rst +++ b/docs/source/deployment/3_unified_hf.rst @@ -69,6 +69,12 @@ 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. +.. note:: + Megatron IQ export currently requires tensor and pipeline model parallel sizes of 1. Packing + happens during export, so a tensor-parallel shard would be packed as if it were a whole + weight, and a pipeline stage holding no IQ layer would not reach the same rejection as its + peers. Expert parallelism is supported, assuming every expert uses the same format. + .. warning:: Megatron fused-MoE IQ export is not currently supported. Its packed tensor would require the deployment consumer to understand diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 06e32dfdeee..9e5facda2c3 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -318,11 +318,23 @@ 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 self._any_rank_uses_iq_quantization() and get_tensor_model_parallel_world_size() != 1: - raise NotImplementedError( - "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model " - "parallel size 1" - ) + if self._any_rank_uses_iq_quantization(): + # Both sizes below are identical on every rank, and the IQ flag is agreed across + # ranks, so these raise everywhere or nowhere. Raising on only a subset would strand + # the rest in the collectives further down. + if get_tensor_model_parallel_world_size() != 1: + raise NotImplementedError( + "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model " + "parallel size 1" + ) + # Requiring PP=1 is also what makes the per-expert fused-MoE rejection safe: with + # every rank holding the same layers, that check runs on all of them rather than + # only the stages that happen to own an MoE block. + if pp_size != 1: + raise NotImplementedError( + "Megatron IQ1_S/IQ2_XS unified export currently requires pipeline model " + "parallel size 1" + ) # Main export process layer_state_dicts = self.layer_state_dicts @@ -1188,12 +1200,15 @@ def _get_iq_weight_state( def _reject_unsupported_fused_iq_export(qformat: str) -> None: """Reject fused-expert IQ payloads until a deployment loader owns their layout. - Raised from inside the per-expert loops, so it is rank-local: a stage owning no fused - expert skips it and blocks in ``save_pretrained``'s collectives while its peers exit. - Hoisting it beside the TP guard needs a rank-uniform way to know the *experts* are IQ; - keying off the architecture's rule table is not it, since that cannot distinguish an - architecture that has fused experts from one whose experts are actually IQ, and would - reject IQ dense-layer exports on GPT-OSS and Llama4. + Raised from inside the per-expert loops, so it only runs on ranks that own an expert. + The guards in ``save_pretrained`` are what make that safe: IQ export requires PP=1 and + TP=1, so every rank holds the same layers and reaches the same loops, and expert + parallelism shards a set of experts quantized alike -- so every rank arrives here with + the same ``qformat`` and they raise together rather than stranding each other in a + collective. + + The one gap left is a rank holding no local expert at all, which needs expert-parallel + size to exceed the expert count. Worth revisiting if that becomes a supported topology. """ if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): raise NotImplementedError( From 68b723003a11ea092b0bc65c52427b52ae6127a0 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 18 Sep 2026 17:43:35 +0000 Subject: [PATCH 11/11] Add the PP>1 export test and record the TEGroupedLinear detection gap test_megatron_iq_export_rejects_pipeline_parallelism mirrors the existing TP test: it patches the parallel-size helpers, so it needs no real Megatron model. Verified that the construction those tests use does trip the guard -- is_enabled True, num_bits "iq2_xs", uses_iq_quantization True -- which also confirms the existing TP test still passes now that the guard routes through _any_rank_uses_iq_quantization. On the TEGroupedLinear finding: confirmed, and broader than the IQ guard. Reproduced with a module shaped like one -- weight0..N parameters and a single GroupedQuantizer under weight_quantizer: weight_attr_names yields: [] get_quantization_format: None uses_iq_quantization: False So this is not a divergence between the guard and the exporter; both are blind to that layout for every format, and an experts-only model reports no format at all. representative_weight_quantizer handles the GroupedQuantizer correctly -- the yield in weight_attr_names is gated on a plain ``weight`` parameter that a TEGroupedLinear does not have. Patching only uses_iq_quantization would make it see IQ where the rest of the export sees nothing, rejecting TP>1 for a model the exporter would then treat as unquantized. The fix belongs in weight_attr_names, which every format's detection shares and which I cannot exercise here -- it needs Transformer Engine and Megatron. Recorded in the docstring instead. Whether this is reachable at all is worth the author's judgement: TEGroupedLinear is the grouped-expert layout, and fused-expert IQ export is rejected outright. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/export/quant_utils.py | 6 ++++ .../export/test_unified_export_megatron.py | 29 +++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index ebe4f6bb27b..d3f2b2f6b56 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -459,6 +459,12 @@ def uses_iq_quantization(module) -> bool: This reads ``num_bits`` directly rather than resolving each layer's full format, so an unrelated unsupported quantizer elsewhere in the model cannot turn the check into an error. + + Known gap, shared with ``get_quantization_format``: ``weight_attr_names`` yields nothing for + a TEGroupedLinear, whose parameters are ``weight0..N`` while its quantizer is a single + ``GroupedQuantizer`` under ``weight_quantizer``. Neither function sees such a module, so an + experts-only IQ model reports no format at all -- not just here. Closing it belongs in + ``weight_attr_names``, where it affects every format, rather than in this helper. """ for weight_name in weight_attr_names(module): weight_quantizer = representative_weight_quantizer(module, weight_name) 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 7894910bf5b..5659e0481e7 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -324,6 +324,35 @@ def test_megatron_iq_export_rejects_tensor_parallelism(): exporter.save_pretrained("unused", "unused") +def test_megatron_iq_export_rejects_pipeline_parallelism(): + """IQ packing requires PP=1 so the fused-MoE rejection reaches every rank. + + The rejection raises from inside the per-expert loops, so a stage owning no expert would + skip it and block in save_pretrained's collectives while its peers exit. PP=1 removes the + divergence rather than trying to detect it. + """ + 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", + ) + ) + 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=2), + patch.object(uem, "get_tensor_model_parallel_rank", return_value=0), + patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1), + pytest.raises(NotImplementedError, match="pipeline model parallel size 1"), + ): + exporter.save_pretrained("unused", "unused") + + def _test_unified_export_megatron( tmp_path, model_type,