From 1efcd42f55d4a267c2a67d85fe95a78fdb33b358 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 24 Jul 2026 19:39:37 +0000 Subject: [PATCH] fix(export): unblock INT4 AWQ export for fused MoE experts (NemotronH) int4_awq PTQ on NVIDIA-Nemotron-3-Super-120B-A12B-BF16 could not produce a checkpoint. Calibration succeeded; export failed. Two independent defects, both only reachable on AWQ formats, and both specific to the fused-experts layout that transformers >= 5.5 uses for NemotronH (NemotronHExperts -> 3-D packed params, wrapped as _QuantNonGatedFusedExperts). FP8/NVFP4 were unaffected. 1) requantize_resmooth_fused_llm_layers() called get_experts_list() for every MoE block, gated on "awq" in quantization_format. That helper addresses experts positionally (module.experts[i].), which only exists for the sequential nn.ModuleList layout, so a fused wrapper raised "TypeError: object of type 'QuantNemotronHExperts' has no len()". Fused blocks are now skipped: resmoothing exists to give every expert in a group one common input pre_quant_scale, and a fused wrapper already shares a single input quantizer across all its experts, so there is nothing to reconcile. 2) With that unblocked, _export_fused_experts()'s uncalibrated-expert fallback assigned a scalar amax. That is only valid for per-tensor formats; for block-wise formats (INT4 AWQ, block_sizes={-1: 128}) the exporter then evaluated weight.shape[-1] // weights_scaling_factor.shape[-1] on a 0-d tensor and raised "IndexError: tuple index out of range". The fallback now derives a per-block amax shaped (out_features, n_blocks), matching what MaxCalibrator produces for a calibrated expert. This path is easy to hit on this model: 512 routed experts at top-22 leaves many experts unrouted for any modest calibration set. Verified end-to-end on 4xB200: hf_ptq.py --qformat int4_awq on the full 120B now writes a complete checkpoint (8 shards, 68G, W4A16_AWQ group_size=128), with expert weight_scale shapes (2688, 8) and (1024, 21) matching 1024/128 and 2688/128 blocks. Adds a regression test that fails on the pre-fix fallback. Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 2 + modelopt/torch/export/moe_utils.py | 29 ++++++++- modelopt/torch/export/unified_export_hf.py | 29 +++++++-- .../plugins/test_fused_experts.py | 64 +++++++++++++++++++ 4 files changed, 119 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 329e5d21f05..1cb367403af 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -12,6 +12,8 @@ Changelog **Bug Fixes** +- Fix AWQ (e.g. ``int4_awq``, ``w4a8_awq``, ``nvfp4_awq``) HF checkpoint export for models whose MoE experts are stored as fused 3-D parameters — Nemotron-H on ``transformers>=5.5`` (``NemotronHExperts``), and any other model wrapped by ``_QuantFusedExperts`` / ``_QuantNonGatedFusedExperts``. Calibration succeeded but export raised ``TypeError: object of type 'QuantNemotronHExperts' has no len()``, because the AWQ-only expert resmoothing step addressed experts positionally (``module.experts[i].``), which only exists for the sequential ``nn.ModuleList`` layout; fused blocks are now skipped, as a fused wrapper already shares one input quantizer (and therefore one ``pre_quant_scale``) across all of its experts. Past that, the uncalibrated-expert fallback in ``_export_fused_experts`` assigned a scalar ``amax``, which is only valid for per-tensor formats — block-wise formats then evaluated ``weight.shape[-1] // weights_scaling_factor.shape[-1]`` on a 0-d tensor and raised ``IndexError: tuple index out of range``. It now derives a per-block ``amax`` shaped ``(out_features, n_blocks)``, matching what ``MaxCalibrator`` produces for a calibrated expert. Experts left unrouted by the calibration set are common on wide MoEs (e.g. 512 experts at top-22), so both paths were easy to hit. Non-AWQ formats (FP8, NVFP4) were unaffected. + 0.46 (2026-08-17) ^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 787e173959e..466b3a29315 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -74,6 +74,33 @@ def _delete_fused_moe_source_attrs(module: nn.Module) -> None: delattr(module, attr) +def _weight_derived_amax(quantizer: nn.Module, weight: torch.Tensor) -> torch.Tensor: + """Amax for an uncalibrated expert weight, at the quantizer's own granularity. + + A scalar (0-d) amax is only meaningful for per-tensor formats. Block-wise formats -- + e.g. INT4 AWQ with ``block_sizes={-1: 128}`` -- need one amax per block, matching what + ``MaxCalibrator`` produces for a calibrated expert. Handing them a 0-d amax makes the + exporter compute ``weight.shape[-1] // weights_scaling_factor.shape[-1]`` on a 0-d + scale and raise "IndexError: tuple index out of range", which is how a MoE expert that + the router never picked during calibration used to break INT4 AWQ export. + """ + block_sizes = getattr(quantizer, "block_sizes", None) + block_size = None + if isinstance(block_sizes, dict): + # Block size on the last (input) axis; recorded under -1 or the positive axis index. + block_size = block_sizes.get(-1, block_sizes.get(weight.dim() - 1)) + if ( + not isinstance(block_size, int) + or block_size <= 0 + or weight.dim() < 2 + or weight.shape[-1] % block_size != 0 + ): + return weight.abs().amax().to(torch.float32) + # (out_features, n_blocks) -- the same layout the per-channel slicing above normalises + # a calibrated block amax into, and what the exporter divides the weight columns by. + return weight.abs().reshape(weight.shape[0], -1, block_size).amax(dim=-1).to(torch.float32) + + def _export_fused_experts( module: nn.Module, dtype: torch.dtype, @@ -258,7 +285,7 @@ def _export_fused_experts( or torch.all(w_quantizer._amax == 0) ) ): - w_quantizer.amax = weight_slice.abs().amax().to(torch.float32) + w_quantizer.amax = _weight_derived_amax(w_quantizer, weight_slice) warnings.warn( f"Expert {idx} {proj_name} weight quantizer was not calibrated " f"(amax missing or zero). Using weight-derived amax as fallback. " diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index cee64c22c05..cb70967814f 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -65,7 +65,7 @@ export_sparse_attention_config = None # Importing the built-in handlers installs their entries in the two registries. -from . import hf_export_handlers as _hf_export_handlers # noqa: F401 +from . import hf_export_handlers as _hf_export_handlers from .convert_hf_config import convert_hf_quant_config_format from .layer_utils import ( get_experts_list, @@ -425,6 +425,20 @@ def _fuse_shared_input_modules( return fused_linears +def _has_fused_experts(module: torch.nn.Module) -> bool: + """Return True when an MoE block stores its experts as fused 3-D parameters. + + ``get_experts_list`` / ``_get_expert_attr`` address experts positionally + (``module.experts[i].``), which only exists for the sequential + ``nn.ModuleList`` layout. Fused wrappers (e.g. NemotronH ``NemotronHExperts`` on + transformers >= 5.5) pack every expert into 3-D parameters and expose per-expert + quantizers as ``_weight_quantizers`` instead, so they have no ``__len__`` and + no per-expert submodules to group. + """ + experts = getattr(module, "experts", None) + return experts is not None and _hf_export_handlers._has_fused_experts_quantizers(experts) + + def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): """Group modules that take the same input and register shared parameters in module.""" # TODO: Handle DBRX MoE @@ -444,9 +458,16 @@ def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): for name, module in model.named_modules(): module_names.add(name) - # For MoE models update pre_quant_scale to average pre_quant_scale amongst experts - if is_moe(module) and ( - quantization_format is not QUANTIZATION_NONE + # For MoE models update pre_quant_scale to average pre_quant_scale amongst experts. + # Fused-experts blocks are skipped: resmoothing exists to give every expert in a group + # one common input pre_quant_scale, but a fused wrapper already shares a single input + # quantizer (and hence a single pre_quant_scale) across all its experts, so there is + # nothing to reconcile. Without the guard, get_experts_list() below raises + # "object of type 'QuantNemotronHExperts' has no len()". + if ( + is_moe(module) + and not _has_fused_experts(module) + and quantization_format is not QUANTIZATION_NONE and ("awq" in quantization_format or quantization_format == QUANTIZATION_NVFP4_SVDQUANT) ): # update_experts_avg_prequant_scale(module) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 9fa836bb620..8a726e7b7ec 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -636,6 +636,70 @@ def _spy_export(wrapper, dtype, **_kwargs): if QuantModuleRegistry.get(expert_type) is not None: QuantModuleRegistry.unregister(expert_type) + def test_uncalibrated_expert_block_quant_amax_is_per_block(self, monkeypatch): + """An uncalibrated expert under a block-wise format needs a per-block amax. + + Regression for INT4 AWQ export of NemotronH (NVBug: "cannot run int4_awq PTQ"). + With 512 routed experts at top-22 a small calibration set leaves some experts + unrouted, so their weight quantizers reach export with ``_amax is None``. The + fallback used to assign a scalar ``weight.abs().amax()``, which is only valid + for per-tensor formats; for ``block_sizes={-1: N}`` the exporter then evaluates + ``weight.shape[-1] // weights_scaling_factor.shape[-1]`` on a 0-d tensor and + raised ``IndexError: tuple index out of range``. + """ + block_size = HIDDEN_DIM // 2 + experts = _SyntheticNonGatedFusedExperts() + expert_type = type(experts) + if QuantModuleRegistry.get(expert_type) is None: + QuantModuleRegistry.register({expert_type: "test.SyntheticNonGatedFusedExperts"})( + _QuantNonGatedFusedExperts + ) + try: + converted = QuantModuleRegistry.convert(experts) + + # Enable every expert weight quantizer as block-quantized, and leave them + # all uncalibrated so the export-time fallback is what supplies the amax. + block_cfg = QuantizerAttributeConfig( + num_bits=4, block_sizes={-1: block_size, "type": "static"} + ) + for quantizers in ( + converted.up_proj_weight_quantizers, + converted.down_proj_weight_quantizers, + ): + for q in quantizers: + q.set_from_attribute_config(block_cfg) + q._disabled = False + + seen = [] + + def _spy_export(wrapper, dtype, **_kwargs): + seen.append( + (tuple(wrapper.weight.shape), wrapper.weight_quantizer._amax.detach().clone()) + ) + + monkeypatch.setattr( + "modelopt.torch.export.unified_export_hf._export_quantized_weight", + _spy_export, + ) + + _export_fused_experts(converted, torch.float16) + + assert seen, "no projections were exported" + for weight_shape, amax in seen: + out_features, in_features = weight_shape + assert amax.dim() == 2, ( + f"uncalibrated block-quantized expert got a {amax.dim()}-d amax " + f"{tuple(amax.shape)}; block formats need one amax per block" + ) + assert tuple(amax.shape) == (out_features, in_features // block_size), ( + f"expected per-block amax {(out_features, in_features // block_size)}, " + f"got {tuple(amax.shape)}" + ) + assert torch.all(amax > 0), "weight-derived amax must be positive" + finally: + if QuantModuleRegistry.get(expert_type) is not None: + QuantModuleRegistry.unregister(expert_type) + # --------------------------------------------------------------------------- # Tests for tied-experts dedup in _export_fused_experts