Support efficient TEGroupedMLP (moe_grouped_gemm=True) for Minitron Pruning

CHANGELOG.rst

@@ -32,6 +32,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 NVFP4 Four-Over-Six (4/6) weight quantization (``mtq.NVFP4_FOUR_OVER_SIX_CFG``): MSE weight calibration picks, per block, between an M=6 and an M=4 dynamic range (the choice is folded into the FP8 per-block scales), with the ``four_over_six: true`` flag normalizing those scales by 256 (vs 448) for M=4 headroom. Supported via ``mtq.quantize`` and HF / Megatron export only -- **not** ``mtq.compress``, which does not preserve the per-block M=4/M=6 choice

examples/megatron_bridge/prune_minitron.py

@@ -93,6 +93,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(
@@ -368,7 +376,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).

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
@@ -741,6 +744,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:
@@ -749,6 +757,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."""
@@ -819,8 +951,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)
 
@@ -1055,8 +1186,12 @@ def __getattribute__(self, name):
             return mixer.d_inner
         if name in ("nheads_local_tp", "nheads_local_tpcp"):
             return mixer.nheads
-        if name == "conv1d_cp1":
+        if name == "conv1d_cp1":  # nemo:26.06 and earlier: conv is a module
             return mixer.conv1d
+        if name == "conv1d_weight_cp1":  # nemo:26.08+: raw conv parameters (dynamically sliced)
+            return mixer.conv1d_weight
+        if name == "conv1d_bias_cp1":  # nemo:26.08+
+            return mixer.conv1d_bias
         if name == "dt_bias_cp1":
             return mixer.dt_bias
         if name == "A_log_cp1":
@@ -1122,11 +1257,19 @@ def _setup(self, *, hidden_size: TracedHp):
         DMRegistry.convert(self.in_proj, input_size=hidden_size, output_size=in_proj_output_size)
 
         conv_dim = build_concat_hp([d_inner, bc])  # z, B, C
-        DMRegistry.convert(self.conv1d)
-        self.conv1d.in_channels = conv_dim
-        self.conv1d.out_channels = conv_dim
-        ks = self.conv1d.get_hparam("kernel_size")
-        ks.choices = [ks.original]
+        if hasattr(self, "conv1d"):  # nemo:26.06 and earlier: a depthwise `nn.Conv1d` module.
+            DMRegistry.convert(self.conv1d)
+            self.conv1d.in_channels = conv_dim
+            self.conv1d.out_channels = conv_dim
+            ks = self.conv1d.get_hparam("kernel_size")
+            ks.choices = [ks.original]
+        else:  # nemo:26.08+: the conv is stored as raw parameters
+
+            def _slice_conv(mod, val, _hp=conv_dim):
+                return get_sliced_tensor_by_slices(val, [_hp.active_slice])
+
+            self._register_dynamic_attribute("conv1d_weight", _slice_conv)  # [conv_dim, 1, d_conv]
+            self._register_dynamic_attribute("conv1d_bias", _slice_conv)  # [conv_dim]
 
         if self.rmsnorm:
             DMRegistry.convert(self.norm)
@@ -1151,7 +1294,8 @@ def export(self) -> torch.nn.Module:
         """Export the dynamic module to a torch.nn.Module."""
         self.in_proj.export()
         self.out_proj.export()
-        self.conv1d.export()
+        if hasattr(self, "conv1d"):  # nemo:26.06 and earlier
+            self.conv1d.export()
         if self.rmsnorm:
             self.norm.export()
         return super().export()

modelopt/torch/prune/plugins/mcore_minitron.py

@@ -66,6 +66,7 @@
     _DynamicMoELayer,
     _DynamicSelfAttention,
     _DynamicSequentialMLP,
+    _DynamicTEGroupedMLP,
     _DynamicTransformerLayer,
 )
 from modelopt.torch.nas.plugins.megatron_model_stats import (
@@ -138,7 +139,7 @@ def drop_mcore_language_model_layers(model: nn.Module, *, layers_to_drop: list[i
     assert isinstance(model, supported_model_types), (
         f"Model should have one of {supported_model_types} submodule, got {model}"
     )
-    print_rank_0(f"Dropping decoder layers {layers_to_drop} from model.")
+    print_rank_0(f"Dropping decoder layers {layers_to_drop} (1-indexed) from model.")
 
     # get the number of layers remaining in each pp rank
     layers_remaining_per_pp = torch.zeros(
@@ -937,6 +938,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)
 
@@ -1365,6 +1368,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/import_utils.py

@@ -15,6 +15,7 @@
 
 """Handles suppressing import errors for third-party modules that may or may not be available."""
 
+import inspect
 from contextlib import contextmanager
 
 from .logging import warn_rank_0
@@ -23,6 +24,11 @@
 @contextmanager
 def import_plugin(plugin_name, msg_if_missing=None, verbose=True, success_msg=None):
     """Context manager to import a plugin and suppress ModuleNotFoundError."""
+    # Capture the ``with import_plugin(...)`` call site so warnings point at the plugin
+    # that actually failed rather than at this helper. When ``__enter__`` runs the
+    # generator body, the stack is [0]=here, [1]=contextlib.__enter__, [2]=the caller.
+    caller = inspect.stack()[2]
+    caller_loc = f"{caller.filename}:{caller.lineno}"
     try:
         yield
         if verbose and success_msg is not None:
@@ -33,6 +39,6 @@ def import_plugin(plugin_name, msg_if_missing=None, verbose=True, success_msg=No
     except Exception as e:
         if verbose:
             warn_rank_0(
-                f"Failed to import modelopt {plugin_name} plugin due to: {e!r}. "
-                "You may ignore this warning if you do not need this plugin."
+                f"Failed to import modelopt {plugin_name} plugin (from {caller_loc}) due to: "
+                f"{e!r}. You may ignore this warning if you do not need this plugin."
             )