Support efficient TEGroupedMLP (moe_grouped_gemm=True) for Minitron Pruning

CHANGELOG.rst

@@ -50,6 +50,7 @@ Changelog
   - **GatedDeltaNet** (linear attention) and **gated attention** (``attention_output_gate``), such as Qwen3.5 (hybrid GatedDeltaNet + gated-attention) language models, including MoE variants — attention / linear-attention heads are not pruned.
   - **Multi-Latent Attention (MLA)**, such as DeepSeek — MLA latent ranks are not pruned.
   - **Latent MoE**, such as Nemotron-3-Super — ``hidden_size`` pruning resizes the latent projections while the MoE latent dim itself is not pruned.
+- Optimize Minitron pruning support for MoE models using the fused **grouped GEMM** experts (``TEGroupedMLP``) in addition to the existing ``SequentialMLP`` path. ``examples/megatron_bridge/prune_minitron.py`` now uses grouped GEMM by default (pass ``--no_moe_grouped_gemm`` to fall back to ``TESequentialMLP``).
 - Add Minitron pruning support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/prune_minitron.py``. The language model is pruned while the vision tower is left intact and the full VLM is saved back; ``hidden_size`` is not pruned if it is shared with the vision projector. Pruning importance is estimated from image-text calibration (the full VLM forward over vision-conditioned activations) by default, or from a text dataset for text-only ablations.
 - Add PTQ support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/quantize.py``. Only the language model is quantized (vision tower + projector left in full precision) and the full VLM is saved as a Megatron checkpoint. The calibration modality is inferred from ``--calib_dataset_name``: an image-text dataset drives the full VLM forward (vision-conditioned activations), while a text dataset runs text-only calibration of the language model. Image-text calibration shards across data-parallel ranks (context parallelism is supported only for text-only calibration). HuggingFace unified export of a quantized VLM is not yet supported.
 - Add Megatron-Bridge distillation and Quantization-Aware Distillation (QAD) support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/distill.py``.

examples/megatron_bridge/prune_minitron.py

@@ -107,6 +107,14 @@ def get_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument("--hf_model_name_or_path", type=str, required=True)
     parser.add_argument("--trust_remote_code", action="store_true")
+    parser.add_argument(
+        "--no_moe_grouped_gemm",
+        action="store_true",
+        help=(
+            "Use SequentialMLP for MoE experts instead of the (default) efficient fused "
+            "TEGroupedMLP (grouped GEMM). Only affects MoE models."
+        ),
+    )
 
     target_group = parser.add_mutually_exclusive_group(required=True)
     target_group.add_argument(
@@ -382,7 +390,7 @@ def main(args: argparse.Namespace):
             "mtp_num_layers": 0,  # MTP is not supported during calibration
         },
         init_model_parallel=True,
-        moe_grouped_gemm=False,
+        moe_grouped_gemm=not args.no_moe_grouped_gemm,
     )
 
     # TODO: Support pruning with MTP heads enabled (e.g. Qwen3.5 mtp_num_hidden_layers=1).
@@ -503,18 +511,9 @@ def main(args: argparse.Namespace):
         )
 
         match = re.fullmatch(r"mmlu_(\d+)pct_bs(\d+)", args.prune_score_func)
-        legacy_match = re.fullmatch(r"mmlu_(\d+)pct", args.prune_score_func)
         if match:
             mmlu_frac = float(match.group(1)) / 100.0
             batch_size = int(match.group(2))
-        elif legacy_match:
-            warn_rank_0(
-                f"Score function '{args.prune_score_func}' uses the deprecated format "
-                "'mmlu_<N>pct'. Use 'mmlu_<N>pct_bs<bs>' to specify the evaluation batch size. "
-                "Falling back to batch_size=1."
-            )
-            mmlu_frac = float(legacy_match.group(1)) / 100.0
-            batch_size = 1
         else:
             raise ValueError(
                 f"Invalid score function: {args.prune_score_func}. "

modelopt/torch/nas/plugins/megatron.py

@@ -19,15 +19,18 @@
 import types
 from abc import ABC
 from collections.abc import Callable, Sequence
+from functools import partial
 
 import torch
 import torch.nn as nn
 import transformer_engine as te
 from megatron.core.extensions.transformer_engine import (
+    TEColumnParallelGroupedLinear,
     TEColumnParallelLinear,
     TEDotProductAttention,
     TELayerNormColumnParallelLinear,
     TELinear,
+    TERowParallelGroupedLinear,
     TERowParallelLinear,
 )
 from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
@@ -44,7 +47,7 @@
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.mlp import MLP
 from megatron.core.transformer.moe import moe_utils
-from megatron.core.transformer.moe.experts import SequentialMLP
+from megatron.core.transformer.moe.experts import SequentialMLP, TEGroupedMLP
 from megatron.core.transformer.moe.moe_layer import MoELayer
 from megatron.core.transformer.moe.router import TopKRouter
 from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
@@ -759,6 +762,11 @@ def _setup(self, *, hidden_size: TracedHp):
         for expert in self.local_experts:
             DMRegistry.convert(expert, hidden_size=hidden_size, hp_name="moe_ffn_hidden_size")
 
+    def modify(self, ffn_hidden_size_divisor: int = 1, **kwargs) -> None:
+        """Modify each expert's moe_ffn_hidden_size hparam choices based on search space config."""
+        for expert in self.local_experts:
+            expert.modify(ffn_hidden_size_divisor=ffn_hidden_size_divisor)
+
     def export(self) -> torch.nn.Module:
         """Export the dynamic module to a standard SequentialMLP."""
         for expert in self.local_experts:
@@ -767,6 +775,130 @@ def export(self) -> torch.nn.Module:
         return super().export()
 
 
+@DMRegistry.register(
+    {
+        TEColumnParallelGroupedLinear: (
+            "megatron.core.extensions.transformer_engine.TEColumnParallelGroupedLinear"
+        ),
+        TERowParallelGroupedLinear: (
+            "megatron.core.extensions.transformer_engine.TERowParallelGroupedLinear"
+        ),
+    }
+)
+class _DynamicTEGroupedLinear(DynamicModule):
+    """A TEGroupedLinear (column/row parallel) with dynamic hyperparams for grouped-GEMM MoE.
+
+    TEGroupedMLP fuses all local experts into two grouped linears, each storing the per-expert
+    weights as separate ``weight0..weight{num_gemms-1}`` params (shape ``[out, in]``, optional
+    ``bias{i}``). ``moe_ffn_hidden_size`` / ``hidden_size`` slice each expert weight by
+    ``output_size`` (rows) / ``input_size`` (cols) like a normal linear; ``num_local_experts``
+    reorders/drops experts by remapping position ``j`` to the ``j``-th most important expert and
+    exposing ``num_gemms = num_local_experts.active`` so TE only reads the kept experts.
+    """
+
+    def _setup(self, *, input_size: TracedHp, output_size: TracedHp, num_local_experts: TracedHp):
+        assert not self.single_grouped_weight, (
+            "moe_single_grouped_weight=True is not supported for grouped-GEMM pruning yet."
+        )
+        # input_size/output_size/num_local_experts are all shared with the sibling grouped linear
+        # (and num_local_experts additionally with the router) via _DynamicTEGroupedMLP.
+        self._register_hparam("input_size", input_size)
+        self._register_hparam("output_size", output_size)
+        self._register_hparam("num_local_experts", num_local_experts)
+
+        self._register_dynamic_attribute("num_gemms", lambda mod, val: num_local_experts.active)
+        self._register_dynamic_attribute("in_features", lambda mod, val: input_size.active)
+        self._register_dynamic_attribute("out_features", lambda mod, val: output_size.active)
+        for j in range(self.num_gemms):
+            self._register_dynamic_attribute(f"weight{j}", partial(self._get_expert_param, pos=j))
+            if self.use_bias:
+                self._register_dynamic_attribute(f"bias{j}", partial(self._get_expert_param, pos=j))
+
+    @staticmethod
+    def _get_expert_param(mod: "_DynamicTEGroupedLinear", val: torch.Tensor, *, pos: int):
+        """Dynamic getter for weight{pos}/bias{pos}: map position -> ranked expert, then slice."""
+        hp = mod.get_hparam("num_local_experts")
+        max_experts = hp.max
+        assert isinstance(max_experts, int)
+        order = hp._slice_order.tolist() if hp._slice_order is not None else range(max_experts)
+        e = order[pos]
+        is_weight = val.dim() == 2
+        raw = mod._parameters[f"{'weight' if is_weight else 'bias'}{e}"]
+        slices = [mod.get_hparam("output_size").active_slice]
+        if is_weight:
+            slices.append(mod.get_hparam("input_size").active_slice)
+        return get_sliced_tensor_by_slices(raw, slices)
+
+    def export(self) -> torch.nn.Module:
+        """Export to a standard TEGroupedLinear with the kept experts sliced + reordered in place."""
+        # Read all sliced/reordered params (via the dynamic getters) before mutating any, then drop
+        # the per-expert weight/bias attrs so the base export only folds num_gemms/in/out_features.
+        active = self.get_hparam("num_local_experts").active
+        assert isinstance(active, int)
+        weights = [getattr(self, f"weight{j}").detach().clone() for j in range(active)]
+        biases = [
+            getattr(self, f"bias{j}").detach().clone() for j in range(active) if self.use_bias
+        ]
+        for name in [n for n in list(self._parameters) if n.startswith(("weight", "bias"))]:
+            delattr(self, name)
+
+        super().export()  # num_gemms -> active, in/out_features -> sliced sizes, class un-patched
+
+        for j, weight in enumerate(weights):
+            self.register_parameter(f"weight{j}", torch.nn.Parameter(weight))
+        for j, bias in enumerate(biases):
+            self.register_parameter(f"bias{j}", torch.nn.Parameter(bias))
+        return self
+
+
+@DMRegistry.register({TEGroupedMLP: "megatron.core.transformer.moe.experts.TEGroupedMLP"})
+class _DynamicTEGroupedMLP(DynamicModule):
+    """A TEGroupedMLP (grouped-GEMM MoE experts) with dynamic hyperparams.
+
+    Mirrors ``_DynamicSequentialMLP`` but the experts are two fused ``TEGroupedLinear`` layers rather
+    than an ``nn.ModuleList`` of per-expert MLPs. Since Minitron prunes homogeneously, all experts
+    share a single ``moe_ffn_hidden_size`` hparam (unlike the SequentialMLP path which registers one per expert).
+    """
+
+    def _setup(self, *, hidden_size: TracedHp):
+        """Setup the TEGroupedMLP dynamic module with global hidden_size hparam."""
+        num_local_experts = TracedHp(list(range(1, self.num_local_experts + 1)))
+        self._register_hparam("num_local_experts", num_local_experts)
+
+        moe_ffn_hidden_size = TracedHp(list(range(1, self.config.moe_ffn_hidden_size + 1)))
+        self._register_hparam("moe_ffn_hidden_size", moe_ffn_hidden_size)
+
+        linear_fc1_output_size = (
+            build_concat_hp([moe_ffn_hidden_size] * 2)
+            if self.config.gated_linear_unit
+            else moe_ffn_hidden_size
+        )
+        DMRegistry.convert(  # _DynamicTEGroupedLinear
+            self.linear_fc1,
+            input_size=hidden_size,
+            output_size=linear_fc1_output_size,
+            num_local_experts=num_local_experts,
+        )
+        DMRegistry.convert(  # _DynamicTEGroupedLinear
+            self.linear_fc2,
+            input_size=moe_ffn_hidden_size,
+            output_size=hidden_size,
+            num_local_experts=num_local_experts,
+        )
+
+    def modify(self, ffn_hidden_size_divisor: int = 1, **kwargs) -> None:
+        """Modify the shared moe_ffn_hidden_size hparam choices based on search space config."""
+        hp = self.get_hparam("moe_ffn_hidden_size")
+        choices = {int(make_divisible(c, ffn_hidden_size_divisor)) for c in hp.choices}  # type: ignore[arg-type]
+        hp.choices = list(set(hp.choices) & choices | {hp.original})
+
+    def export(self) -> torch.nn.Module:
+        """Export the dynamic module to a standard TEGroupedMLP."""
+        self.linear_fc1.export()
+        self.linear_fc2.export()
+        return super().export()
+
+
 @DMRegistry.register({MoELayer: "megatron.core.transformer.moe.moe_layer.MoELayer"})
 class _DynamicMoELayer(DynamicModule):
     """A MoELayer with dynamic hyperparams."""
@@ -837,8 +969,7 @@ def modify(
         expert_hp.choices = list(set(expert_hp.choices) & choices | {expert_hp.original})
 
         # Modify expert FFN hparam choices
-        for expert in self.experts.local_experts:
-            expert.modify(ffn_hidden_size_divisor=ffn_hidden_size_divisor)
+        self.experts.modify(ffn_hidden_size_divisor=ffn_hidden_size_divisor)
         if self.use_shared_expert:
             self.shared_experts.modify(ffn_hidden_size_divisor)
 

modelopt/torch/prune/plugins/mcore_minitron.py

@@ -67,6 +67,7 @@
     _DynamicMoELayer,
     _DynamicSelfAttention,
     _DynamicSequentialMLP,
+    _DynamicTEGroupedMLP,
     _DynamicTransformerLayer,
 )
 from modelopt.torch.nas.plugins.megatron_model_stats import (
@@ -974,6 +975,8 @@ def __init__(self, model: DynamicModule):
                 _register_mlp_importance(module, self)
             elif isinstance(module, _DynamicSequentialMLP):
                 _register_sequential_mlp_importance(module, self)
+            elif isinstance(module, _DynamicTEGroupedMLP):
+                _register_grouped_mlp_importance(module, self)
             elif isinstance(module, _DynamicMambaMixer):
                 _register_mamba_mixer_importance(module, self)
 
@@ -1402,6 +1405,46 @@ def _estimate_expert_importance(mod):
     )
 
 
+def _register_grouped_mlp_importance(
+    module: _DynamicTEGroupedMLP, registry: ImportanceEstimatorRegistry
+) -> None:
+    """Register importance estimators for TEGroupedMLP (grouped-GEMM MoE experts) modules.
+
+    Mirrors the SequentialMLP path: ``num_local_experts`` reuses the expert-L2 hook (TEGroupedMLP
+    shares SequentialMLP's forward signature), and ``moe_ffn_hidden_size`` is a single shared score
+    from the fused ``linear_fc2`` input activations (all experts' tokens), since experts prune
+    homogeneously.
+    """
+    # Expert importance for num_local_experts; also creates module._activations (the dict saved and
+    # restored by the per-rank score checkpoint). We stash the ffn score in it so re-pruning from a
+    # checkpoint recovers it without re-running the forward loop.
+    _register_sequential_mlp_importance(module, registry)
+    module._activations["ffn_activations"] = None
+
+    def _grouped_fc2_forward_hook(mod, module_inner, input, output):
+        """Collect ffn-channel activations from the fused linear_fc2 input (all experts' tokens)."""
+        # input[0] is the permuted intermediate [total_tokens, moe_ffn_hidden_size] (no batch dim)
+        acts = gather_from_tensor_model_parallel_region(input[0]).detach()[:, None, :]
+        acts = acts.to(torch.float32).abs().mean(dim=0).pow(2).sum(dim=0)  # [moe_ffn_hidden_size]
+        prev = mod._activations["ffn_activations"]
+        mod._activations["ffn_activations"] = acts if prev is None else prev + acts
+
+    def _estimate_grouped_ffn_importance(mod):
+        """Return the activation magnitude-based importance (L2 norm) of moe_ffn_hidden_size."""
+        acts = mod._activations["ffn_activations"]
+        assert acts is not None, "No activations collected for importance estimation."
+        return acts.pow(0.5)
+
+    registry.register_hook(
+        module.linear_fc2,
+        partial(_grouped_fc2_forward_hook, module),
+        hook_type="forward",
+    )
+    registry.register_importance(
+        module, "moe_ffn_hidden_size", lambda: _estimate_grouped_ffn_importance(module)
+    )
+
+
 def _register_mamba_mixer_importance(
     module: _DynamicMambaMixer, registry: ImportanceEstimatorRegistry
 ) -> None:

modelopt/torch/utils/plugins/mbridge.py

@@ -95,10 +95,9 @@ def load_mbridge_model_from_hf(
             assert hasattr(provider, key), f"{type(provider)} does not have attribute {key}"
             setattr(provider, key, value)
 
-    # Pruning does not support grouped GEMM yet, so disable it for MoE models. Set the flag on the
-    # provider (the bridge's native, possibly custom/hybrid spec reads it at build time) rather than
-    # replacing the whole layer spec -- overwriting it would drop custom layers (e.g. Qwen3.5's
-    # GatedDeltaNet + gated-attention or Gemma3's custom spec).
+    # Set moe_grouped_gemm on the provider (the bridge's native, possibly custom/hybrid spec reads
+    # it at build time) rather than replacing the whole layer spec -- overwriting it would drop
+    # custom layers (e.g. Qwen3.5's GatedDeltaNet + gated-attention or Gemma3's custom spec).
     if HAS_HYBRID and isinstance(provider, (HybridModelProvider)):
         provider.hybrid_stack_spec = get_te_hybrid_stack_spec(moe_grouped_gemm=moe_grouped_gemm)
         provider.moe_grouped_gemm = moe_grouped_gemm

modelopt/torch/utils/plugins/megatron_generate.py

@@ -88,6 +88,25 @@ def cp_gather_logits(local_logits: torch.Tensor, cp_group, global_seq_len: int)
     return torch.cat(chunks, dim=1)
 
 
+def _assert_mamba_within_int32_indexing(model: MegatronModule, batch_size: int, seq_length: int):
+    """Guard Mamba2 calibration against int32 activation-indexing overflow.
+
+    The SSD Triton kernels index activations with int32, so ``batch * seq * mamba in_proj`` must stay
+    < 2**31 or they hit a cryptic CUDA 'illegal memory access'. Fail fast with the max safe batch.
+    """
+    d = getattr(model, "_modelopt_max_mamba_in_proj", None)
+    if d is None:
+        d = max(
+            (m.in_proj.weight.shape[0] for m in model.modules() if hasattr(m, "in_proj")), default=0
+        )
+        model._modelopt_max_mamba_in_proj = d
+    if d and batch_size * seq_length * d >= 2**31:
+        raise ValueError(
+            f"{batch_size=} x {seq_length=} x mamba in_proj={d} overflows int32 in the Mamba2 kernels. "
+            f"Reduce calibration batch size to <= {2**31 // (seq_length * d)}."
+        )
+
+
 def get_current_memory_info():
     """Get current memory usage."""
     remaining_mem, total_mem = torch.cuda.mem_get_info()
@@ -157,6 +176,9 @@ def megatron_prefill(
 
     padded_seq_len = tokens.shape[-1]
 
+    # Fail fast with a clear message if the Mamba activation would overflow int32 kernel indexing.
+    _assert_mamba_within_int32_indexing(model, batch_size, padded_seq_len)
+
     cp_size = mpu.get_context_parallel_world_size()
 
     # Under CP a local triu mask would be wrong for the per-rank zigzag chunks; pass None and let

tests/examples/megatron_bridge/test_prune_minitron.py

@@ -60,8 +60,9 @@ def test_prune_minitron(tmp_path, num_gpus, create_teacher, megatron_format):
         if megatron_format
         else {"output_hf_path": pruned_path}
     )
+    # TODO: Dont enable grouped GEMM for MoE models until nemo:26.08 container
     prune_command_parts = extend_cmd_parts(
-        ["torchrun", f"--nproc_per_node={num_gpus}", "prune_minitron.py"],
+        ["torchrun", f"--nproc_per_node={num_gpus}", "prune_minitron.py", "--no_moe_grouped_gemm"],
         hf_model_name_or_path=teacher_hf_path,
         pp_size=num_gpus,
         calib_dataset_name="cnn_dailymail",
@@ -126,8 +127,9 @@ def test_prune_minitron_vlm(tmp_path, num_gpus, create_teacher):
     prune_target_params = int(language_model_params * 0.7)
 
     pruned_model_path = tmp_path / "pruned"
+    # TODO: Dont enable grouped GEMM for MoE models until nemo:26.08 container
     prune_command_parts = extend_cmd_parts(
-        ["torchrun", f"--nproc_per_node={num_gpus}", "prune_minitron.py"],
+        ["torchrun", f"--nproc_per_node={num_gpus}", "prune_minitron.py", "--no_moe_grouped_gemm"],
         hf_model_name_or_path=teacher_hf_path,
         output_hf_path=pruned_model_path,
         pp_size=num_gpus,