diff --git a/CHANGELOG.rst b/CHANGELOG.rst index df3cdefdb5e..36931b71f36 100755 --- a/CHANGELOG.rst +++ b/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 diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index bdd40b3ad8f..813f90ca766 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/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). diff --git a/modelopt/torch/nas/plugins/megatron.py b/modelopt/torch/nas/plugins/megatron.py index 4ad31567a05..f34e5af222f 100644 --- a/modelopt/torch/nas/plugins/megatron.py +++ b/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() diff --git a/modelopt/torch/prune/plugins/mcore_minitron.py b/modelopt/torch/prune/plugins/mcore_minitron.py index cee43d85807..a9a8fad7a12 100644 --- a/modelopt/torch/prune/plugins/mcore_minitron.py +++ b/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: diff --git a/modelopt/torch/utils/import_utils.py b/modelopt/torch/utils/import_utils.py index 8229da51e51..15bc023ffa9 100644 --- a/modelopt/torch/utils/import_utils.py +++ b/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." ) diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py index e53e6516992..ae3b79cb298 100644 --- a/modelopt/torch/utils/plugins/mbridge.py +++ b/modelopt/torch/utils/plugins/mbridge.py @@ -128,10 +128,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 isinstance(provider, MambaModelProvider): provider.mamba_stack_spec = get_te_mamba_stack_spec(moe_grouped_gemm=moe_grouped_gemm) elif (provider.num_moe_experts or 0) > 0: diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 8ca09a77967..09df25f2f0f 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -59,8 +59,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", @@ -121,8 +122,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, diff --git a/tests/gpu_megatron/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py b/tests/gpu_megatron/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py index 158b6cafacd..5df4c2fa79e 100644 --- a/tests/gpu_megatron/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py +++ b/tests/gpu_megatron/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py @@ -36,6 +36,8 @@ _DynamicMoELayer, _DynamicSelfAttention, _DynamicSequentialMLP, + _DynamicTEGroupedLinear, + _DynamicTEGroupedMLP, _DynamicTELayerNormColumnParallelLinear, _DynamicTEProjRowParallelLinear, _DynamicTEQKVLayerNormColumnParallelLinear, @@ -231,7 +233,7 @@ def test_gpt_self_attention_head_sorting(distributed_setup_size_1): destroy_model_parallel() -def _test_gpt_moe_search_space(rank, size): +def _test_gpt_moe_search_space(moe_grouped_gemm, rank, size): channel_divisor = 4 num_layers = min(size * 2, 8) @@ -258,6 +260,7 @@ def _test_gpt_moe_search_space(rank, size): activation_func="squared_relu", transformer_impl="transformer_engine", num_moe_experts=num_moe_experts, + moe_grouped_gemm=moe_grouped_gemm, moe_ffn_hidden_size=moe_ffn_hidden_size, moe_shared_expert_intermediate_size=moe_shared_expert_intermediate_size, ).cuda() @@ -280,11 +283,16 @@ def _test_gpt_moe_search_space(rank, size): moe = model.decoder.layers[0].mlp assert isinstance(moe, _DynamicMoELayer) assert isinstance(moe.router, _DynamicTopKRouter) - assert isinstance(moe.experts, _DynamicSequentialMLP) - assert isinstance(moe.experts.local_experts, DynamicModuleList) - for expert in moe.experts.local_experts: - assert isinstance(expert, _DynamicMLP) assert isinstance(moe.shared_experts, _DynamicMLP) + if moe_grouped_gemm: + assert isinstance(moe.experts, _DynamicTEGroupedMLP) + assert isinstance(moe.experts.linear_fc1, _DynamicTEGroupedLinear) + assert isinstance(moe.experts.linear_fc2, _DynamicTEGroupedLinear) + else: + assert isinstance(moe.experts, _DynamicSequentialMLP) + assert isinstance(moe.experts.local_experts, DynamicModuleList) + for expert in moe.experts.local_experts: + assert isinstance(expert, _DynamicMLP) # NOTE: `search_space_size` does not reduce across TP/PP groups ss_size_per_pp = search_space_size(model) @@ -293,15 +301,12 @@ def _test_gpt_moe_search_space(rank, size): moe_shared_ffn_choices = moe_shared_expert_intermediate_size // channel_divisor hidden_size_choices = hidden_size // channel_divisor num_layers_per_pp = num_layers // size - # SequentialMLP has per-expert moe_ffn_hidden_size hparams + # SequentialMLP has one moe_ffn_hidden_size hparam per expert (moe_ffn_choices**num_moe_experts); + # TEGroupedMLP shares a single one (moe_ffn_choices). + moe_ffn_ss = moe_ffn_choices if moe_grouped_gemm else moe_ffn_choices**num_moe_experts assert ( ss_size_per_pp - == ( - num_heads_choices - * num_moe_experts - * moe_ffn_choices**num_moe_experts - * moe_shared_ffn_choices - ) + == (num_heads_choices * num_moe_experts * moe_ffn_ss * moe_shared_ffn_choices) ** num_layers_per_pp * num_layers * hidden_size_choices @@ -319,5 +324,6 @@ def _test_gpt_moe_search_space(rank, size): assert not any(named_dynamic_modules(model)) -def test_gpt_moe_search_space(dist_workers): - dist_workers.run(_test_gpt_moe_search_space) +@pytest.mark.parametrize("moe_grouped_gemm", [False, True]) +def test_gpt_moe_search_space(dist_workers, moe_grouped_gemm): + dist_workers.run(partial(_test_gpt_moe_search_space, moe_grouped_gemm)) diff --git a/tests/gpu_megatron/torch/nas/plugins/test_megatron_mamba_dynamic_modules.py b/tests/gpu_megatron/torch/nas/plugins/test_megatron_mamba_dynamic_modules.py index db8b9e10ba6..c81c3eeb7f2 100644 --- a/tests/gpu_megatron/torch/nas/plugins/test_megatron_mamba_dynamic_modules.py +++ b/tests/gpu_megatron/torch/nas/plugins/test_megatron_mamba_dynamic_modules.py @@ -100,7 +100,8 @@ def _test_mamba_search_space(rank, size): assert isinstance(layer.mixer, _DynamicMambaMixer) assert isinstance(layer.mixer.in_proj, _DynamicTELayerNormColumnParallelLinear) assert isinstance(layer.mixer.out_proj, _DynamicTERowParallelLinear) - assert isinstance(layer.mixer.conv1d, _DynamicConvNd) + if hasattr(layer.mixer, "conv1d"): # nemo:26.06 and earlier + assert isinstance(layer.mixer.conv1d, _DynamicConvNd) if layer.mixer.rmsnorm: assert isinstance(layer.mixer.norm, _DynamicExtendedRMSNorm) if is_pipeline_last_stage(): diff --git a/tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py b/tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py index cb66b11683b..031381d61aa 100644 --- a/tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py +++ b/tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py @@ -356,7 +356,7 @@ def test_mcore_gpt_pruning( ) -def _test_mcore_gpt_moe_parameter_sorting(rank, size): +def _test_mcore_gpt_moe_parameter_sorting(moe_grouped_gemm, rank, size): set_seed(SEED) # Use relatively bigger model here for more accurate test for sorting channel_divisor = 64 @@ -385,6 +385,7 @@ def _test_mcore_gpt_moe_parameter_sorting(rank, size): activation_func="squared_relu", transformer_impl="transformer_engine", num_moe_experts=num_moe_experts, + moe_grouped_gemm=moe_grouped_gemm, moe_ffn_hidden_size=moe_ffn_hidden_size, moe_shared_expert_intermediate_size=moe_shared_expert_intermediate_size, bf16=False, @@ -408,9 +409,11 @@ def _test_mcore_gpt_moe_parameter_sorting(rank, size): sortable_per_pp = [ n for n, hp in dynamic_space.named_hparams(configurable=True) if hp.importance is not None ] - # (num_moe_experts + 3) hps per layer + 1 for hidden_size (num_layers is not sorted!) - # Per layer: num_attention_heads, num_moe_experts, moe_ffn (per expert), moe_shared_ffn - assert len(sortable_per_pp) == (num_moe_experts + 3) * num_layers // size + 1 + # (moe_ffn_count + 3) hps per layer + 1 for hidden_size (num_layers is not sorted!) + # Per layer: num_attention_heads, num_moe_experts, moe_ffn, moe_shared_ffn. + # SequentialMLP registers one moe_ffn hparam per expert; TEGroupedMLP shares a single one. + moe_ffn_count = 1 if moe_grouped_gemm else num_moe_experts + assert len(sortable_per_pp) == (moe_ffn_count + 3) * num_layers // size + 1 # sanity check if the model functionality is preserved after sorting export_searchspace(model, mtn.get_subnet_config(model)) @@ -418,11 +421,12 @@ def _test_mcore_gpt_moe_parameter_sorting(rank, size): compare_outputs(y1, y2, rtol=1e-5, atol=1e-3) -def test_mcore_gpt_moe_parameter_sorting(dist_workers): - dist_workers.run(_test_mcore_gpt_moe_parameter_sorting) +@pytest.mark.parametrize("moe_grouped_gemm", [False, True]) +def test_mcore_gpt_moe_parameter_sorting(dist_workers, moe_grouped_gemm): + dist_workers.run(partial(_test_mcore_gpt_moe_parameter_sorting, moe_grouped_gemm)) -def _test_mcore_gpt_pruning_moe(ckpt_dir, rank, size): +def _test_mcore_gpt_pruning_moe(ckpt_dir, moe_grouped_gemm, rank, size): channel_divisor = 4 num_layers = size @@ -446,6 +450,7 @@ def _get_model(initialize_megatron=True): activation_func="squared_relu", transformer_impl="transformer_engine", num_moe_experts=num_moe_experts, + moe_grouped_gemm=moe_grouped_gemm, moe_ffn_hidden_size=moe_ffn_hidden_size, moe_shared_expert_intermediate_size=moe_shared_expert_intermediate_size, ).cuda() @@ -483,10 +488,24 @@ def _get_model(initialize_megatron=True): assert moe.router.expert_bias.shape == (pruned_num_moe_experts,) assert moe.router.weight.shape == (pruned_num_moe_experts, pruned_hidden_size) assert moe.experts.num_local_experts == pruned_num_moe_experts - assert len(moe.experts.local_experts) == pruned_num_moe_experts - for expert in moe.experts.local_experts: - assert expert.linear_fc1.weight.shape == (pruned_moe_ffn, pruned_hidden_size) - assert expert.linear_fc2.weight.shape == (pruned_hidden_size, pruned_moe_ffn) + if moe_grouped_gemm: + # TEGroupedMLP fuses experts into two grouped linears with per-expert weight{i} params + assert moe.experts.linear_fc1.num_gemms == pruned_num_moe_experts + assert moe.experts.linear_fc2.num_gemms == pruned_num_moe_experts + for i in range(pruned_num_moe_experts): + assert getattr(moe.experts.linear_fc1, f"weight{i}").shape == ( + pruned_moe_ffn, + pruned_hidden_size, + ) + assert getattr(moe.experts.linear_fc2, f"weight{i}").shape == ( + pruned_hidden_size, + pruned_moe_ffn, + ) + else: + assert len(moe.experts.local_experts) == pruned_num_moe_experts + for expert in moe.experts.local_experts: + assert expert.linear_fc1.weight.shape == (pruned_moe_ffn, pruned_hidden_size) + assert expert.linear_fc2.weight.shape == (pruned_hidden_size, pruned_moe_ffn) assert moe.shared_experts.linear_fc1.weight.shape == ( pruned_moe_shared_ffn, pruned_hidden_size, @@ -519,8 +538,11 @@ def _get_model(initialize_megatron=True): ) -def test_mcore_gpt_pruning_moe(dist_workers, tmp_path): - dist_workers.run(partial(_test_mcore_gpt_pruning_moe, tmp_path / "minitron_scores")) +@pytest.mark.parametrize("moe_grouped_gemm", [False, True]) +def test_mcore_gpt_pruning_moe(dist_workers, tmp_path, moe_grouped_gemm): + dist_workers.run( + partial(_test_mcore_gpt_pruning_moe, tmp_path / "minitron_scores", moe_grouped_gemm) + ) def _build_and_prune_variant(size, export_config, *, num_attention_heads=4, **model_kwargs): diff --git a/tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py b/tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py index 93ef70ac40e..239f68d543a 100644 --- a/tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py +++ b/tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py @@ -210,7 +210,10 @@ def forward_loop(m): assert mixer.headdim == pruned_mamba_head_dim assert mixer.d_inner == pruned_mamba_num_heads * pruned_mamba_head_dim assert mixer.out_proj.out_features == pruned_hidden_size - assert mixer.conv1d.in_channels == mixer.conv1d.out_channels == mixer.d_inner + bc + if hasattr(mixer, "conv1d"): # nemo:26.06 and earlier + assert mixer.conv1d.in_channels == mixer.conv1d.out_channels == mixer.d_inner + bc + else: # nemo:26.08+ + assert mixer.conv1d_weight.shape[0] == mixer.conv1d_bias.shape[0] == mixer.d_inner + bc # Assert model.config is updated for correct save/restoring assert model.config.ffn_hidden_size == pruned_ffn_hidden_size @@ -255,13 +258,14 @@ def test_mcore_mamba_hybrid_pruning(dist_workers, tmp_path): } -def _make_nas_hybrid_model(size): +def _make_nas_hybrid_model(size, moe_grouped_gemm=False): return get_mcore_mamba_hybrid_model( tensor_model_parallel_size=1, pipeline_model_parallel_size=size, initialize_megatron=True, transformer_impl="transformer_engine", bf16=False, + moe_grouped_gemm=moe_grouped_gemm, **_NAS_MODEL_KWARGS, ).cuda() @@ -332,7 +336,8 @@ def _assert_top_k_candidates(searcher_state, constraint_key, expected_top_k, k=1 def _test_mcore_mamba_hybrid_pruning_nas_params(rank, size, ckpt_dir): set_seed(SEED) - model = _make_nas_hybrid_model(size) + # Covers grouped-GEMM (TEGroupedMLP) MoE pruning; the memory_mb test below covers SequentialMLP. + model = _make_nas_hybrid_model(size, moe_grouped_gemm=True) baseline_params, baseline_active = mcore_param_count( model.config, @@ -427,7 +432,8 @@ def _test_mcore_mamba_hybrid_pruning_nas_memory_mb(rank, size, ckpt_dir): set_seed(SEED) dtype_bytes = 2 sequence_length = 128 - model = _make_nas_hybrid_model(size) + # Covers SequentialMLP MoE pruning; the params test above covers grouped-GEMM (TEGroupedMLP). + model = _make_nas_hybrid_model(size, moe_grouped_gemm=False) _, _, _, baseline_memory_mb = mcore_memory_footprint_mb( model.config,