[Cleanup] Remove the deprecated GroupedMLP

gpt_builders.py

@@ -132,7 +132,6 @@ def _get_transformer_layer_spec(use_te, config):
             args.qk_layernorm,
             args.multi_latent_attention,
             args.experimental_attention_variant,
-            moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm,
             qk_l2_norm=args.qk_l2_norm,
             use_kitchen=config.use_kitchen,
             use_te_activation_func=config.use_te_activation_func,
@@ -152,7 +151,6 @@ def _get_transformer_layer_spec(use_te, config):
             args.qk_layernorm,
             args.multi_latent_attention,
             args.experimental_attention_variant,
-            moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm,
             normalization=args.normalization,
             use_kitchen=config.use_kitchen,
             use_kitchen_attention=config.use_kitchen_attention,

megatron/core/extensions/transformer_engine_spec_provider.py

@@ -20,7 +20,6 @@
 from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
 from megatron.core.transformer.mlp import MLPSubmodules, TEActivationFunctionBuilder
 from megatron.core.transformer.moe.experts import (
-    GroupedMLP,
     SequentialMLP,
     TEGroupedMLP,
     TEGroupedMLPSubmodules,
@@ -66,27 +65,16 @@ def core_attention(self) -> type:
         return TEDotProductAttention
 
     def grouped_mlp_modules(
-        self, moe_use_grouped_gemm: bool, moe_use_legacy_grouped_gemm: bool
+        self, moe_use_grouped_gemm: bool
     ) -> (
         tuple[type[TEGroupedMLP], TEGroupedMLPSubmodules]
         | tuple[type[SequentialMLP], MLPSubmodules]
-        | tuple[type[GroupedMLP], None]
     ):
         """Which module and submodules to use for grouped mlp"""
-        if (
-            moe_use_grouped_gemm
-            and TEColumnParallelGroupedLinear is not None
-            and not moe_use_legacy_grouped_gemm
-        ):
+        if moe_use_grouped_gemm and TEColumnParallelGroupedLinear is not None:
             return TEGroupedMLP, TEGroupedMLPSubmodules(
                 linear_fc1=TEColumnParallelGroupedLinear, linear_fc2=TERowParallelGroupedLinear
             )
-        elif moe_use_grouped_gemm:
-            warnings.warn(
-                'The legacy GroupedMLP will be deprecated in Megatron-Core v0.12.0. '
-                'Please update the TransformerEngine to version>=1.7.0 and use TEGroupedMLP.'
-            )
-            return GroupedMLP, None
         else:
             if not is_te_min_version("1.7.0.dev0"):
                 warnings.warn(

megatron/core/models/backends.py

@@ -13,7 +13,6 @@
 from megatron.core.transformer.dot_product_attention import DotProductAttention
 from megatron.core.transformer.mlp import MLPSubmodules, TEActivationFunctionBuilder
 from megatron.core.transformer.moe.experts import (
-    GroupedMLP,
     InferenceGroupedMLP,
     SequentialMLP,
     TEGroupedMLPSubmodules,
@@ -84,7 +83,7 @@ def core_attention(self) -> type:
 
     @abstractmethod
     def grouped_mlp_modules(
-        self, moe_use_grouped_gemm: bool, moe_use_legacy_grouped_gemm: bool
+        self, moe_use_grouped_gemm: bool
     ) -> tuple[type, MLPSubmodules | TEGroupedMLPSubmodules | None]:
         """Which module and submodules to use for grouped mlp"""
         ...
@@ -128,19 +127,12 @@ def core_attention(self) -> type:
         return DotProductAttention
 
     def grouped_mlp_modules(
-        self, moe_use_grouped_gemm: bool, moe_use_legacy_grouped_gemm: bool
-    ) -> tuple[type[GroupedMLP], None] | tuple[type[SequentialMLP], MLPSubmodules]:
+        self, moe_use_grouped_gemm: bool
+    ) -> tuple[type[SequentialMLP], MLPSubmodules]:
         """Which module and submodules to use for grouped mlp"""
-        if moe_use_grouped_gemm:
-            warnings.warn(
-                "The legacy GroupedMLP will be deprecated in Megatron-Core v0.12.0. "
-                "Please update the TransformerEngine to version>=1.7.0 and use TEGroupedMLP."
-            )
-            return GroupedMLP, None
-        else:
-            return SequentialMLP, MLPSubmodules(
-                linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear
-            )
+        return SequentialMLP, MLPSubmodules(
+            linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear
+        )
 
     def activation_func(self) -> TEActivationFunctionBuilder | None:
         """Which module to use for activation function"""
@@ -190,7 +182,7 @@ def activation_func(self) -> TEActivationFunctionBuilder | None:
         return cast(TEActivationFunctionBuilder, TEActivationOp)
 
     def grouped_mlp_modules(
-        self, moe_use_grouped_gemm: bool, moe_use_legacy_grouped_gemm: bool
+        self, moe_use_grouped_gemm: bool
     ) -> Tuple[type, Optional[MLPSubmodules]]:
         """Which module and submodules to use for grouped mlp"""
         return InferenceGroupedMLP, MLPSubmodules(

megatron/core/models/gpt/experimental_attention_variant_module_specs.py

@@ -392,7 +392,6 @@ def _get_self_attention_module_spec(
         moe_grouped_gemm=config.moe_grouped_gemm,
         qk_layernorm=config.qk_layernorm,
         multi_latent_attention=config.multi_latent_attention,
-        moe_use_legacy_grouped_gemm=config.moe_use_legacy_grouped_gemm,
         qk_l2_norm=config.qk_l2_norm,
         use_kitchen=config.use_kitchen,
         use_te_activation_func=config.use_te_activation_func,
@@ -444,7 +443,6 @@ def _get_moe_module_spec(
         backend=backend,
         num_experts=config.num_moe_experts,
         moe_grouped_gemm=config.moe_grouped_gemm,
-        moe_use_legacy_grouped_gemm=config.moe_use_legacy_grouped_gemm,
         use_te_activation_func=config.use_te_activation_func,
     )
     moe_spec.metainfo["fuse_pre_mlp_layernorm"] = False

megatron/core/models/gpt/gpt_layer_specs.py

@@ -93,7 +93,6 @@ def get_gpt_layer_with_inference_submodules(
         backend=backend,
         num_experts=num_experts,
         moe_grouped_gemm=moe_grouped_gemm,
-        moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
         use_te_op_fuser=False,
         use_te_activation_func=False,
     )
@@ -179,7 +178,6 @@ def get_gpt_layer_with_transformer_engine_submodules(
     qk_layernorm: Optional[bool] = False,
     multi_latent_attention: Optional[bool] = False,
     fp8: Optional[str] = None,  # pylint: disable=unused-argument
-    moe_use_legacy_grouped_gemm: Optional[bool] = False,
     qk_l2_norm: Optional[bool] = False,
     use_te_op_fuser: Optional[bool] = False,
     use_kitchen: bool = False,
@@ -197,8 +195,6 @@ def get_gpt_layer_with_transformer_engine_submodules(
         qk_layernorm (bool, optional): To use layernorm for queries/keys. Defaults to False.
         multi_latent_attention (bool, optional): To use MLA. Defaults to False.
         fp8 (str, optional): Deprecated. For temporary Nemo compatibility.
-        moe_use_legacy_grouped_gemm (bool, optional): Force use the legacy GroupedMLP.
-                                                      Defaults to False.
         qk_l2_norm (bool, optional): To use l2 norm for queries/keys. Defaults to False.
         use_te_op_fuser (bool, optional): Use Transformer Engine's operation-based API, which may
                                           enable certain operation fusions. Defaults to False.
@@ -231,7 +227,6 @@ def get_gpt_layer_with_transformer_engine_submodules(
         backend=backend,
         num_experts=num_experts,
         moe_grouped_gemm=moe_grouped_gemm,
-        moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
         use_te_op_fuser=use_te_op_fuser,
         use_te_activation_func=use_te_activation_func,
     )
@@ -318,7 +313,6 @@ def get_gpt_layer_local_submodules(
     qk_layernorm: Optional[bool] = False,
     multi_latent_attention: Optional[bool] = False,
     fp8: Optional[str] = None,  # pylint: disable=unused-argument
-    moe_use_legacy_grouped_gemm: Optional[bool] = False,
     normalization: Optional[str] = None,
     qk_l2_norm: Optional[bool] = False,
     use_kitchen: bool = False,
@@ -334,8 +328,6 @@ def get_gpt_layer_local_submodules(
         qk_layernorm (bool, optional): To use layernorm for queries/keys. Defaults to False.
         multi_latent_attention (bool, optional): To use MLA. Defaults to False.
         fp8 (str, optional): Deprecated. For temporary Nemo compatibility.
-        moe_use_legacy_grouped_gemm (bool, optional): Force use the legacy GroupedMLP.
-                                                      Defaults to False.
         qk_l2_norm (bool, optional): To use l2 norm for queries/keys. Defaults to False.
 
     Returns:
@@ -366,10 +358,7 @@ def get_gpt_layer_local_submodules(
         )
 
     mlp = get_mlp_module_spec_for_backend(
-        backend=backend,
-        num_experts=num_experts,
-        moe_grouped_gemm=moe_grouped_gemm,
-        moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
+        backend=backend, num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm
     )
 
     if multi_latent_attention:
@@ -438,19 +427,14 @@ def _get_mlp_module_spec(
     num_experts: Optional[int] = None,
     moe_grouped_gemm: Optional[bool] = False,
     fp8: Optional[str] = None,  # pylint: disable=unused-argument
-    moe_use_legacy_grouped_gemm: Optional[bool] = False,
 ):
     warnings.warn(
         """This private function is on a deprecation track. Please switch to `get_mlp_module_spec`
         since it will be removed in a future release."""
     )
 
     return get_mlp_module_spec(
-        use_te=use_te,
-        num_experts=num_experts,
-        moe_grouped_gemm=moe_grouped_gemm,
-        fp8=fp8,
-        moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
+        use_te=use_te, num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm, fp8=fp8
     )
 
 
@@ -459,7 +443,6 @@ def get_mlp_module_spec(
     num_experts: Optional[int] = None,
     moe_grouped_gemm: Optional[bool] = False,
     fp8: Optional[str] = None,  # pylint: disable=unused-argument
-    moe_use_legacy_grouped_gemm: Optional[bool] = False,
     use_te_op_fuser: Optional[bool] = False,
 ) -> ModuleSpec:
     """Helper function to get module spec for MLP/MoE"""
@@ -482,7 +465,6 @@ def get_mlp_module_spec(
         backend=TESpecProvider() if use_te else LocalSpecProvider(),
         num_experts=num_experts,
         moe_grouped_gemm=moe_grouped_gemm,
-        moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
         use_te_op_fuser=use_te_op_fuser,
     )
 
@@ -491,7 +473,6 @@ def get_mlp_module_spec_for_backend(
     backend: BackendSpecProvider,
     num_experts: Optional[int] = None,
     moe_grouped_gemm: Optional[bool] = False,
-    moe_use_legacy_grouped_gemm: Optional[bool] = False,
     use_te_op_fuser: Optional[bool] = False,
     use_te_activation_func: bool = False,
 ) -> ModuleSpec:
@@ -520,7 +501,6 @@ def get_mlp_module_spec_for_backend(
             backend=backend,
             num_experts=num_experts,
             moe_grouped_gemm=moe_grouped_gemm,
-            moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
             use_te_activation_func=use_te_activation_func,
         )
 
@@ -541,7 +521,6 @@ def get_gpt_decoder_layer_specs(
             moe_grouped_gemm=False,
             qk_layernorm=config.qk_layernorm,
             multi_latent_attention=config.multi_latent_attention,
-            moe_use_legacy_grouped_gemm=config.moe_use_legacy_grouped_gemm,
             qk_l2_norm=qk_l2_norm,
             use_kitchen=config.use_kitchen,
             use_te_activation_func=config.use_te_activation_func,
@@ -553,7 +532,6 @@ def get_gpt_decoder_layer_specs(
             moe_grouped_gemm=config.moe_grouped_gemm,
             qk_layernorm=config.qk_layernorm,
             multi_latent_attention=config.multi_latent_attention,
-            moe_use_legacy_grouped_gemm=config.moe_use_legacy_grouped_gemm,
             qk_l2_norm=qk_l2_norm,
             use_kitchen=config.use_kitchen,
             use_te_activation_func=config.use_te_activation_func,
@@ -582,7 +560,6 @@ def get_gpt_decoder_layer_specs(
             moe_grouped_gemm=False,
             qk_layernorm=config.qk_layernorm,
             multi_latent_attention=config.multi_latent_attention,
-            moe_use_legacy_grouped_gemm=config.moe_use_legacy_grouped_gemm,
             normalization=normalization,
             qk_l2_norm=qk_l2_norm,
             use_kitchen=config.use_kitchen,
@@ -594,7 +571,6 @@ def get_gpt_decoder_layer_specs(
             moe_grouped_gemm=config.moe_grouped_gemm,
             qk_layernorm=config.qk_layernorm,
             multi_latent_attention=config.multi_latent_attention,
-            moe_use_legacy_grouped_gemm=config.moe_use_legacy_grouped_gemm,
             normalization=normalization,
             qk_l2_norm=qk_l2_norm,
             use_kitchen=config.use_kitchen,

megatron/core/models/gpt/moe_module_specs.py

@@ -19,7 +19,6 @@ def get_moe_module_spec(
     use_te: Optional[bool] = True,
     num_experts: Optional[int] = None,
     moe_grouped_gemm: Optional[bool] = False,
-    moe_use_legacy_grouped_gemm: Optional[bool] = False,
 ) -> ModuleSpec:
     """Helper function to get module spec for MoE.
 
@@ -37,18 +36,14 @@ def get_moe_module_spec(
     else:
         backend = LocalSpecProvider()
     return get_moe_module_spec_for_backend(
-        backend=backend,
-        num_experts=num_experts,
-        moe_grouped_gemm=moe_grouped_gemm,
-        moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm,
+        backend=backend, num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm
     )
 
 
 def get_moe_module_spec_for_backend(
     backend: BackendSpecProvider,
     num_experts: Optional[int] = None,
     moe_grouped_gemm: Optional[bool] = False,
-    moe_use_legacy_grouped_gemm: Optional[bool] = False,
     use_te_activation_func: bool = False,
 ) -> ModuleSpec:
     """Helper function to get module spec for MoE"""
@@ -63,8 +58,7 @@ def get_moe_module_spec_for_backend(
     )
 
     expert_module, expert_submodule = backend.grouped_mlp_modules(
-        moe_grouped_gemm is not None and moe_grouped_gemm,
-        moe_use_legacy_grouped_gemm is not None and moe_use_legacy_grouped_gemm,
+        moe_grouped_gemm is not None and moe_grouped_gemm
     )
     if expert_submodule is not None:
         expert_submodule.activation_func = activation_func
@@ -95,7 +89,7 @@ def get_inference_optimized_moe_spec() -> ModuleSpec:
     backend = InferenceSpecProvider()
     activation_func = backend.activation_func()
 
-    expert_module, expert_submodule = backend.grouped_mlp_modules(True, False)
+    expert_module, expert_submodule = backend.grouped_mlp_modules(True)
     if expert_submodule is not None:
         expert_submodule.activation_func = activation_func
 

megatron/core/models/mamba/mamba_layer_specs.py

@@ -41,7 +41,6 @@
     use_te=True,
     num_experts=8,  # Can be any positive integer (must not be None).
     moe_grouped_gemm=True,
-    moe_use_legacy_grouped_gemm=False,
 )
 
 # Inference-optimized MoE spec