[OMNIML-5003] Support non-gated fused MoE experts (NemotronH) in HF PTQ

CHANGELOG.rst

@@ -111,6 +111,7 @@ Changelog
 
 **Bug Fixes**
 
+- Support non-gated fused MoE experts in unified HuggingFace export. Nemotron-H MoE models (transformers 5.x ``NemotronHExperts``) store experts as fused 3-D ``up_proj`` / ``down_proj`` parameters with no ``gate_up_proj``; the fused-experts detection previously keyed on ``gate_up_proj``, so these were never wrapped as ``_QuantFusedExperts`` and export raised ``NotImplementedError: MoE model with experts type 'NemotronHExperts' is not supported``. The fused-experts path now also recognizes the non-gated layout (new ``_QuantNonGatedFusedExperts``) and exports a single ``up_proj`` per expert; the gated path is unchanged.
 - In Megatron-Core only do EP amax sync for routed expert weights if ``sync_expert_weight_amax=True``. Previously EP amax sync would sync routed expert weights across EP ranks even when ``sync_expert_weight_amax`` was False.
 - Fix Megatron-Core HF importer to load fused ``TELayerNormColumnParallelLinear.layer_norm_weight`` from HF for GPT-family models (Qwen3 etc.) under ``--export-default-te-spec``. Importer now prefers per-context keys ``fused_input_layernorm`` / ``fused_pre_mlp_layernorm`` (fallback ``fused_norm`` for Nemotron-H backward compatibility); ``mcore_qwen.py`` provides the new rules. Without this fix, post-prune MMLU sat at chance.
 - Fix ONNX AutoCast ``keep_io_types=True`` sanity-check failure (``Unexpected type in I/O tensor ...``) when a network input/output is an empty tensor (a dimension of size 0). Such tensors were "fake-cast" (retyped in place) to the low precision type; because the value-info map aliases the ``graph.input``/``graph.output`` ``ValueInfoProto``, this silently changed the model's I/O type. AutoCast now inserts a real ``Cast`` for protected I/O tensors instead.

modelopt/torch/export/layer_utils.py

@@ -987,8 +987,10 @@ def module_match_name_list(module, name_list):
     # Structural detection: after _export_fused_experts, fused expert modules
     # have per-expert submodules with gate_proj/up_proj/down_proj.
     # Also handles models that originally used this naming (Qwen, DeepSeek, etc.).
-    if hasattr(module, "experts") and hasattr(module.experts, "gate_up_proj_weight_quantizers"):
-        return ["gate_up_proj", "down_proj"]
+    if hasattr(module, "experts"):
+        first_proj_attr = getattr(module.experts, "_first_proj_attr", "gate_up_proj")
+        if hasattr(module.experts, f"{first_proj_attr}_weight_quantizers"):
+            return [first_proj_attr, "down_proj"]
 
     if module_match_name_list(
         module,
@@ -1004,7 +1006,7 @@ def module_match_name_list(module, name_list):
     elif module_match_name_list(module, ["MixtralSparseMoeBlock"]):
         # Old-style Mixtral (iterable experts) uses w1/w2/w3.
         # Fused Mixtral (transformers 5.0+) is already handled by the
-        # structural gate_up_proj_weight_quantizers check above.
+        # structural first-projection quantizer check above.
         return ["w1", "w2", "w3"]
     elif module_match_name_list(module, ["MixtralMoeSparseMoeBlock"]):
         # Older transformers naming for Mixtral

modelopt/torch/export/moe_utils.py

@@ -59,11 +59,14 @@ def _delete_fused_moe_source_attrs(module: nn.Module) -> None:
     aliases or via the full unpack/pack path) so the redundant fused form
     doesn't appear in the exported state_dict alongside the per-expert form.
     """
+    first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj")
+    first_proj_weight_quantizers_attr = f"{first_proj_attr}_weight_quantizers"
+    first_proj_input_quantizer_attr = f"{first_proj_attr}_input_quantizer"
     for attr in (
-        "gate_up_proj",
+        first_proj_attr,
+        first_proj_weight_quantizers_attr,
+        first_proj_input_quantizer_attr,
         "down_proj",
-        "gate_up_proj_weight_quantizers",
-        "gate_up_proj_input_quantizer",
         "down_proj_weight_quantizers",
         "down_proj_input_quantizer",
     ):
@@ -79,28 +82,29 @@ def _export_fused_experts(
 ) -> None:
     """Split fused MoE expert weights and export per-expert quantization scales.
 
-    Works with any module wrapped by ``_QuantFusedExperts`` — i.e. any HF
-    transformers 5.0+ fused expert container that stores ``gate_up_proj`` and
-    ``down_proj`` as 3-D ``nn.Parameter`` tensors with per-expert quantizer
-    ``nn.ModuleList`` s.
+    Works with any module wrapped by ``_QuantFusedExperts`` (gated, with a fused
+    ``gate_up_proj``) or ``_QuantNonGatedFusedExperts`` (non-gated, with a single
+    ``up_proj`` — e.g. NemotronH). Both store their projections as 3-D
+    ``nn.Parameter`` tensors with per-expert quantizer ``nn.ModuleList`` s.
 
     Steps:
 
-    1. Handle amax fallback for uncalibrated expert input quantizers.
-    2. Split fused 3-D weights into per-expert 2-D projections
-       (``gate_proj``, ``up_proj``, ``down_proj``).
+    1. Handle amax fallback for uncalibrated expert weight quantizers.
+    2. Split fused 3-D weights into per-expert 2-D projections — gated:
+       (``gate_proj``, ``up_proj``, ``down_proj``); non-gated: (``up_proj``,
+       ``down_proj``).
     3. Call ``_export_quantized_weight`` on each projection.
     4. Register results under the standard naming convention::
 
-           {E}.gate_proj.weight, {E}.gate_proj.weight_scale, ...
+           {E}.gate_proj.weight, {E}.gate_proj.weight_scale, ...  # gated only
            {E}.up_proj.weight, {E}.up_proj.weight_scale, ...
            {E}.down_proj.weight, {E}.down_proj.weight_scale, ...
 
     Tied-experts dedup is opt-in via ``_moe_tied_cache``: when multiple
     fused-expert modules share their 3-D source params via HF
     ``_tied_weights_keys``, the unpacking creates fresh per-expert tensors
     that break the tie. With ``_moe_tied_cache`` provided (tuple-keyed by
-    ``(gate_up_proj.data_ptr(), down_proj.data_ptr())``), the alias step
+    ``(<first_proj>.data_ptr(), down_proj.data_ptr())``), the alias step
     at the end re-points the per-expert ``weight`` / ``weight_scale`` /
     ``weight_scale_2`` / ``input_scale`` buffers at a previously-processed
     module sharing the same source memory. ``_tied_cache`` (int-keyed) is
@@ -114,13 +118,20 @@ def _export_fused_experts(
     from modelopt.torch.quantization.plugins.huggingface import _get_fused_expert_intermediate_dim
 
     n = module.num_experts
-    expert_dim = _get_fused_expert_intermediate_dim(module)
+    # Gated experts fuse gate+up into ``gate_up_proj`` and must be split on export;
+    is_gated = getattr(module, "_is_gated", True)
+    first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj")
+    # Only the gated split needs the per-expert intermediate dim (gate|up boundary).
+    expert_dim = _get_fused_expert_intermediate_dim(module) if is_gated else None
 
     # Capture source tensor identities BEFORE unpacking (the source
     # attrs are deleted at the end of this function).
-    _source_key = (module.gate_up_proj.data_ptr(), module.down_proj.data_ptr())
+    _source_key = (
+        getattr(module, first_proj_attr).data_ptr(),
+        module.down_proj.data_ptr(),
+    )
 
-    # Tied-experts fast path: if this exact (gate_up, down) source-tensor pair
+    # Tied-experts fast path: if this exact (first_proj, down) source-tensor pair
     # has been processed before, alias all per-expert buffers directly from the
     # prior module — no unpacking, no per-expert packing, no transient buffers
     # thrown away. Cache miss falls through to the full unpack/pack below and
@@ -133,14 +144,15 @@ def _export_fused_experts(
             return
 
     # 1. Shared input quantizers — one per projection type, shared across all experts.
-    gate_up_input_q = module.gate_up_proj_input_quantizer
+    first_proj_input_q = getattr(module, f"{first_proj_attr}_input_quantizer")
+    first_proj_weight_quantizers = getattr(module, f"{first_proj_attr}_weight_quantizers")
     down_input_q = module.down_proj_input_quantizer
 
-    gate_up = module.gate_up_proj.data
+    first_proj = getattr(module, first_proj_attr).data  # gate_up_proj or up_proj
     down = module.down_proj.data
 
     # 2-3. Split + export each per-expert projection.
-    fused_dim0 = gate_up.shape[1]  # 2 * expert_dim
+    fused_dim0 = first_proj.shape[1]  # gated: 2 * expert_dim; non-gated: expert_dim
 
     for idx in range(n):
         expert = nn.Module()
@@ -153,13 +165,19 @@ def _export_fused_experts(
         # fallback further down would otherwise compute amax independently from
         # each half — gate's max and up's max generally differ — producing
         # mismatched weight_scale_2 and garbled MoE output at inference.
-        gate_up_q = module.gate_up_proj_weight_quantizers[idx]
-        if getattr(gate_up_q, "is_enabled", False) and (
-            not hasattr(gate_up_q, "_amax")
-            or gate_up_q._amax is None
-            or torch.all(gate_up_q._amax == 0)
+        # Non-gated experts have no gate/up fusion, so this shared-amax step is
+        # skipped — their single up_proj uses the generic per-projection fallback.
+        first_proj_q = first_proj_weight_quantizers[idx]
+        if (
+            is_gated
+            and getattr(first_proj_q, "is_enabled", False)
+            and (
+                not hasattr(first_proj_q, "_amax")
+                or first_proj_q._amax is None
+                or torch.all(first_proj_q._amax == 0)
+            )
         ):
-            gate_up_q.amax = gate_up[idx].abs().amax().to(torch.float32)
+            first_proj_q.amax = first_proj[idx].abs().amax().to(torch.float32)
             warnings.warn(
                 f"Expert {idx} gate_up_proj weight quantizer was not calibrated "
                 f"(amax missing or zero). Using fused-tensor amax as fallback "
@@ -168,22 +186,38 @@ def _export_fused_experts(
                 stacklevel=2,
             )
 
-        projections = [
-            ("gate_proj", gate_up[idx, :expert_dim, :], 0, fused_dim0, True),
-            ("up_proj", gate_up[idx, expert_dim:, :], expert_dim, fused_dim0, True),
-            ("down_proj", down[idx], 0, down.shape[1], False),
-        ]
-
-        for proj_name, weight_slice, fused_start, fused_total, is_gate_up in projections:
+        if is_gated:
+            projections = [
+                ("gate_proj", first_proj[idx, :expert_dim, :], 0, fused_dim0, True),
+                ("up_proj", first_proj[idx, expert_dim:, :], expert_dim, fused_dim0, True),
+                ("down_proj", down[idx], 0, down.shape[1], False),
+            ]
+        else:
+            # Non-gated: the single up_proj maps 1:1 to its weight quantizer, so it
+            # is exported whole (no dim-0 split, no shared gate/up weight_scale_2).
+            projections = [
+                ("up_proj", first_proj[idx], 0, fused_dim0, True),
+                ("down_proj", down[idx], 0, down.shape[1], False),
+            ]
+
+        for (
+            proj_name,
+            weight_slice,
+            fused_start,
+            fused_total,
+            uses_first_proj_quantizers,
+        ) in projections:
             w_quantizer_src = (
-                module.gate_up_proj_weight_quantizers[idx]
-                if is_gate_up
+                first_proj_weight_quantizers[idx]
+                if uses_first_proj_quantizers
                 else module.down_proj_weight_quantizers[idx]
             )
-            i_quantizer = gate_up_input_q if is_gate_up else down_input_q
+            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.
-            w_quantizer = copy.deepcopy(w_quantizer_src) if is_gate_up else w_quantizer_src
+            w_quantizer = (
+                copy.deepcopy(w_quantizer_src) if uses_first_proj_quantizers else w_quantizer_src
+            )
 
             # For per-channel amax (dim >= 1), proportionally slice dim-0
             # to match the split weight.

modelopt/torch/export/plugins/vllm_fakequant_hf.py

@@ -161,8 +161,9 @@ def _fakequant_fused_experts_weights(
     expert) that the base loop skips, leaving the fused 3-D weight unquantized
     in the export and breaking weight-fold round-trips.
     """
+    first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj")
     for w_attr, q_attr in (
-        ("gate_up_proj", "gate_up_proj_weight_quantizers"),
+        (first_proj_attr, f"{first_proj_attr}_weight_quantizers"),
         ("down_proj", "down_proj_weight_quantizers"),
     ):
         quantizers = getattr(module, q_attr, None)

modelopt/torch/export/quant_utils.py

@@ -1495,21 +1495,24 @@ def sync_tied_input_amax(model: nn.Module) -> int:
     YOCO-style models). Must run BEFORE per-module export so the merged
     amax flows into ``input_scale`` derivation. Handles both dense
     Linears (keyed by ``weight.data_ptr()``) and fused MoE (keyed by
-    ``(gate_up_proj, down_proj)`` data_ptr tuple). Returns the number of
+    ``(<first_proj>, down_proj)`` data_ptr tuple). Returns the number of
     tied groups merged.
     """
     from collections import defaultdict
 
     by_dp: dict = defaultdict(list)
     for _, m in model.named_modules():
         # Fused MoE: 3-D source tensors with shared input quantizers
+        first_proj_attr = getattr(m, "_first_proj_attr", "gate_up_proj")
+        first_proj = getattr(m, first_proj_attr, None)
+        first_proj_input_quantizer_attr = f"{first_proj_attr}_input_quantizer"
         if (
-            hasattr(m, "gate_up_proj_input_quantizer")
-            and hasattr(m, "gate_up_proj")
+            hasattr(m, first_proj_input_quantizer_attr)
+            and first_proj is not None
             and hasattr(m, "down_proj")
-            and m.gate_up_proj.dim() == 3
+            and first_proj.dim() == 3
         ):
-            key = ("moe", m.gate_up_proj.data_ptr(), m.down_proj.data_ptr())
+            key = ("moe", first_proj.data_ptr(), m.down_proj.data_ptr())
             by_dp[key].append(m)
         # Dense quantized Linear with an input_quantizer
         elif (
@@ -1549,7 +1552,8 @@ def _merge(quantizers: list) -> bool:
         if len(modules) < 2:
             continue
         if key[0] == "moe":
-            for q_name in ("gate_up_proj_input_quantizer", "down_proj_input_quantizer"):
+            first_proj_attr = getattr(modules[0], "_first_proj_attr", "gate_up_proj")
+            for q_name in (f"{first_proj_attr}_input_quantizer", "down_proj_input_quantizer"):
                 if _merge([getattr(m, q_name, None) for m in modules]):
                     synced += 1
         elif _merge([m.input_quantizer for m in modules]):

modelopt/torch/export/unified_export_hf.py

@@ -830,10 +830,12 @@ def _process_quantized_modules(
         ):
             sub_module.unpack_weight()
 
-        if hasattr(sub_module, "gate_up_proj_weight_quantizers"):
-            # _QuantFusedExperts uses plural `gate_up_proj_weight_quantizers` (ModuleList),
-            # which get_quantization_format's singular-weight_quantizer check misses. Handle
-            # it explicitly before the format gate so fused-experts get split + quantized.
+        first_proj_attr = getattr(sub_module, "_first_proj_attr", "gate_up_proj")
+        if hasattr(sub_module, f"{first_proj_attr}_weight_quantizers"):
+            # _QuantFusedExperts uses plural `<first_proj>_weight_quantizers`
+            # (ModuleList), which get_quantization_format's singular-weight_quantizer
+            # check misses. Handle it explicitly before the format gate so fused-experts
+            # get split + quantized.
             with fsdp2_aware_weight_update(model, sub_module, reshard=False):
                 _export_fused_experts(
                     sub_module,
@@ -937,6 +939,10 @@ def _export_transformers_checkpoint(
     for _, sub_module in model.named_modules():
         if is_moe(sub_module) and hasattr(sub_module, "experts"):
             expert_linear_names = get_expert_linear_names(sub_module)
+            first_proj_attr = getattr(sub_module.experts, "_first_proj_attr", "gate_up_proj")
+            has_fused_experts_quantizers = hasattr(
+                sub_module.experts, f"{first_proj_attr}_weight_quantizers"
+            )
             for linear_name in expert_linear_names:
                 # Handle DBRX experts specifically
                 if "QuantDbrxExperts" in type(sub_module.experts).__name__:
@@ -949,7 +955,7 @@ def _export_transformers_checkpoint(
                                 modules=list(linear_modulelist),
                                 quantizer_attrs=["input_quantizer"],
                             )
-                elif hasattr(sub_module.experts, "gate_up_proj_weight_quantizers"):
+                elif has_fused_experts_quantizers:
                     # _QuantFusedExperts: amax fallback is handled in _export_fused_experts
                     break
                 elif (

modelopt/torch/quantization/algorithms.py

@@ -80,6 +80,12 @@ def _is_hf_quant_fused_experts_module(module: nn.Module) -> bool:
     "down_proj_input_quantizer",
     "down_proj_weight_quantizer",
 )
+_NON_GATED_FUSED_EXPERTS_REPLAY_QUANTIZER_ATTRS = (
+    "up_proj_input_quantizer",
+    "up_proj_weight_quantizer",
+    "down_proj_input_quantizer",
+    "down_proj_weight_quantizer",
+)
 
 
 def _get_replay_quantizer_attr(attr_name: str) -> str:
@@ -97,7 +103,11 @@ def _get_quantizer_attrs(module: nn.Module) -> tuple[str, ...]:
     For standard Linear-derived QuantModules, returns the canonical trio.
     """
     if _is_hf_quant_fused_experts_module(module):
-        return _FUSED_EXPERTS_QUANTIZER_ATTRS
+        try:
+            from .plugins.huggingface import _get_fused_experts_quantizer_attr_names
+        except ImportError:
+            return _FUSED_EXPERTS_QUANTIZER_ATTRS
+        return _get_fused_experts_quantizer_attr_names(module)
     return _STD_QUANTIZER_ATTRS
 
 
@@ -1524,7 +1534,11 @@ def _cfg_to_dict(v):
         for pattern in _as_list(search_state.get("disabled_layers"))
     )
     per_module_entries: list[dict] = []
-    _per_module_attrs = (*_STD_QUANTIZER_ATTRS, *_FUSED_EXPERTS_REPLAY_QUANTIZER_ATTRS)
+    _per_module_attrs = (
+        *_STD_QUANTIZER_ATTRS,
+        *_FUSED_EXPERTS_REPLAY_QUANTIZER_ATTRS,
+        *_NON_GATED_FUSED_EXPERTS_REPLAY_QUANTIZER_ATTRS,
+    )
     # Track global (non per-module) recipe entries.  Last recipe wins for each pattern.
     global_entries: dict[str, dict] = {}
 

modelopt/torch/quantization/model_calib.py

@@ -720,8 +720,9 @@ def _warn_local_hessian_fallback(name, weight, weight_quantizer, block_size, war
 
 def _is_quant_fused_experts(module: nn.Module) -> bool:
     """Whether ``module`` is a converted HF fused-MoE-experts wrapper with per-expert quantizers."""
+    first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj")
     return hasattr(module, "_current_expert_idx") and hasattr(
-        module, "gate_up_proj_weight_quantizers"
+        module, f"{first_proj_attr}_weight_quantizers"
     )
 
 
@@ -765,11 +766,12 @@ def _dense_hook(linear, args):
             handles.append(module.register_forward_pre_hook(_dense_hook))
         elif _is_quant_fused_experts(module):
             with enable_weight_access_and_writeback(module, model, name_to_module):
+                first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj")
                 for weight_name, quantizers_name, input_q_name in (
                     (
-                        "gate_up_proj",
-                        "gate_up_proj_weight_quantizers",
-                        "gate_up_proj_input_quantizer",
+                        first_proj_attr,
+                        f"{first_proj_attr}_weight_quantizers",
+                        f"{first_proj_attr}_input_quantizer",
                     ),
                     ("down_proj", "down_proj_weight_quantizers", "down_proj_input_quantizer"),
                 ):