diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 62e6d19ddec..6ea6bd06c84 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,8 +8,8 @@ Changelog *Quantization* -- Add IQ1_S and IQ2_XS weight-only fake quantization with GGML-compatible 256-value block encoders and built-in ``iq1_s`` / ``iq2_xs`` PTQ recipes. -- Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. +- Add IQ1_S and IQ2_XS weight-only fake quantization with GGML-compatible 256-value block encoders and built-in ``iq1_s`` / ``iq2_xs`` PTQ recipes. Unified checkpoint export stores supported Hugging Face and TP=1 Megatron weights as packed ``uint8`` tensors. +- Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8, NVFP4, IQ1_S and IQ2_XS on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration. - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md `_ for details. diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst index ccef639d00e..970aa8fc25e 100644 --- a/docs/source/deployment/3_unified_hf.rst +++ b/docs/source/deployment/3_unified_hf.rst @@ -50,6 +50,47 @@ The unified HF export API supports the following quantization formats: 4. NVFP4_AWQ - NVIDIA 4-bit floating point with AWQ optimization 5. INT4_AWQ - 4-bit integer with AWQ optimization 6. W4A8_AWQ - 4-bit weights and 8-bit activations with AWQ optimization +7. IQ1_S - GGML-compatible block format, 1.5625 bits per weight +8. IQ2_XS - GGML-compatible block format, 2.3125 bits per weight + +.. note:: + GGML has no equivalent for ModelOpt's per-tensor FP8 weight-and-activation format. In particular, + GGML does not define a first-class FP8 tensor type with the corresponding per-tensor weight and + activation scale semantics. Converting a ModelOpt FP8 checkpoint to GGUF therefore requires + conversion to another GGML-supported tensor type rather than a lossless FP8 encoding. + +IQ weight representation +~~~~~~~~~~~~~~~~~~~~~~~~ + +For IQ1_S and IQ2_XS, unified export replaces each floating-point ``.weight`` with a +``uint8`` tensor containing byte-exact GGML blocks. Its shape is +``[*logical_shape[:-1], logical_shape[-1] // 256, payload_bytes]``, where ``payload_bytes`` is 50 +for IQ1_S and 74 for IQ2_XS. No separate shape tensor is stored: a loader recovers the logical +shape as ``[*weight.shape[:-2], weight.shape[-2] * 256]``. This is unambiguous because IQ export +requires the logical last dimension to be divisible by 256. + +Each 50-byte IQ1_S block represents 256 logical weights: + +* bytes 0--1 are the little-endian FP16 super-block scale ``d``; +* bytes 2--33 are the low eight bits of 32 codebook indices, one per group of eight weights; and +* bytes 34--49 are eight little-endian ``uint16`` metadata words, one per group of 32 weights. Each + word stores the remaining three bits of four codebook indices, a 3-bit local-scale code, and one + delta-shift bit. + +The canonical 2048-by-8 IQ1_S codebook is part of the implementation rather than the checkpoint. +The complete block therefore costs ``50 * 8 / 256 = 1.5625`` bits per logical weight. + +Each 74-byte IQ2_XS block represents 256 logical weights: + +* bytes 0--1 are the little-endian FP16 super-block scale ``d``; +* bytes 2--65 are 32 little-endian ``uint16`` codes, one per group of eight weights. Each code + contains a 9-bit codebook index and seven stored sign bits; the eighth sign bit is derived from + parity; and +* bytes 66--73 contain sixteen 4-bit local-scale codes, packed two per byte. Each local scale is + shared by two adjacent eight-weight groups. + +The canonical 512-by-8 IQ2_XS codebook is part of the implementation rather than the checkpoint. +The complete block therefore costs ``74 * 8 / 256 = 2.3125`` bits per logical weight. Minimum Framework Versions -------------------------- diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index 251a37cd076..0dff17f660e 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -19,6 +19,8 @@ from collections import defaultdict from typing import Any +from .quant_format import IQ_FORMATS, iq_format_spec + def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) -> dict[str, Any]: """Map a per-layer quant_algo string to compressed-tensors config group details. @@ -117,6 +119,20 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) }, "weights": {"dynamic": False, "num_bits": 8, "type": "float", "group_size": gs}, } + elif quant_algo.lower() in IQ_FORMATS: + spec = iq_format_spec(quant_algo) + # ``num_bits`` and ``type`` are nominal format labels, not a scalar integer decode + # contract. Consumers must honor ``packing`` and ``block_payload_bytes``. + return { + "weights": { + "dynamic": False, + "num_bits": spec["num_bits"], + "type": "int", + "group_size": spec["group_size"], + "packing": spec["packing"], + "block_payload_bytes": spec["block_payload_bytes"], + } + } else: warnings.warn( f"Unsupported quantization algorithm '{quant_algo}' in " @@ -209,6 +225,10 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An "targets": ["Linear"], } new_config["config_groups"] = {"group_0": config_group_details} + elif quant_algo_value and quant_algo_value.lower() in IQ_FORMATS: + config_group_details = _quant_algo_to_group_config(quant_algo_value) + config_group_details["targets"] = ["Linear"] + new_config["config_groups"] = {"group_0": config_group_details} elif quant_algo_value == "NVFP4_SVD": # NVFP4 + SVDQuant: NVFP4 weights/activations plus an AWQ-style # pre_quant_scale and a low-rank residual (svdquant_lora_a/b) stored as diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 550fb0dbd6a..6ebc02a304f 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -38,12 +38,14 @@ def _export_weight( module: nn.Module, ctx: ExportContext, weight_name: str = "weight", + module_name: str | None = None, ) -> None: # Imported lazily to avoid a cycle: unified_export_hf imports this module to # install the built-in handlers while retaining this legacy helper's import path. from .unified_export_hf import _export_quantized_weight - _export_quantized_weight(module, ctx.dtype, weight_name) + describe_as = f"{module_name}.{weight_name}" if module_name else None + _export_quantized_weight(module, ctx.dtype, weight_name, describe_as) # Preparation handlers are registered in the same precedence as the legacy MoE prepass. @@ -155,7 +157,7 @@ def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContex Tied experts are packed independently and their duplicate keys are dropped by name in postprocess_state_dict; no per-module dedup cache is used. """ - _export_fused_experts(module, ctx.dtype) + _export_fused_experts(module, ctx.dtype, module_name=name) @ExportModuleRegistry.register(predicate=is_quantlinear) @@ -168,7 +170,7 @@ def _export_quant_linear(name: str, module: nn.Module, ctx: ExportContext) -> No if get_quantization_format(module) == QUANTIZATION_NONE: return try: - _export_weight(module, ctx) + _export_weight(module, ctx, module_name=name) except AssertionError as e: raise AssertionError( f"Failed to export module '{name}' (type={type(module).__name__}): {e}" @@ -198,7 +200,7 @@ def _export_quant_embedding(name: str, module: nn.Module, ctx: ExportContext) -> ) return try: - _export_weight(module, ctx) + _export_weight(module, ctx, module_name=name) except AssertionError as e: raise AssertionError( f"Failed to export embedding '{name}' (type={type(module).__name__}): {e}" @@ -224,4 +226,4 @@ def _export_bmm_experts(name: str, module: nn.Module, ctx: ExportContext) -> Non quantizer_attrs=["gate_up_proj_input_quantizer", "down_proj_input_quantizer"], ) for weight_name in ["gate_up_proj", "down_proj"]: - _export_weight(module, ctx, weight_name) + _export_weight(module, ctx, weight_name, module_name=name) diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index 9767d5635a4..4aa14b6eced 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/modelopt/torch/export/layerwise_export.py @@ -42,6 +42,7 @@ from .quant_utils import ( _get_kv_cache_postprocess_config, _postprocess_single_tensor, + _validate_iq_export_support, get_quant_config, get_quantization_format, ) @@ -137,6 +138,7 @@ def assert_formats_supported(module: nn.Module, scope: str) -> None: def assert_layerwise_export_supported(model: nn.Module) -> None: """Raise unless per-layer export is valid for this model.""" assert_formats_supported(model, "before calibration") + _validate_iq_export_support(model) tied = _tied_quantized_modules(model) if tied: diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 734302690f1..892267b64ed 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -15,13 +15,15 @@ """Utilities for Mixture-of-Experts (MoE) model export.""" -import copy import warnings +from copy import deepcopy from pathlib import Path import torch import torch.nn as nn +from .quant_format import IQ_FORMATS + def _delete_fused_moe_source_attrs(module: nn.Module) -> None: """Remove the 3-D fused source params and per-expert quantizer ModuleLists. @@ -48,6 +50,7 @@ def _delete_fused_moe_source_attrs(module: nn.Module) -> None: def _export_fused_experts( module: nn.Module, dtype: torch.dtype, + module_name: str | None = None, ) -> None: """Split fused MoE expert weights and export per-expert quantization scales. @@ -114,6 +117,7 @@ def _export_fused_experts( first_proj_q = first_proj_weight_quantizers[idx] if ( is_gated + and getattr(first_proj_q, "num_bits", None) not in IQ_FORMATS and getattr(first_proj_q, "is_enabled", False) and ( not hasattr(first_proj_q, "_amax") @@ -156,17 +160,23 @@ def _export_fused_experts( if uses_first_proj_quantizers else module.down_proj_weight_quantizers[idx] ) + is_iq = getattr(w_quantizer_src, "num_bits", None) in IQ_FORMATS i_quantizer = first_proj_input_q if uses_first_proj_quantizers else down_input_q - # gate/up share a weight quantizer — clone so each gets independent amax. + # Formats that export amax need independent gate/up copies before slicing it. IQ + # packing derives metadata from each weight directly, so retaining the source avoids + # duplicating a potentially populated reconstruction cache. w_quantizer = ( - copy.deepcopy(w_quantizer_src) if uses_first_proj_quantizers else w_quantizer_src + deepcopy(w_quantizer_src) + if uses_first_proj_quantizers and not is_iq + else w_quantizer_src ) # For per-channel amax (dim >= 1), proportionally slice dim-0 # to match the split weight. if ( - hasattr(w_quantizer, "_amax") + not is_iq + and hasattr(w_quantizer, "_amax") and w_quantizer._amax is not None and w_quantizer._amax.dim() >= 1 ): @@ -194,7 +204,8 @@ def _export_fused_experts( # If the weight quantizer was never calibrated, compute amax from weights. if ( - hasattr(w_quantizer, "is_enabled") + not is_iq + and hasattr(w_quantizer, "is_enabled") and w_quantizer.is_enabled and ( not hasattr(w_quantizer, "_amax") @@ -215,7 +226,10 @@ def _export_fused_experts( wrapper.weight_quantizer = w_quantizer wrapper.input_quantizer = i_quantizer - _export_quantized_weight(wrapper, dtype) + describe_as = ( + f"{module_name}.{idx}.{proj_name}.weight" if module_name is not None else None + ) + _export_quantized_weight(wrapper, dtype, describe_as=describe_as) proj = nn.Module() proj.weight = wrapper.weight diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index 832342391d0..858b8bfee09 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -71,6 +71,8 @@ def _merge_quantizer_states(objs: list) -> dict: class VllmFqGPTModelExporter(GPTModelExporter): """VLLM fakequant GPTModel exporter.""" + _packs_iq_weights = False + @staticmethod def _pop_quantizer_keys(state_dict: dict) -> None: """Remove quantizer tensors from an export shard (OrderedDict-safe).""" @@ -114,6 +116,7 @@ def _get_quantized_state( module: torch.nn.Module, dtype: torch.dtype = torch.float16, prefix: str = "", + describe_as: str | None = None, ) -> tuple[dict[str, torch.Tensor], str, int]: """Return a state_dict, quantization format, and block_size of the module. @@ -124,6 +127,8 @@ def _get_quantized_state( Args: module: The target module to perform real quantization. dtype: The default data type. + prefix: The exported module prefix. + describe_as: Optional diagnostic label accepted for base-exporter compatibility. Returns: Tuple: state_dict, quantization format, and block_size of the module. diff --git a/modelopt/torch/export/quant_format.py b/modelopt/torch/export/quant_format.py index f241880f26f..27ab3e7de50 100644 --- a/modelopt/torch/export/quant_format.py +++ b/modelopt/torch/export/quant_format.py @@ -19,6 +19,15 @@ constants, for example, are in :mod:`modelopt.torch.export.trtllm.model_config`. """ +from typing import Any + +from modelopt.torch.quantization.ggml import ( + IQ1_S_BLOCK_BYTES, + IQ1_S_BLOCK_SIZE, + IQ2_XS_BLOCK_BYTES, + IQ2_XS_BLOCK_SIZE, +) + QUANTIZATION_NONE = None QUANTIZATION_FP8 = "fp8" QUANTIZATION_INT8_SQ = "int8_sq" @@ -36,11 +45,47 @@ QUANTIZATION_FP8_PB_REAL = "fp8_pb_real" QUANTIZATION_FP8_PB_WO = "fp8_pb_wo" QUANTIZATION_FP8_PC_PT = "fp8_pc_pt" +QUANTIZATION_IQ1_S = "iq1_s" +QUANTIZATION_IQ2_XS = "iq2_xs" + +IQ_FORMATS = frozenset({QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS}) +IQ_FORMAT_SPECS: dict[str, dict[str, Any]] = { + QUANTIZATION_IQ1_S: { + "quant_algo": "IQ1_S", + "num_bits": 1, + "group_size": IQ1_S_BLOCK_SIZE, + "block_payload_bytes": IQ1_S_BLOCK_BYTES, + "packing": "ggml", + }, + QUANTIZATION_IQ2_XS: { + "quant_algo": "IQ2_XS", + "num_bits": 2, + "group_size": IQ2_XS_BLOCK_SIZE, + "block_payload_bytes": IQ2_XS_BLOCK_BYTES, + "packing": "ggml", + }, +} + + +def iq_format_spec(quantization_format: str) -> dict[str, Any]: + """Return a copy of the canonical export metadata for one IQ format.""" + try: + return IQ_FORMAT_SPECS[quantization_format.lower()].copy() + except KeyError: + raise ValueError(f"Unsupported IQ quantization format: {quantization_format}") from None + # Formats whose scales are purely per-module, so export never merges them across the q/k/v # and gate/up groups that share an input. Every other format unifies input_amax (and, for # NVFP4, weight_scale_2) across such a group, which only a whole-model forward can discover. -FUSION_FREE_FORMATS = frozenset({QUANTIZATION_FP8, QUANTIZATION_NONE, QUANTIZATION_FP8_PB_REAL}) +FUSION_FREE_FORMATS = frozenset( + { + QUANTIZATION_FP8, + *IQ_FORMATS, + QUANTIZATION_NONE, + QUANTIZATION_FP8_PB_REAL, + } +) KV_CACHE_FP8 = "FP8" KV_CACHE_FP8_K_NVFP4_V = "FP8_K_NVFP4_V" diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 35a9cd29d29..b3eec78480a 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -26,6 +26,7 @@ from modelopt import __version__ from modelopt.torch.models import get_spec, list_all_possible +from modelopt.torch.quantization.ggml import quantize_iq1_s, quantize_iq2_xs from modelopt.torch.quantization.model_calib import ( enable_stats_collection, finish_stats_collection, @@ -47,9 +48,15 @@ ) from modelopt.torch.utils import clear_cuda_cache -from ..quantization.nn import NVFP4StaticQuantizer, SequentialQuantizer, TensorQuantizer +from ..quantization.nn import ( + GroupedQuantizer, + NVFP4StaticQuantizer, + SequentialQuantizer, + TensorQuantizer, +) from .model_utils import TiedWeightMap, get_language_model_from_vl from .quant_format import ( + IQ_FORMATS, KV_CACHE_FP8, KV_CACHE_FP8_K_NVFP4_V, KV_CACHE_INT8, @@ -62,6 +69,8 @@ QUANTIZATION_INT4_AWQ, QUANTIZATION_INT8_SQ, QUANTIZATION_INT8_WO, + QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_XS, QUANTIZATION_MXFP4, QUANTIZATION_MXFP8, QUANTIZATION_NONE, @@ -72,11 +81,271 @@ QUANTIZATION_W4A8_MXFP4_FP8, QUANTIZATION_W4A8_NVFP4_FP8, QUANTIZATION_W4A16_NVFP4, + iq_format_spec, ) logger = logging.getLogger(__name__) +def _iq_export_weight_quantizer( + module: nn.Module, weight_name: str +) -> tuple[nn.Module | None, bool]: + """Return a representative quantizer and whether ``weight_name`` uses a plural layout.""" + quantizer = representative_weight_quantizer(module, weight_name) + quantizer_name = quantizer_attr_names(weight_name).weight_quantizer + plural_quantizers = getattr(module, quantizer_name + "s", None) + if quantizer is not None: + return quantizer, isinstance(plural_quantizers, nn.ModuleList) + + if weight_name == getattr(module, "_first_proj_attr", None) and weight_name != "gate_up_proj": + # Older non-gated fused-expert wrappers kept the first-projection quantizers under + # the gate_up_proj sentinel name rather than the projection's actual attribute name. + plural_quantizers = getattr(module, "gate_up_proj_weight_quantizers", None) + if isinstance(plural_quantizers, nn.ModuleList) and len(plural_quantizers) > 0: + return plural_quantizers[0], True + return None, False + + +def _iq_export_weights( + model: nn.Module, +) -> Generator[tuple[nn.Module, str, str, str, TensorQuantizer, bool], None, None]: + """Yield each selected IQ weight and the metadata needed by export preflight.""" + for module_name, module in model.named_modules(): + inspected: set[str] = set() + for weight_name in weight_attr_names(module): + qualified_name = f"{module_name}.{weight_name}".lstrip(".") + if qualified_name in inspected: + continue + inspected.add(qualified_name) + quantizer, uses_plural_quantizers = _iq_export_weight_quantizer(module, weight_name) + if ( + isinstance(quantizer, TensorQuantizer) + and quantizer.is_enabled + and quantizer.num_bits in IQ_FORMATS + ): + yield ( + module, + module_name, + weight_name, + weight_name, + quantizer, + uses_plural_quantizers, + ) + + grouped_quantizer = getattr(module, "weight_quantizer", None) + if isinstance(grouped_quantizer, GroupedQuantizer) and len(grouped_quantizer) > 0: + num_gemms = int(getattr(module, "num_gemms", len(grouped_quantizer))) + for index in range(max(num_gemms, len(grouped_quantizer))): + quantizer = grouped_quantizer[min(index, len(grouped_quantizer) - 1)] + weight_name = f"weight{index}" + qualified_name = f"{module_name}.{weight_name}".lstrip(".") + if qualified_name in inspected: + continue + inspected.add(qualified_name) + if ( + isinstance(quantizer, TensorQuantizer) + and quantizer.is_enabled + and quantizer.num_bits in IQ_FORMATS + ): + # Grouped linears are rebuilt as standard per-expert modules before packing. + yield module, module_name, weight_name, "weight", quantizer, True + + +def _iq_export_weight_shape_error( + module: nn.Module, + module_name: str, + weight_name: str, + quantizer: TensorQuantizer, + *, + allow_nonstandard_name: bool = False, +) -> str | None: + """Return one IQ export compatibility error, if any, for ``weight_name``.""" + qualified_name = f"{module_name}.{weight_name}".lstrip(".") + if weight_name != "weight" and not allow_nonstandard_name: + return f"{qualified_name}: nonstandard weight attributes are not supported" + + weight = getattr(module, weight_name, None) + if not isinstance(weight, torch.Tensor): + return None + group_size = iq_format_spec(quantizer.num_bits)["group_size"] + if weight.dim() == 0 or weight.shape[-1] % group_size: + return f"{qualified_name}: {tuple(weight.shape)}" + return None + + +def _iq_export_weight_shape_errors(model: nn.Module) -> list[str]: + """Return selected IQ weights that cannot be handled by unified export.""" + incompatible: list[str] = [] + for ( + module, + module_name, + weight_name, + _, + quantizer, + uses_plural_quantizers, + ) in _iq_export_weights(model): + # Fused-expert plural and grouped quantizers are rebuilt as standard per-expert modules. + error = _iq_export_weight_shape_error( + module, + module_name, + weight_name, + quantizer, + allow_nonstandard_name=uses_plural_quantizers, + ) + if error is not None: + incompatible.append(error) + return incompatible + + +def _validate_iq_export_support(model: nn.Module) -> None: + """Reject unsupported IQ weights and quantizer settings before export mutates the model.""" + if not any( + isinstance(module, TensorQuantizer) and module.is_enabled and module.num_bits in IQ_FORMATS + for module in model.modules() + ): + return + incompatible = _iq_export_weight_shape_errors(model) + incompatible.extend(_iq_export_quantizer_config_errors(model)) + + if incompatible: + details = "\n - ".join(incompatible) + raise ValueError( + "IQ export cannot handle the following selected weights or quantizer settings:" + f"\n - {details}" + ) + + +def _iq_quantizer_config_error( + module: nn.Module, + quantization_format: str, + weight_name: str = "weight", + describe_as: str | None = None, + quantizer: nn.Module | None = None, +) -> Exception | None: + """Return an IQ export configuration error without mutating the module.""" + where = f" for '{describe_as}'" if describe_as else "" + if quantizer is None: + quantizer = representative_weight_quantizer(module, weight_name) + if not isinstance(quantizer, TensorQuantizer): + return ValueError( + f"{quantization_format.upper()} export requires one TensorQuantizer{where}" + ) + if quantizer.num_bits != quantization_format or quantizer.backend != "ggml": + return ValueError( + f"{quantization_format.upper()} export requires the matching ggml quantizer{where}" + ) + extra_args = quantizer.backend_extra_args or {} + unknown_keys = set(extra_args) - {"search_impl"} + if unknown_keys: + return ValueError( + f"{quantization_format.upper()} export received unsupported backend_extra_args keys " + f"{sorted(unknown_keys)}{where}" + ) + search_impl = extra_args.get("search_impl", "auto") + if search_impl != "auto": + return NotImplementedError( + f"{quantization_format.upper()} export only supports search_impl='auto'{where}" + ) + quantizer_names = quantizer_attr_names(weight_name) + for quantizer_name in (quantizer_names.input_quantizer, quantizer_names.output_quantizer): + activation_quantizer = getattr(module, quantizer_name, None) + if activation_quantizer is not None and ( + getattr(activation_quantizer, "is_enabled", False) + or getattr(activation_quantizer, "pre_quant_scale", None) is not None + ): + return NotImplementedError( + f"{quantization_format.upper()} export is weight-only{where}; activation " + "quantization " + "and pre_quant_scale are not represented in the checkpoint" + ) + return None + + +def _validate_iq_quantizer_config( + module: nn.Module, + quantization_format: str, + weight_name: str = "weight", + describe_as: str | None = None, +) -> None: + """Reject IQ export settings that the checkpoint encoder cannot reproduce.""" + error = _iq_quantizer_config_error(module, quantization_format, weight_name, describe_as) + if error is not None: + raise error + + +def _iq_export_quantizer_config_errors(model: nn.Module) -> list[str]: + """Return selected IQ quantizer configurations unsupported by checkpoint export.""" + incompatible: list[str] = [] + for ( + module, + module_name, + weight_name, + quantizer_weight_name, + quantizer, + _, + ) in _iq_export_weights(model): + qualified_name = f"{module_name}.{weight_name}".lstrip(".") + error = _iq_quantizer_config_error( + module, + quantizer.num_bits, + quantizer_weight_name, + qualified_name, + quantizer=quantizer, + ) + if error is not None: + incompatible.append(str(error)) + return incompatible + + +def _pack_iq_weight( + weight: torch.Tensor, + quantization_format: str, + module: nn.Module | None = None, + weight_name: str = "weight", + describe_as: str | None = None, +) -> torch.Tensor: + """Pack one IQ weight, using ``describe_as`` only for failure attribution.""" + if quantization_format not in IQ_FORMATS: + raise ValueError(f"Unsupported IQ quantization format: {quantization_format}") + if module is not None: + if describe_as is None: + describe_as = f"{type(module).__name__}.{weight_name}" + _validate_iq_quantizer_config( + module, quantization_format, weight_name, describe_as=describe_as + ) + try: + if quantization_format == QUANTIZATION_IQ1_S: + packed_weight, _ = quantize_iq1_s(weight) + elif quantization_format == QUANTIZATION_IQ2_XS: + packed_weight, _ = quantize_iq2_xs(weight) + else: + raise NotImplementedError( + f"No checkpoint packer is registered for {quantization_format.upper()}" + ) + except (TypeError, ValueError) as exc: + if describe_as is None: + describe_as = f"state_dict.{weight_name}" + raise type(exc)( + f"Failed to pack {quantization_format.upper()} weight " + f"'{describe_as}' with shape {tuple(weight.shape)}: {exc}" + ) from exc + + spec = iq_format_spec(quantization_format) + expected_shape = ( + *weight.shape[:-1], + weight.shape[-1] // spec["group_size"], + spec["block_payload_bytes"], + ) + if packed_weight.shape != expected_shape: + where = describe_as or f"state_dict.{weight_name}" + raise RuntimeError( + f"{quantization_format.upper()} packing for weight '{where}' produced " + f"{tuple(packed_weight.shape)}, expected {expected_shape}; logical-shape recovery " + "is no longer valid" + ) + return packed_weight + + def _has_large_fp8_scale(value: torch.Tensor) -> bool: """Return whether an FP8 scale exceeds the recommended threshold.""" return bool(torch.any(value > 0.5)) @@ -474,6 +743,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 IQ_FORMATS: + 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 +996,12 @@ def process_layer_quant_config(layer_config_dict): "quant_algo": "MXFP8", "group_size": block_size_value, } + elif v in IQ_FORMATS: + spec = iq_format_spec(v) + layer_config = { + key: spec[key] + for key in ("quant_algo", "group_size", "block_payload_bytes", "packing") + } else: layer_config = {"quant_algo": v} diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 8670af77403..50d3a903af9 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -94,6 +94,7 @@ ) from .quant_format import ( FUSION_FREE_FORMATS, + IQ_FORMATS, QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, QUANTIZATION_FP8_PC_PT, @@ -108,6 +109,8 @@ ) from .quant_utils import ( _get_kv_cache_postprocess_config, + _pack_iq_weight, + _validate_iq_export_support, fuse_prequant_layernorm, fuse_prequant_to_linear, get_activation_scaling_factor, @@ -594,6 +597,7 @@ def _export_quantized_weight( sub_module: nn.Module, dtype: torch.dtype, weight_name: str = "weight", + describe_as: str | None = None, ): """For the given weight attr of the sub_module, export the quantization info of it. @@ -621,6 +625,20 @@ def _export_quantized_weight( "which dispatches to the streaming writer that materialises weights layer-by-layer." ) + if quantization_format in IQ_FORMATS: + if weight_name != "weight": + where = describe_as or f"{type(sub_module).__name__}.{weight_name}" + raise NotImplementedError( + "IQ unified export currently supports modules with a standard 'weight' " + f"attribute, got '{where}'" + ) + packed_weight = _pack_iq_weight( + weight.to(dtype), quantization_format, sub_module, weight_name, describe_as + ) + setattr(sub_module, weight_name, nn.Parameter(packed_weight, requires_grad=False)) + maybe_clear_cuda_cache() + return + weight_quantizer: TensorQuantizer | SequentialQuantizer = getattr( sub_module, quantizer_attrs.weight_quantizer ) @@ -970,6 +988,7 @@ def _prepare_model_for_export(model, dtype, is_modelopt_qlora): the quant config. """ dtype = _resolve_export_dtype(model, dtype) + _validate_iq_export_support(model) # One tied-weight map for the whole export (amax sync + final dedup in postprocess_state_dict). # Sourced from HF's name-based all_tied_weights_keys, so it is correct even under FSDP/offload. tied_map = TiedWeightMap(model) @@ -1264,6 +1283,10 @@ def _export_diffusers_checkpoint( module_components = { name: comp for name, comp in all_components.items() if isinstance(comp, nn.Module) } + # Validate every component before processing the first one so a later incompatible weight + # cannot leave earlier components partially packed. + for component in module_components.values(): + _validate_iq_export_support(component) # Best-effort diffusers pipeline check (kept for folder layout + model_index.json behavior) is_diffusers_pipe = False diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index 32477173e51..3b5c6ac7e81 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -48,6 +48,7 @@ from .quant_utils import ( _get_kv_cache_postprocess_config, _postprocess_single_tensor, + _validate_iq_export_support, get_quant_config, ) from .registry import ExportContext @@ -366,6 +367,7 @@ def _export_transformers_checkpoint_streaming( # --- Same model-level setup as _export_transformers_checkpoint --- dtype = _resolve_export_dtype(model, dtype) + _validate_iq_export_support(model) _prepare_moe_inputs(model, dtype, is_modelopt_qlora) requantize_resmooth_fused_llm_layers(model) diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 37c896fd609..721c08cefaf 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -35,7 +35,10 @@ from safetensors.torch import save_file from modelopt import __version__ -from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer +from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( + GroupedQuantizer, + TensorQuantizer, +) from modelopt.torch.utils import import_plugin, warn_rank_0 from .convert_hf_config import convert_hf_quant_config_format @@ -56,6 +59,7 @@ ) from .plugins.megatron_importer import GPTModelImporter, _get_mamba_conv1d from .quant_format import ( + IQ_FORMATS, KV_CACHE_FP8, KV_CACHE_NVFP4, QUANTIZATION_FP8, @@ -64,8 +68,13 @@ QUANTIZATION_NONE, QUANTIZATION_NVFP4, QUANTIZATION_W4A16_NVFP4, + iq_format_spec, ) from .quant_utils import ( + _iq_export_quantizer_config_errors, + _iq_export_weight_shape_errors, + _pack_iq_weight, + _validate_iq_quantizer_config, get_activation_scaling_factor, get_kv_cache_dtype, get_kv_cache_scaling_factor, @@ -94,6 +103,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 @@ -129,6 +139,8 @@ class GPTModelExporter: clamp_kv_cache_scales: Whether to clamp FP8 KV cache scaling factors to at least 1.0. """ + _packs_iq_weights = True + def __init__( self, model: torch.nn.Module, @@ -264,6 +276,7 @@ def save_pretrained_extra_modules( save_directory: str | os.PathLike, ): """Save a EAGLE or Medusa checkpoints which can be deployed by vLLM and TensorRT-LLM.""" + self._validate_iq_export() # We use the last PP rank to write the config because # medusa_heads and eagle_module only exist in the last stage. pp_rank = get_pipeline_model_parallel_rank() @@ -287,6 +300,94 @@ def _is_sidecar_writer_rank(is_last_stage_main_rank: bool) -> bool: and get_expert_model_parallel_rank() == 0 ) + @staticmethod + def _model_has_iq_quantizer(model: torch.nn.Module) -> bool: + """Return whether this rank owns any enabled IQ weight quantizer.""" + return any( + isinstance(module, TensorQuantizer) + and module.is_enabled + and module.num_bits in IQ_FORMATS + for module in model.modules() + ) + + def _model_has_packed_expert_iq_quantizer(self) -> bool: + """Return whether this rank would route an IQ expert through a packed-expert rule.""" + packed_projection_names = { + rule_name.rsplit(".", 1)[-1] + for rule_name, rule in getattr(self, "rules", {}).items() + if getattr(rule, "_modelopt_export_func_name", None) + in {"pack_name_remapping", "pack_name_remapping_gpt_oss"} + } + if not packed_projection_names: + return False + + for module_name, module in self.model.named_modules(): + path = module_name.split(".") + if ( + "local_experts" in path + and path[-1] in packed_projection_names + and self._model_has_iq_quantizer(module) + ): + return True + return False + + @staticmethod + def _collective_iq_export_flags( + local_iq: bool, local_error: bool, local_packed_expert_iq: bool + ) -> tuple[bool, bool, bool]: + """Combine IQ presence, preflight-error, and packed-expert flags across ranks.""" + if not torch.distributed.is_initialized(): + return local_iq, local_error, local_packed_expert_iq + device = ( + torch.device("cuda", torch.cuda.current_device()) + if torch.distributed.get_backend() == torch.distributed.Backend.NCCL + else torch.device("cpu") + ) + flags = torch.tensor( + [int(local_iq), int(local_error), int(local_packed_expert_iq)], + dtype=torch.int32, + device=device, + ) + torch.distributed.all_reduce(flags, op=torch.distributed.ReduceOp.MAX) + return bool(flags[0].item()), bool(flags[1].item()), bool(flags[2].item()) + + def _validate_iq_export(self) -> None: + """Reject unsupported IQ export configurations in one collective preflight.""" + if not self._packs_iq_weights: + return + + local_iq = self._model_has_iq_quantizer(self.model) + tp_size = get_tensor_model_parallel_world_size() + local_errors = [] + if local_iq and tp_size == 1: + local_errors.extend(_iq_export_weight_shape_errors(self.model)) + local_errors.extend(_iq_export_quantizer_config_errors(self.model)) + local_packed_expert_iq = local_iq and self._model_has_packed_expert_iq_quantizer() + any_iq, any_error, any_packed_expert_iq = self._collective_iq_export_flags( + local_iq, bool(local_errors), local_packed_expert_iq + ) + if not any_iq: + return + if tp_size != 1: + raise NotImplementedError( + "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model " + "parallel size 1" + ) + if any_packed_expert_iq: + raise NotImplementedError( + "Megatron IQ1_S/IQ2_XS export does not yet support packed expert weights " + "because the HF transpose changes the 256-value block axis" + ) + if not any_error: + return + if local_errors: + details = "\n - ".join(local_errors) + raise ValueError( + "IQ export rejected unsupported selected weights or quantizer settings:" + f"\n - {details}" + ) + raise ValueError("Another distributed rank owns an IQ weight that failed export preflight") + def save_pretrained( self, save_directory: str | os.PathLike, @@ -297,6 +398,8 @@ def save_pretrained( Args: save_directory: Directory to which to save. Will be created if it doesn't exist. """ + self._validate_iq_export() + pp_rank = get_pipeline_model_parallel_rank() pp_size = get_pipeline_model_parallel_world_size() tp_rank = get_tensor_model_parallel_rank() @@ -309,10 +412,10 @@ 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) # Main export process layer_state_dicts = self.layer_state_dicts - quantization_format = self._get_quantization_format(self.model) quantization = None if quantization_format in ( QUANTIZATION_FP8_PB_REAL, @@ -325,6 +428,8 @@ def save_pretrained( quantization = "NVFP4" elif quantization_format == QUANTIZATION_W4A16_NVFP4: quantization = "W4A16_NVFP4" + elif quantization_format in IQ_FORMATS: + quantization = quantization_format.upper() if is_last_stage_main_rank: if is_writer_rank: @@ -388,6 +493,11 @@ def save_pretrained( } if quantization in ("NVFP4", "W4A16_NVFP4"): # update block size quantization_config["group_size"] = 16 + elif quantization_format in IQ_FORMATS: + spec = iq_format_spec(quantization_format) + quantization_config.update( + {key: spec[key] for key in ("group_size", "block_payload_bytes", "packing")} + ) if gathered_kv_cache_dtype is not None: quantization_config["kv_cache_quant_algo"] = gathered_kv_cache_dtype @@ -1009,9 +1119,14 @@ def _custom_mapping_to_lambda(mapping): func = method_map[mapping.func_name] prefix = mapping.target_name_or_prefix func_kwargs = mapping.func_kwargs - return lambda m, *args, **kwargs: func( - m, prefix.format(*args), **{**func_kwargs, **kwargs} - ) + + def apply_mapping(m, *args, **kwargs): + return func(m, prefix.format(*args), **{**func_kwargs, **kwargs}) + + # Preflight uses this marker to identify packed-expert rules before any rank enters + # per-layer conversion, where a local exception would strand other pipeline ranks. + setattr(apply_mapping, "_modelopt_export_func_name", mapping.func_name) + return apply_mapping for arch, mappings in all_mcore_hf_export_mapping.items(): all_rules[arch] = { @@ -1027,6 +1142,7 @@ def _get_weight_bias( module: torch.nn.Module, dtype: torch.dtype = torch.float16, name_to_value: dict[str, torch.Tensor] | None = None, + keep_weight_device: bool = False, ) -> dict[str, torch.Tensor]: """Get the weight and bias of the module. @@ -1035,6 +1151,7 @@ def _get_weight_bias( dtype: The data type of the weight and bias. name_to_value: The dictionary to store the weight and bias. A new dict is created if not provided. + keep_weight_device: Keep the weight on its current device instead of moving it to CPU. Returns: The dictionary containing the weight and bias. @@ -1045,7 +1162,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: @@ -1065,6 +1184,7 @@ def _get_quantized_state( module: torch.nn.Module, dtype: torch.dtype = torch.float16, prefix: str = "", + describe_as: str | None = None, ) -> tuple[dict[str, torch.Tensor], str, int]: """Return a state_dict, quantization format, and block_size of the module. @@ -1072,6 +1192,7 @@ def _get_quantized_state( module: The target module to perform real quantization. dtype: The default data type. prefix: The prefix of the layer. + describe_as: Optional logical weight label used only in IQ diagnostics. Returns: Tuple: state_dict, quantization format, and block_size of the module. @@ -1082,13 +1203,25 @@ 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 IQ_FORMATS + if is_iq: + weight_label = describe_as or (f"{prefix}weight" if prefix else "weight") + _validate_iq_quantizer_config(module, qformat, describe_as=weight_label) + name_to_value = self._get_weight_bias( + module, dtype, name_to_value, keep_weight_device=is_iq + ) 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 its current device until its + # final HF layout has been produced: this avoids a CPU round-trip in slicing paths and + # leaves this seam ready for a device-side packer. + if is_iq: + return name_to_value, qformat, block_size # Getting the weight scales weight_scale = get_weight_scaling_factor(module) weight_scale_2 = get_weight_scaling_factor_2(module) @@ -1124,6 +1257,19 @@ 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, module: torch.nn.Module + ) -> dict[str, torch.Tensor]: + """Validate and pack one final-layout weight into the IQ checkpoint representation.""" + if get_tensor_model_parallel_world_size() != 1: + raise NotImplementedError( + "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model " + "parallel size 1" + ) + packed_weight = _pack_iq_weight(weight, qformat, module=module, describe_as=weight_key) + return {weight_key: packed_weight.detach().cpu()} + def _record_layer_quant_config(self, prefix: str, qformat: str | None, block_size: int | None): """Record per-HF-layer quantization metadata for mixed precision exports.""" if qformat in (None, QUANTIZATION_NONE): @@ -1188,7 +1334,11 @@ 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 IQ_FORMATS: + self._state_dict.update( + self._get_iq_weight_state(prefix + "weight", weight, qformat, module) + ) + elif weight_scale is None: self._state_dict[prefix + "weight"] = weight else: self._state_dict[prefix + "weight"] = to_quantized_weight( @@ -1233,7 +1383,18 @@ 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 IQ_FORMATS: + self._state_dict.update( + self._get_iq_weight_state( + gate_proj_prefix + "weight", gate_proj_weight, qformat, module + ) + ) + self._state_dict.update( + self._get_iq_weight_state( + up_proj_prefix + "weight", up_proj_weight, qformat, module + ) + ) + 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: @@ -1388,18 +1549,27 @@ def _grouped_mlp_slicing( # weight_scale_2 derivation asserts one. Max-calibration weight amax # is exactly max(|W|), so compute it from this expert's weight. _wq = module.weight_quantizer - if getattr(_wq, "_amax", None) is None and getattr(_wq, "is_enabled", False): + if ( + getattr(_wq, "num_bits", None) not in IQ_FORMATS + and getattr(_wq, "_amax", None) is None + and getattr(_wq, "is_enabled", False) + ): _wq.amax = module.weight.detach().abs().max().float() temp_amax_wqs.append(_wq) name_to_value, qformat, block_size = self._get_quantized_state( - module, self.dtype, prefix=prefix + module, + self.dtype, + prefix=prefix, + describe_as=expert_prefix + "weight", ) weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) name_to_value.pop("weight", None) 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 IQ_FORMATS: + weight = weight.cpu() weight_scale_cpu = ( weight_scale.detach().cpu().clone() if weight_scale is not None else None ) @@ -1430,7 +1600,13 @@ def _grouped_mlp_slicing( ] for shard_prefix, shard_weight, shard_scale in shards: - if shard_scale is None: + if qformat in IQ_FORMATS: + local_expert_state.update( + self._get_iq_weight_state( + shard_prefix + "weight", shard_weight, qformat, module + ) + ) + elif shard_scale is None: local_expert_state[shard_prefix + "weight"] = shard_weight else: local_expert_state[shard_prefix + "weight"] = to_quantized_weight( @@ -1593,7 +1769,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 IQ_FORMATS: + for key, weight in zip(proj_keys, proj_weights): + self._state_dict.update(self._get_iq_weight_state(key, weight, qformat, module)) + elif weight_scale is None: for key, weight in zip(proj_keys, proj_weights): self._state_dict[key] = weight else: @@ -1708,7 +1887,17 @@ 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 IQ_FORMATS: + for proj_prefix, proj_weight in zip(proj_prefixes, proj_weights): + if proj_prefix in keep_bf16: + self._state_dict[proj_prefix + "weight"] = proj_weight.cpu() + else: + self._state_dict.update( + self._get_iq_weight_state( + proj_prefix + "weight", proj_weight, qformat, in_proj + ) + ) + elif weight_scale is None: for key, proj_weight in zip(proj_keys, proj_weights): self._state_dict[key] = proj_weight else: @@ -1804,6 +1993,11 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False, tr name_to_value, qformat, block_size = self._get_quantized_state( getattr(expert, layer_type), self.dtype, prefix=prefix ) + if qformat in IQ_FORMATS: + raise NotImplementedError( + "Megatron IQ1_S/IQ2_XS export does not yet support packed expert weights " + "because the HF transpose changes the 256-value block axis" + ) weight = name_to_value.pop("weight") weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) input_scale = ( @@ -1873,6 +2067,11 @@ def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None, is_mtp=F name_to_value, qformat, block_size = self._get_quantized_state( getattr(expert, layer_type), self.dtype, prefix=prefix ) + if qformat in IQ_FORMATS: + raise NotImplementedError( + "Megatron IQ1_S/IQ2_XS export does not yet support packed expert weights " + "because the HF transpose changes the 256-value block axis" + ) weight = name_to_value.pop("weight") bias = name_to_value.pop("bias", None) weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) diff --git a/modelopt_recipes/general/ptq/iq1_s.yaml b/modelopt_recipes/general/ptq/iq1_s.yaml index 0eb80fcf29d..9a402494b1b 100644 --- a/modelopt_recipes/general/ptq/iq1_s.yaml +++ b/modelopt_recipes/general/ptq/iq1_s.yaml @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# IQ1_S weight-only PTQ. +# IQ1_S weight-only PTQ and unified-checkpoint export. imports: preset: configs/ptq/presets/model/iq1_s @@ -10,6 +10,7 @@ metadata: recipe_type: ptq description: >- Applies uniform GGML-compatible IQ1_S weight-only quantization to eligible linear layers. - This is not a mixed per-tensor precision preset. No calibration data is required. + Unified export writes the packed block payload to each quantized Linear weight. This is not + a mixed per-tensor precision preset. No calibration data is required. quantize: $import: preset diff --git a/modelopt_recipes/general/ptq/iq2_xs.yaml b/modelopt_recipes/general/ptq/iq2_xs.yaml index c0edc2d5351..d5e2e36bad6 100644 --- a/modelopt_recipes/general/ptq/iq2_xs.yaml +++ b/modelopt_recipes/general/ptq/iq2_xs.yaml @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# IQ2_XS weight-only PTQ. +# IQ2_XS weight-only PTQ and unified-checkpoint export. imports: preset: configs/ptq/presets/model/iq2_xs @@ -10,6 +10,7 @@ metadata: recipe_type: ptq description: >- Applies uniform GGML-compatible IQ2_XS weight-only quantization to eligible linear layers. - This is not a mixed per-tensor precision preset. No calibration data is required. + Unified export writes the packed block payload to each quantized Linear weight. This is not + a mixed per-tensor precision preset. No calibration data is required. quantize: $import: preset diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index b2a08308507..6d25ff1424c 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -138,11 +138,12 @@ activations and tensor-core math are what deliver the throughput. activations. Needs no calibration forward pass; the QAT starting point for the GPT-OSS family (see `examples/gpt-oss`). - **`iq1_s` / `iq2_xs`** — GGML-compatible IQ1_S or IQ2_XS weights on all linear - layers, with BF16 activations. No calibration data is required. Prefer weights whose final - dimension is divisible by 256: the fake-quant path pads other shapes with zeros, which - prevents reconstruction-cache reuse and can degrade the fit in the final partial block. These - recipes configure simulated weight quantization only; packed checkpoint export is added - separately. + layers, with BF16 activations. Unified export replaces each logical weight with a shaped + `uint8` tensor whose final axis contains 50 or 74 bytes per 256-value block; there is no + separate shape tensor. No calibration data is required. Every selected weight's final + dimension must be divisible by 256 for packed export. Exclude other layers: the fake-quant + path pads them with zeros, which prevents reconstruction-cache reuse and can degrade the fit + in the final partial block. --- diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py index 84244a8fec1..396dda8ce36 100644 --- a/tests/gpu/torch/export/test_layerwise_export.py +++ b/tests/gpu/torch/export/test_layerwise_export.py @@ -139,6 +139,24 @@ def _fp8_cfg(): return copy.deepcopy(mtq.FP8_DEFAULT_CFG) +def _iq2_cfg(): + return { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": { + "num_bits": "iq2_xs", + "block_sizes": {-1: 256}, + "backend": "ggml", + "backend_extra_args": {"search_impl": "auto"}, + }, + }, + ], + "algorithm": None, + } + + def _kv_cache_cfg(): return mtq.update_quant_cfg_with_kv_cache_quant( copy.deepcopy(mtq.FP8_DEFAULT_CFG), copy.deepcopy(mtq.FP8_KV_CFG["quant_cfg"]) @@ -180,6 +198,40 @@ def test_layerwise_export_resolves_mixed_kv_cache_formats(tmp_path): } +def test_iq2_layerwise_export_matches_whole_model_export(tmp_path): + """The per-layer path writes the same packed IQ tensors and metadata.""" + + def build_model(): + model = ( + get_tiny_llama( + hidden_size=256, + intermediate_size=256, + num_hidden_layers=2, + vocab_size=256, + ) + .cuda() + .eval() + ) + model.config.architectures = ["LlamaForCausalLM"] + return model + + baseline_dir = tmp_path / "baseline" + baseline = mtq.quantize(build_model(), _iq2_cfg()) + export_hf_checkpoint(baseline, export_dir=baseline_dir) + + layerwise_dir = tmp_path / "layerwise" + layerwise = mtq.quantize(build_model(), _iq2_cfg()) + exporter = LayerwiseExporter(layerwise, layerwise_dir) + layers = list(layerwise.model.layers) + exporter.bind(layers) + for layer_idx, layer in enumerate(layers): + exporter.export_layer(layer_idx, layer) + exporter.finalize() + + _assert_same_checkpoint(_load_checkpoint(baseline_dir), _load_checkpoint(layerwise_dir)) + _assert_same_quant_config(baseline_dir, layerwise_dir) + + def _nvfp4_cfg(): """NVFP4 with o_proj left unquantized. diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py index 2da25e90e15..a08372a4d15 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -45,7 +45,11 @@ import modelopt.torch.quantization as mtq import modelopt.torch.speculative as mtsp from modelopt.torch.export import KV_CACHE_FP8, export_mcore_gpt_to_hf, import_mcore_gpt_from_hf +from modelopt.torch.export.plugins.vllm_fakequant_megatron import VllmFqGPTModelExporter from modelopt.torch.export.unified_export_megatron import GPTModelExporter +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.ggml import dequantize_iq1_s, dequantize_iq2_xs +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 +90,341 @@ 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 = {} + + with patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1): + exporter._name_remapping(linear, "model.layers.0.mlp.down_proj.") + + packed_key = "model.layers.0.mlp.down_proj.weight" + assert exporter._state_dict[packed_key].shape == (2, 1, payload_bytes) + 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_megatron_iq_export_guard_uses_one_collective_for_all_flags(): + def report_remote_iq(flags, **_kwargs): + flags.copy_(torch.tensor([1, 0, 1], dtype=flags.dtype, device=flags.device)) + + with ( + patch.object(torch.distributed, "is_initialized", return_value=True), + patch.object(torch.distributed, "get_backend", return_value=torch.distributed.Backend.GLOO), + patch.object(torch.distributed, "all_reduce", side_effect=report_remote_iq), + ): + assert GPTModelExporter._collective_iq_export_flags(False, False, False) == ( + True, + False, + True, + ) + + +def test_megatron_iq_tensor_parallel_guard_finds_later_mixed_format(): + fp8 = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16) + fp8.weight_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3))) + iq = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16) + iq.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + ) + + assert GPTModelExporter._model_has_iq_quantizer(torch.nn.Sequential(fp8, iq)) + + +def test_megatron_iq_packed_expert_guard_finds_selected_projection(): + expert = torch.nn.Module() + expert.linear_fc1 = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16) + expert.linear_fc1.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + ) + model = torch.nn.Module() + model.local_experts = torch.nn.ModuleList([expert]) + + def packed_rule(*_args, **_kwargs): + pass + + setattr(packed_rule, "_modelopt_export_func_name", "pack_name_remapping") + exporter = object.__new__(GPTModelExporter) + exporter.model = model + exporter.rules = {"local_experts.linear_fc1": packed_rule} + + assert exporter._model_has_packed_expert_iq_quantizer() + + +def test_megatron_iq_tensor_parallel_guard_covers_extra_modules(): + exporter = object.__new__(GPTModelExporter) + exporter.model = torch.nn.Module() + + with ( + patch.object(uem, "get_tensor_model_parallel_world_size", return_value=2), + patch.object(exporter, "_model_has_iq_quantizer", return_value=True), + patch.object(exporter, "_collective_iq_export_flags", return_value=(True, False, False)), + pytest.raises(NotImplementedError, match="tensor model parallel size 1"), + ): + exporter.save_pretrained_extra_modules("unused") + + +def test_vllm_fakequant_export_skips_iq_packing_guard(): + exporter = object.__new__(VllmFqGPTModelExporter) + exporter.model = torch.nn.Module() + + with ( + patch.object(uem, "get_tensor_model_parallel_world_size", return_value=2), + patch.object(exporter, "_collective_iq_export_flags") as collective, + ): + exporter._validate_iq_export() + + collective.assert_not_called() + + +def test_megatron_non_iq_export_skips_shape_scan(): + exporter = object.__new__(GPTModelExporter) + exporter.model = torch.nn.Module() + + with ( + patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(exporter, "_model_has_iq_quantizer", return_value=False), + patch.object( + exporter, "_collective_iq_export_flags", return_value=(False, False, False) + ) as collective, + patch.object(uem, "_iq_export_weight_shape_errors") as shape_errors, + ): + exporter._validate_iq_export() + + collective.assert_called_once_with(False, False, False) + shape_errors.assert_not_called() + + +def test_megatron_iq_shape_guard_fails_all_ranks_before_packing(): + exporter = object.__new__(GPTModelExporter) + exporter.model = torch.nn.Linear(192, 2, bias=False, dtype=torch.bfloat16) + exporter.model.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + ) + + with ( + patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1), + patch.object( + exporter, "_collective_iq_export_flags", return_value=(True, True, False) + ) as collective, + pytest.raises(ValueError, match=r"weight: \(2, 192\)"), + ): + exporter._validate_iq_export() + + collective.assert_called_once_with(True, True, False) + + +def test_megatron_iq_shape_guard_reports_another_rank(): + exporter = object.__new__(GPTModelExporter) + exporter.model = torch.nn.Module() + + with ( + patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(exporter, "_collective_iq_export_flags", return_value=(True, True, False)), + pytest.raises(ValueError, match="Another distributed rank"), + ): + exporter._validate_iq_export() + + +def test_megatron_iq_config_guard_fails_collectively_before_conversion(): + exporter = object.__new__(GPTModelExporter) + exporter.model = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16) + exporter.model.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + ) + exporter.model.input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) + + with ( + patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1), + patch.object( + exporter, "_collective_iq_export_flags", return_value=(True, True, False) + ) as collective, + pytest.raises(ValueError, match="weight-only"), + ): + exporter._validate_iq_export() + + collective.assert_called_once_with(True, True, False) + + +def test_megatron_iq_packed_expert_guard_fails_all_ranks_before_conversion(): + exporter = object.__new__(GPTModelExporter) + exporter.model = torch.nn.Module() + + with ( + patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(exporter, "_model_has_iq_quantizer", return_value=False), + patch.object(exporter, "_model_has_packed_expert_iq_quantizer", return_value=False), + patch.object( + exporter, "_collective_iq_export_flags", return_value=(True, False, True) + ) as collective, + pytest.raises(NotImplementedError, match="packed expert weights"), + ): + exporter._validate_iq_export() + + collective.assert_called_once_with(False, False, False) + + +def test_megatron_iq_packer_rejects_tensor_parallelism(): + """Internal packing keeps the TP=1 guard for extra-module export paths.""" + module = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16) + module.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + ) + + with ( + patch.object(uem, "get_tensor_model_parallel_world_size", return_value=2), + pytest.raises(NotImplementedError, match="tensor model parallel size 1"), + ): + GPTModelExporter._get_iq_weight_state("weight", module.weight, "iq2_xs", module) + + +def test_megatron_iq_packer_error_identifies_state_dict_key(): + weight_key = "model.layers.0.mlp.down_proj.weight" + module = torch.nn.Linear(192, 2, bias=False, dtype=torch.bfloat16) + module.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + ) + + with ( + patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1), + pytest.raises(ValueError, match=weight_key), + ): + GPTModelExporter._get_iq_weight_state(weight_key, module.weight, "iq2_xs", module) + + +def test_megatron_iq_packer_revalidates_final_layout_owner(): + """The final packing boundary rejects activation state omitted from the checkpoint.""" + module = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16) + module.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + ) + module.input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) + + with ( + patch.object(uem, "get_tensor_model_parallel_world_size", return_value=1), + pytest.raises(NotImplementedError, match="weight-only"), + ): + GPTModelExporter._get_iq_weight_state("weight", module.weight, "iq2_xs", module) + + +@pytest.mark.parametrize("method_name", ["_pack_name_remapping", "_pack_name_remapping_gpt_oss"]) +def test_megatron_iq_export_rejects_packed_experts(method_name): + """Packed expert layouts cannot preserve the calibrated IQ block axis.""" + linear = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + ) + expert = torch.nn.Module() + expert.linear_fc1 = linear + exporter = object.__new__(GPTModelExporter) + exporter.dtype = torch.bfloat16 + exporter._state_dict = {} + exporter.exclude_modules = [] + exporter.layer_config_dict = {} + + with pytest.raises(NotImplementedError, match="packed expert weights"): + getattr(exporter, method_name)( + torch.nn.ModuleList([expert]), + "model.layers.0.mlp.experts.weight", + layer_type="linear_fc1", + ) + + assert exporter._state_dict == {} + + def _test_unified_export_megatron( tmp_path, model_type, @@ -713,6 +1052,39 @@ def test_grouped_mlp_slicing_maps_local_to_global_expert_ids(): assert "experts.1.gate_up_proj.weight" not in exporter._state_dict +def test_grouped_mlp_slicing_preserves_exclusion_prefix_and_formats_iq_label(): + exporter = _make_exporter_for_grouped_mlp() + calls = [] + exporter._get_quantized_state = lambda *a, **k: ( + calls.append((k["prefix"], k["describe_as"])) or {}, + None, + 0, + ) + module = _FakeTEGroupedMLP(num_gemms=2, local_expert_indices=[4, 5]) + + exporter._grouped_mlp_slicing(module, "experts.{}.gate_up_proj") + + assert calls == [ + ("experts.{}.gate_up_proj", "experts.4.gate_up_proj.weight"), + ("experts.{}.gate_up_proj", "experts.5.gate_up_proj.weight"), + ] + + +def test_vllm_fakequant_grouped_mlp_accepts_diagnostic_label(): + exporter = object.__new__(VllmFqGPTModelExporter) + exporter.dtype = torch.bfloat16 + exporter._state_dict = {} + exporter.exclude_modules = [] + exporter._get_weight_bias = lambda module, _dtype, _values: {"weight": module.weight} + exporter._get_weight_scales = lambda *_args: (None, None) + exporter._record_layer_quant_config = lambda *_args: None + module = _FakeTEGroupedMLP(num_gemms=1, local_expert_indices=[3]) + + exporter._grouped_mlp_slicing(module, "experts.{}.down_proj") + + assert "experts.3.down_proj.weight" in exporter._state_dict + + def test_grouped_mlp_slicing_normalizes_tensor_local_expert_indices(): """local_expert_indices may arrive as a torch.Tensor (Megatron path). It must be normalized to list[int] -- a naive `bool(tensor)` on a multi-element tensor raises. diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index 753c81a4b0e..0b04faf98c2 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -39,6 +39,8 @@ hide_quantizers_from_state_dict, ) from modelopt.torch.export.unified_export_hf import _postprocess_safetensors, export_hf_checkpoint +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer def _load_config(config_path): @@ -127,6 +129,33 @@ def _process_stub(*_args, **_kwargs): assert config_data["quantization_config"] == convert_hf_quant_config_format(dummy_quant_config) +def test_diffusers_iq_shape_preflight_runs_before_component_mutation(tmp_path, monkeypatch): + component = nn.Linear(192, 4, bias=False, dtype=torch.bfloat16) + component.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + ) + original_weight = component.weight.detach().clone() + + monkeypatch.setattr( + unified_export_hf, + "get_diffusion_components", + lambda *_args, **_kwargs: {"transformer": component}, + ) + monkeypatch.setattr(unified_export_hf, "has_quantized_modules", lambda *_args: True) + + with pytest.raises(ValueError, match=r"weight: \(4, 192\)"): + unified_export_hf._export_diffusers_checkpoint( + object(), torch.bfloat16, tmp_path, components=None + ) + + assert component.weight.dtype == torch.bfloat16 + assert torch.equal(component.weight, original_weight) + + def test_flux2_dummy_inputs_shape(): """Verify Flux2-specific dummy input shapes: 4-col RoPE ids, no pooled_projections, guidance.""" model = get_tiny_flux2() diff --git a/tests/unit/torch/export/test_export_registry.py b/tests/unit/torch/export/test_export_registry.py index 4c93388ca08..5a57dbdc79f 100644 --- a/tests/unit/torch/export/test_export_registry.py +++ b/tests/unit/torch/export/test_export_registry.py @@ -40,6 +40,8 @@ _ExportHandlerRegistryCls, ) from modelopt.torch.export.unified_export_hf import _process_quantized_modules +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer class _Experts(nn.Module): @@ -303,6 +305,22 @@ def test_process_quantized_modules_exports_via_registry(): assert weight.dtype == torch.float8_e4m3fn +def test_iq_handler_attributes_configuration_error_to_qualified_weight(): + linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + ) + linear.input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8, axis=None)) + ctx = ExportContext(model=linear, dtype=torch.bfloat16) + + with pytest.raises(NotImplementedError, match=r"model\.layers\.3\.mlp\.down_proj\.weight"): + _export_quant_linear("model.layers.3.mlp.down_proj", linear, ctx) + + def test_export_context_carries_no_resolver(): # Tied-weight dedup is driven by the driver's resolver (fed to sync_tied_input_amax / # postprocess_state_dict), not by handlers, so ExportContext no longer builds or holds diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index 6fc17d982e8..925cdeb5b1e 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -14,16 +14,31 @@ # limitations under the License. +from types import SimpleNamespace + import pytest import torch import torch.nn as nn from _test_utils.torch.export.utils import ToyModel, partial_fp8_config, partial_w4a8_config +import modelopt.torch.export.quant_utils as quant_utils import modelopt.torch.quantization as mtq +from modelopt.torch.export.quant_utils import ( + _pack_iq_weight, + _validate_iq_export_support, + postprocess_state_dict, +) from modelopt.torch.export.unified_export_hf import ( _export_quantized_weight, + _export_transformers_checkpoint, _process_quantized_modules, + export_hf_checkpoint, +) +from modelopt.torch.export.unified_export_hf_streaming import ( + _export_transformers_checkpoint_streaming, ) +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import GroupedQuantizer, TensorQuantizer from modelopt.torch.quantization.utils import quantizer_attr_names @@ -102,6 +117,451 @@ 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 isinstance(linear.weight, nn.Parameter) + assert "packed_weights" not in state_dict + assert "weight_shape" not in state_dict + + +@pytest.mark.parametrize(("num_bits", "payload_bytes"), [("iq1_s", 50), ("iq2_xs", 74)]) +def test_iq_packer_preserves_leading_dimensions(num_bits, payload_bytes): + weight = torch.randn(2, 3, 512, dtype=torch.bfloat16) + + packed = _pack_iq_weight(weight, num_bits) + + assert packed.shape == (2, 3, 2, payload_bytes) + + +def test_postprocess_state_dict_drops_legacy_weight_shape(): + state_dict = { + "layer.weight": torch.ones(2, 2), + "layer.weight_shape": torch.tensor([2, 2]), + } + + processed = postprocess_state_dict(state_dict, maxbound=448, quantization=None) + + assert set(processed) == {"layer.weight"} + + +@pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"]) +def test_export_iq_divisibility_error_identifies_weight(num_bits): + linear = nn.Linear(192, 4, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=num_bits, + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + ) + + with pytest.raises( + ValueError, + match=rf"Failed to pack {num_bits.upper()} weight 'Linear.weight' with shape \(4, 192\)", + ): + _export_quantized_weight(linear, torch.bfloat16) + + +@pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"]) +def test_export_iq_dtype_error_identifies_weight(num_bits): + weight = torch.ones((4, 256), dtype=torch.uint8) + + with pytest.raises( + TypeError, + match=rf"Failed to pack {num_bits.upper()} weight 'model.layers.0.weight'", + ): + _pack_iq_weight(weight, num_bits, describe_as="model.layers.0.weight") + + +def test_export_iq_payload_shape_error_identifies_weight(monkeypatch): + weight = torch.ones((4, 256), dtype=torch.bfloat16) + + def pack_with_wrong_shape(weight): + return torch.zeros((4, 2, 74), dtype=torch.uint8), None + + monkeypatch.setattr(quant_utils, "quantize_iq2_xs", pack_with_wrong_shape) + + with pytest.raises( + RuntimeError, + match=r"packing for weight 'model\.layers\.0\.weight'.*expected \(4, 1, 74\)", + ): + _pack_iq_weight(weight, "iq2_xs", describe_as="model.layers.0.weight") + + +def test_iq_packer_rejects_non_iq_format_before_inspecting_module(): + linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="Unsupported IQ quantization format: fp8"): + _pack_iq_weight(linear.weight, "fp8", module=linear) + + +def test_export_iq_shape_preflight_reports_every_incompatible_weight(): + model = nn.Sequential( + nn.Linear(192, 4, bias=False, dtype=torch.bfloat16), + nn.Linear(256, 4, bias=False, dtype=torch.bfloat16), + nn.Linear(320, 4, bias=False, dtype=torch.bfloat16), + ) + for linear in model: + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + ) + + with pytest.raises(ValueError) as exc_info: + _validate_iq_export_support(model) + + message = str(exc_info.value) + assert "0.weight: (4, 192)" in message + assert "2.weight: (4, 320)" in message + assert "1.weight" not in message + + +def test_export_iq_shape_preflight_skips_detailed_scan_without_iq(monkeypatch): + model = nn.Sequential(nn.Linear(32, 16), nn.ReLU()) + monkeypatch.setattr( + quant_utils, + "_iq_export_weight_shape_errors", + lambda _model: pytest.fail("non-IQ export must skip the detailed weight scan"), + ) + + _validate_iq_export_support(model) + + +def test_export_iq_shape_preflight_reports_grouped_weights(): + class GroupedLinear(nn.Module): + def __init__(self): + super().__init__() + self.num_gemms = 2 + self.weight0 = nn.Parameter(torch.ones(4, 256, dtype=torch.bfloat16)) + self.weight1 = nn.Parameter(torch.ones(4, 192, dtype=torch.bfloat16)) + quantizer_config = QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + self.weight_quantizer = GroupedQuantizer( + TensorQuantizer(quantizer_config), + TensorQuantizer(quantizer_config), + ) + + model = nn.Module() + model.experts = GroupedLinear() + + with pytest.raises(ValueError) as exc_info: + _validate_iq_export_support(model) + + message = str(exc_info.value) + assert "experts.weight1: (4, 192)" in message + assert "experts.weight0" not in message + + +def test_export_iq_shape_preflight_reports_grouped_quantizers_beyond_num_gemms(): + class GroupedLinear(nn.Module): + def __init__(self): + super().__init__() + self.num_gemms = 1 + self.weight0 = nn.Parameter(torch.ones(4, 256, dtype=torch.bfloat16)) + self.weight1 = nn.Parameter(torch.ones(4, 192, dtype=torch.bfloat16)) + quantizer_config = QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + self.weight_quantizer = GroupedQuantizer( + TensorQuantizer(quantizer_config), + TensorQuantizer(quantizer_config), + ) + + model = nn.Module() + model.experts = GroupedLinear() + + with pytest.raises(ValueError, match=r"experts\.weight1: \(4, 192\)"): + _validate_iq_export_support(model) + + +def test_export_iq_shape_preflight_reports_legacy_fused_expert_quantizers(): + class LegacyFusedExperts(nn.Module): + def __init__(self): + super().__init__() + self._first_proj_attr = "up_proj" + self.up_proj = nn.Parameter(torch.ones(4, 192, dtype=torch.bfloat16)) + self.gate_up_proj_weight_quantizers = nn.ModuleList( + [ + TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + ) + ] + ) + + model = nn.Module() + model.experts = LegacyFusedExperts() + + with pytest.raises(ValueError, match=r"experts\.up_proj: \(4, 192\)"): + _validate_iq_export_support(model) + + +def test_export_iq_preflight_rejects_nonstandard_weight_before_mutation(tmp_path): + class CustomWeightModule(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Parameter(torch.ones(4, 256, dtype=torch.bfloat16)) + self.proj_weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + ) + + model = nn.Module() + model.block = CustomWeightModule() + model.config = SimpleNamespace(torch_dtype=torch.bfloat16) + original_weight = model.block.proj.detach().clone() + + with pytest.raises(ValueError, match=r"block\.proj: nonstandard weight"): + export_hf_checkpoint(model, export_dir=tmp_path) + + assert model.block.proj.dtype == torch.bfloat16 + assert torch.equal(model.block.proj, original_weight) + + with pytest.raises(NotImplementedError, match=r"got 'model\.block\.proj'"): + _export_quantized_weight( + model.block, + torch.bfloat16, + weight_name="proj", + describe_as="model.block.proj", + ) + + +def test_export_hf_checkpoint_runs_iq_shape_preflight_before_mutation(tmp_path): + linear = nn.Linear(192, 4, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + ) + linear.config = SimpleNamespace(torch_dtype=torch.bfloat16) + original_weight = linear.weight.detach().clone() + + with pytest.raises(ValueError, match=r"weight: \(4, 192\)"): + export_hf_checkpoint(linear, export_dir=tmp_path) + + assert linear.weight.dtype == torch.bfloat16 + assert torch.equal(linear.weight, original_weight) + + +def test_export_transformers_checkpoint_runs_iq_shape_preflight_before_mutation(): + model = nn.Sequential( + nn.Linear(256, 4, bias=False, dtype=torch.bfloat16), + nn.Linear(192, 4, bias=False, dtype=torch.bfloat16), + ) + model.config = SimpleNamespace(torch_dtype=torch.bfloat16) + for linear in model: + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + ) + original_weights = [linear.weight.detach().clone() for linear in model] + + with pytest.raises(ValueError, match=r"1\.weight: \(4, 192\)"): + _export_transformers_checkpoint(model) + + for linear, original_weight in zip(model, original_weights): + assert linear.weight.dtype == torch.bfloat16 + assert torch.equal(linear.weight, original_weight) + + +def test_export_transformers_checkpoint_runs_iq_config_preflight_before_mutation(): + model = nn.Sequential( + nn.Linear(256, 4, bias=False, dtype=torch.bfloat16), + nn.Linear(256, 4, bias=False, dtype=torch.bfloat16), + ) + model.config = SimpleNamespace(torch_dtype=torch.bfloat16) + for linear in model: + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + ) + model[1].input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) + original_weights = [linear.weight.detach().clone() for linear in model] + + with pytest.raises(ValueError, match=r"weight-only.*1\.weight"): + _export_transformers_checkpoint(model) + + for linear, original_weight in zip(model, original_weights): + assert linear.weight.dtype == torch.bfloat16 + assert torch.equal(linear.weight, original_weight) + + +def test_export_transformers_checkpoint_streaming_runs_iq_shape_preflight_before_mutation( + tmp_path, +): + linear = nn.Linear(192, 4, bias=False, dtype=torch.bfloat16) + linear.config = SimpleNamespace(torch_dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + ) + original_weight = linear.weight.detach().clone() + + with pytest.raises(ValueError, match=r"weight: \(4, 192\)"): + _export_transformers_checkpoint_streaming(linear, export_dir=tmp_path) + + assert linear.weight.dtype == torch.bfloat16 + assert torch.equal(linear.weight, original_weight) + assert not any(tmp_path.iterdir()) + + +@pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"]) +def test_export_iq_rejects_an_unhandled_search_impl(num_bits): + linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=num_bits, + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "reference"}, + ) + ) + + with pytest.raises(NotImplementedError, match="only supports search_impl='auto'"): + _export_quantized_weight(linear, torch.bfloat16) + + +@pytest.mark.parametrize( + ("quantizer_name", "mode"), + [ + ("input_quantizer", "enabled"), + ("input_quantizer", "pre_quant_scale"), + ("output_quantizer", "enabled"), + ], +) +@pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"]) +def test_export_iq_rejects_non_weight_only_config(num_bits, quantizer_name, mode): + linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=num_bits, + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + ) + activation_quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=8, axis=None, enable=mode == "enabled") + ) + setattr(linear, quantizer_name, activation_quantizer) + if mode == "pre_quant_scale": + activation_quantizer.pre_quant_scale = torch.ones(256) + + with pytest.raises(NotImplementedError, match="export is weight-only"): + _export_quantized_weight(linear, torch.bfloat16) + + +def test_export_iq_config_preflight_reports_weight_owner(): + linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + ) + linear.input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) + model = nn.Module() + model.block = linear + + errors = quant_utils._iq_export_quantizer_config_errors(model) + + assert len(errors) == 1 + assert "block.weight" in errors[0] + assert "weight-only" in errors[0] + + +def test_export_iq_config_preflight_uses_unprefixed_grouped_activation_quantizer(): + class GroupedLinear(nn.Module): + def __init__(self): + super().__init__() + self.num_gemms = 2 + self.weight0 = nn.Parameter(torch.ones(4, 256, dtype=torch.bfloat16)) + self.weight1 = nn.Parameter(torch.ones(4, 256, dtype=torch.bfloat16)) + quantizer_config = QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + self.weight_quantizer = GroupedQuantizer( + TensorQuantizer(quantizer_config), + TensorQuantizer(quantizer_config), + ) + self.input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) + + model = nn.Module() + model.experts = GroupedLinear() + + errors = quant_utils._iq_export_quantizer_config_errors(model) + + assert len(errors) == 2 + assert all("weight-only" in error for error in errors) + assert any("experts.weight0" in error for error in errors) + assert any("experts.weight1" in error for error in errors) + + +def test_iq_config_error_identifies_weight(): + linear = nn.Linear(256, 4, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + ) + ) + linear.input_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8, axis=None)) + + with pytest.raises(NotImplementedError, match=r"model\.layers\.3\.mlp\.down_proj\.weight"): + _pack_iq_weight( + linear.weight, + "iq2_xs", + module=linear, + describe_as="model.layers.3.mlp.down_proj.weight", + ) + + class QuantMoELinear(nn.Module): def __init__(self): super().__init__() diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index 8b9670e5576..382b646a027 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,50 @@ def __init__(self): self.v_bmm_quantizer = TensorQuantizer() +@pytest.mark.parametrize( + ("num_bits", "quantization_format", "payload_bytes", "storage_bits"), + [ + ("iq1_s", QUANTIZATION_IQ1_S, 50, 1), + ("iq2_xs", QUANTIZATION_IQ2_XS, 74, 2), + ], +) +def test_iq_quantization_config(num_bits, quantization_format, payload_bytes, storage_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 == { + "dynamic": False, + "num_bits": storage_bits, + "type": "int", + "group_size": 256, + "packing": "ggml", + "block_payload_bytes": payload_bytes, + } + + 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..402130942fc 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -15,6 +15,8 @@ """Tests for _QuantFusedExperts: generic fused MoE quantization and export.""" +import warnings + import pytest import torch import torch.nn as nn @@ -28,6 +30,7 @@ import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module from modelopt.torch.export.moe_utils import _export_fused_experts from modelopt.torch.export.quant_utils import get_quant_config, get_quantization_format +from modelopt.torch.export.unified_export_hf import export_hf_checkpoint from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.conversion import _normalize_fused_experts_quantizer_name from modelopt.torch.quantization.model_calib import local_hessian_calibrate @@ -502,6 +505,148 @@ def forward_loop(m): self._cleanup_registry(expert_type) + def test_export_preserves_packed_iq_weights(self, monkeypatch): + """Packed expert weights must remain present in the exported state dict.""" + model = _TinyMoEModel() + expert_type = type(model.moe.experts) + self._cleanup_registry(expert_type) + register_fused_experts_on_the_fly(model) + + try: + converted = QuantModuleRegistry.convert(model.moe.experts) + config = QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + quantizers = list(converted.gate_up_proj_weight_quantizers) + list( + converted.down_proj_weight_quantizers + ) + for quantizer in quantizers: + quantizer.set_from_attribute_config(config) + + labels = [] + + def fake_pack(weight, *_args): + labels.append(_args[-1]) + return torch.zeros( + (weight.numel() // 256, 74), dtype=torch.uint8, device=weight.device + ) + + monkeypatch.setattr( + "modelopt.torch.export.unified_export_hf._pack_iq_weight", fake_pack + ) + + _export_fused_experts( + converted, torch.float16, module_name="model.layers.0.mlp.experts" + ) + + 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 + assert isinstance( + getattr(getattr(converted, str(idx)), projection).weight, nn.Parameter + ) + assert labels == [ + f"model.layers.0.mlp.experts.{idx}.{projection}.weight" + for idx in range(NUM_EXPERTS) + for projection in ("gate_proj", "up_proj", "down_proj") + ] + finally: + self._cleanup_registry(expert_type) + + def test_export_hf_checkpoint_preflight_allows_fused_iq_weights(self, monkeypatch, tmp_path): + """The model-wide preflight must preserve the fused-expert split path.""" + model = _TinyMoEModel() + expert_type = type(model.moe.experts) + self._cleanup_registry(expert_type) + register_fused_experts_on_the_fly(model) + + try: + converted = QuantModuleRegistry.convert(model.moe.experts) + converted.gate_up_proj = nn.Parameter(torch.ones(NUM_EXPERTS, 512, 256)) + converted.down_proj = nn.Parameter(torch.ones(NUM_EXPERTS, 256, 256)) + config = QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + quantizers = list(converted.gate_up_proj_weight_quantizers) + list( + converted.down_proj_weight_quantizers + ) + for quantizer in quantizers: + quantizer.set_from_attribute_config(config) + + export_called = False + + def fake_export(*_args, **_kwargs): + nonlocal export_called + export_called = True + return {}, {} + + monkeypatch.setattr( + "modelopt.torch.export.unified_export_hf._export_transformers_checkpoint", + fake_export, + ) + monkeypatch.setattr( + model, "save_pretrained", lambda *_args, **_kwargs: None, raising=False + ) + monkeypatch.setattr( + "modelopt.torch.export.unified_export_hf._write_hf_export_config", + lambda *_args, **_kwargs: None, + ) + + export_hf_checkpoint(model, export_dir=tmp_path) + + assert export_called + finally: + self._cleanup_registry(expert_type) + + def test_iq_export_does_not_create_or_warn_about_amax(self, monkeypatch): + """IQ expert packing is weight-only and must not enter amax fallback paths.""" + model = _TinyMoEModel() + expert_type = type(model.moe.experts) + self._cleanup_registry(expert_type) + register_fused_experts_on_the_fly(model) + + try: + converted = QuantModuleRegistry.convert(model.moe.experts) + quantizers = list(converted.gate_up_proj_weight_quantizers) + list( + converted.down_proj_weight_quantizers + ) + config = QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + for quantizer in quantizers: + quantizer.set_from_attribute_config(config) + assert not hasattr(quantizer, "_amax") + + monkeypatch.setattr( + "modelopt.torch.export.unified_export_hf._export_quantized_weight", + lambda *_args, **_kwargs: None, + ) + + def reject_copy(*_args, **_kwargs): + raise AssertionError("IQ export must reuse its source quantizer") + + monkeypatch.setattr("modelopt.torch.export.moe_utils.deepcopy", reject_copy) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _export_fused_experts(converted, torch.float16) + + assert not [warning for warning in caught if "amax" in str(warning.message)] + assert all(not hasattr(quantizer, "_amax") for quantizer in quantizers) + finally: + self._cleanup_registry(expert_type) + def test_uncalibrated_expert_gate_up_share_amax(self, monkeypatch): """gate_proj and up_proj must share weight_scale_2 even when an expert was never routed during calibration.