fix(export): unblock AWQ export for fused MoE experts (Nemotron-H)

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].<linear>``), 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)
 ^^^^^^^^^^^^^^^^^
 

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. "

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].<linear>``), 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 ``<proj>_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)

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