minor cleanup

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Author: kevalmorabia97

CHANGELOG.rst

@@ -24,6 +24,7 @@ NVIDIA Model Optimizer Changelog
 - Add ``get_auto_quantize_config`` API to extract a flat quantization config from ``auto_quantize`` search results, enabling re-quantization at different effective bit targets without re-running calibration.
 - Improve ``auto_quantize`` checkpoint/resume: calibration state is now saved and restored across runs, avoiding redundant calibration when resuming a search.
 - Add NemotronH MoE expert support in ``auto_quantize`` grouping and scoring rules.
+- Support full Transformer Engine spec for Minitron pruning (``mcore_minitron``). Now we no longer need to use custom ModelOpt spec. Note that from user's perspective, this is only internal implementation improvement and does not affect the usage of the pruning workflow.
 
 **Misc**
 

modelopt/torch/nas/plugins/megatron.py

@@ -15,12 +15,14 @@
 
 """Plugin to add NAS/Pruning support for megatron-core Language models like GPT and Mamba."""
 
+import copy
 import types
 from abc import ABC
 from collections.abc import Callable, Sequence
 
 import torch
 import torch.nn as nn
+import transformer_engine as te
 from megatron.core.extensions.transformer_engine import (
     TEColumnParallelLinear,
     TEDotProductAttention,
@@ -29,6 +31,7 @@
 )
 from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
 from megatron.core.models.gpt import GPTModel
+from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec
 from megatron.core.parallel_state import is_pipeline_first_stage, is_pipeline_last_stage
 from megatron.core.tensor_parallel.layers import (
     ColumnParallelLinear,
@@ -43,6 +46,7 @@
 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
+from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_layer import TransformerLayer
 
 from modelopt.torch.nas.modules import DynamicModuleList
@@ -53,7 +57,7 @@
 from modelopt.torch.utils import make_divisible
 
 from ..hparams.concat import build_concat_hp
-from ..modules import _DynamicLayerNorm  # noqa: F401 (re-exported for tests)
+from ..modules import _DynamicLayerNorm
 from ..modules.utils import get_sliced_tensor, get_sliced_tensor_by_slices
 from ..registry import DMRegistry
 from ..traced_hp import TracedHp
@@ -63,6 +67,9 @@
 try:
     import mamba_ssm  # noqa: F401
     from megatron.core.models.mamba import MambaModel
+    from megatron.core.models.mamba.mamba_layer_specs import (
+        mamba_stack_spec as _te_mamba_stack_spec,
+    )
     from megatron.core.ssm.mamba_layer import MambaLayer
     from megatron.core.ssm.mamba_mixer import ExtendedRMSNorm, MambaMixer
 
@@ -72,7 +79,23 @@
 except ImportError:
     HAS_MAMBA = False
 
-__all__ = []
+__all__ = ["get_te_mamba_stack_spec"]
+
+
+# TODO: Maybe upstream this to Megatron-LM
+def get_te_mamba_stack_spec(moe_grouped_gemm: bool = False) -> ModuleSpec:
+    """Return the TE Mamba stack spec."""
+    assert HAS_MAMBA
+    if moe_grouped_gemm:
+        return _te_mamba_stack_spec
+
+    # The upstream TE mamba stack spec hardcodes TEGroupedMLP for MoE.
+    # Replace it with SequentialMLP (TE linear layers, no grouped gemm dependency).
+    te_mamba_stack_spec = copy.deepcopy(_te_mamba_stack_spec)
+    te_mamba_stack_spec.submodules.moe_layer.submodules.mlp = get_moe_module_spec(
+        use_te=True, num_experts=8, moe_grouped_gemm=False
+    )
+    return te_mamba_stack_spec
 
 
 # Local Parallel Linear DynamicModules ##########################################################################
@@ -242,6 +265,22 @@ def export(self) -> torch.nn.Module:
         return super().export()
 
 
+# TE Normalization DynamicModule ###################################################################
+@DMRegistry.register(
+    {te.pytorch.LayerNorm: "te.pytorch.LayerNorm", te.pytorch.RMSNorm: "te.pytorch.RMSNorm"}
+)
+class _DynamicTENorm(_DynamicLayerNorm):
+    """A ``te.pytorch.{Layer/RMS}Norm`` layer with dynamic hyperparams."""
+
+    def _setup(self, *, num_features: TracedHp):
+        """Setup the TENorm dynamic module with pre-defined num_features hparam."""
+        self._register_hparam("num_features", num_features)
+        # register dynamic attributes
+        self._register_dynamic_attribute("weight", self._cut_to_active_features)
+        if hasattr(self, "bias"):  # Bias is not present in RMSNorm
+            self._register_dynamic_attribute("bias", self._cut_to_active_features)
+
+
 # MLP DynamicModule ################################################################################
 @DMRegistry.register(
     {

modelopt/torch/nas/plugins/transformer_engine.py

@@ -23,12 +23,9 @@
 from megatron.bridge.data.builders.hf_dataset import HFDatasetConfig
 from megatron.bridge.data.loaders import setup_data_iterators
 from megatron.bridge.data.utils import get_dataset_provider
-from megatron.bridge.models.gpt_provider import GPTModelProvider, modelopt_transformer_layer_spec
+from megatron.bridge.models.gpt_provider import GPTModelProvider
 from megatron.bridge.models.hf_pretrained.utils import is_safe_repo
-from megatron.bridge.models.mamba.mamba_provider import (
-    MambaModelProvider,
-    modelopt_mamba_stack_spec,
-)
+from megatron.bridge.models.mamba.mamba_provider import MambaModelProvider
 from megatron.bridge.models.nemotronh.nemotron_h_provider import NemotronHModelProvider
 from megatron.bridge.training.config import (
     CheckpointConfig,
@@ -50,6 +47,7 @@
 from megatron.core.utils import unwrap_model
 from transformers import AutoTokenizer
 
+from modelopt.torch.nas.plugins.megatron import get_te_mamba_stack_spec
 from modelopt.torch.utils import get_dataset_samples, print_rank_0, warn_rank_0
 
 __all__ = ["get_hf_mbridge_calibration_loop", "load_mbridge_model_from_hf"]
@@ -94,12 +92,9 @@ def load_mbridge_model_from_hf(
             assert hasattr(provider, key), f"{type(provider)} does not have attribute {key}"
             setattr(provider, key, value)
 
-    print_rank_0("Setting ModelOpt spec for model provider")
     if isinstance(provider, MambaModelProvider):
-        provider.mamba_stack_spec = modelopt_mamba_stack_spec
-    else:
-        provider.transformer_layer_spec = modelopt_transformer_layer_spec
-
+        # disable moe_grouped_gemm in default TE spec until its supported
+        provider.mamba_stack_spec = get_te_mamba_stack_spec(moe_grouped_gemm=False)
     provider.finalize()
     if init_model_parallel:
         provider.initialize_model_parallel(seed=0)

modelopt/torch/utils/plugins/mbridge.py

@@ -12,7 +12,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import copy
 from warnings import warn
 
 import torch
@@ -25,14 +24,14 @@
     get_gpt_layer_local_spec,
     get_gpt_layer_with_transformer_engine_spec,
 )
-from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec
 from megatron.core.models.mamba import MambaModel
 from megatron.core.parallel_state import is_pipeline_first_stage, is_pipeline_last_stage
 from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.transformer_config import TransformerConfig
 
 from modelopt.torch.export.unified_export_megatron import import_mcore_gpt_from_hf
+from modelopt.torch.nas.plugins.megatron import get_te_mamba_stack_spec
 
 try:
     from megatron.core.extensions.transformer_engine import TENorm
@@ -44,19 +43,9 @@
     HAS_TE = False
 
 try:
-    from megatron.core.models.mamba.mamba_layer_specs import (
-        mamba_stack_spec as _te_mamba_stack_spec,
-    )
     from megatron.core.post_training.modelopt.mamba.model_specs import get_mamba_stack_modelopt_spec
     from megatron.core.ssm.mamba_layer import MambaLayer  # noqa: F401
 
-    # The upstream TE mamba stack spec hardcodes TEGroupedMLP for MoE.
-    # Replace it with SequentialMLP (TE linear layers, no grouped gemm dependency).
-    te_mamba_stack_spec = copy.deepcopy(_te_mamba_stack_spec)
-    te_mamba_stack_spec.submodules.moe_layer.submodules.mlp = get_moe_module_spec(
-        use_te=True, num_experts=8, moe_grouped_gemm=False
-    )
-
     HAS_MAMBA = True
 except ImportError as e:
     warn(f"Mamba not installed: {e}")
@@ -152,6 +141,7 @@ def get_mcore_gpt_model(
     bf16: bool = True,
     use_te: bool = False,
     # MoE-specific parameters
+    moe_grouped_gemm: bool = False,
     moe_ffn_hidden_size: int | None = None,
     moe_shared_expert_intermediate_size: int | None = None,
     num_moe_experts: int | None = None,
@@ -195,6 +185,7 @@ def squared_relu(x):
         bf16=bf16,
         # MoE-specific parameters
         moe_router_dtype=None,
+        moe_grouped_gemm=moe_grouped_gemm,
         moe_ffn_hidden_size=moe_ffn_hidden_size,
         moe_shared_expert_intermediate_size=moe_shared_expert_intermediate_size,
         moe_router_enable_expert_bias=True,
@@ -217,6 +208,7 @@ def squared_relu(x):
         assert HAS_APEX, "Apex not installed"
         transformer_layer_spec = get_gpt_layer_local_spec(
             num_experts=num_moe_experts,
+            moe_grouped_gemm=moe_grouped_gemm,
             normalization=normalization,
         )
     else:
@@ -320,6 +312,7 @@ def get_mcore_mamba_hybrid_model(
     mamba_num_groups: int = 2,
     # MoE-specific parameters
     skip_moe: bool = False,
+    moe_grouped_gemm: bool = False,
     moe_ffn_hidden_size: int | None = 64,
     moe_shared_expert_intermediate_size: int | None = 32,
     num_moe_experts: int | None = 8,
@@ -353,6 +346,7 @@ def get_mcore_mamba_hybrid_model(
         mamba_head_dim=mamba_head_dim,
         mamba_num_groups=mamba_num_groups,
         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,
         add_bias_linear=False,
@@ -391,7 +385,7 @@ def get_mcore_mamba_hybrid_model(
     print(f"Using `{hybrid_override_pattern=}` for building MambaModel")
 
     if transformer_impl == "transformer_engine":
-        mamba_spec = te_mamba_stack_spec
+        mamba_spec = get_te_mamba_stack_spec(moe_grouped_gemm=moe_grouped_gemm)
     else:
         mamba_spec = get_mamba_stack_modelopt_spec(remap_te_layernorm=True)
 

tests/_test_utils/torch/megatron/models.py

@@ -31,11 +31,11 @@
     _DynamicColumnParallelLinear,
     _DynamicEmbedding,
     _DynamicExtendedRMSNorm,
-    _DynamicLayerNorm,
     _DynamicMambaLayer,
     _DynamicMambaMixer,
     _DynamicMCoreLanguageModel,
     _DynamicTELayerNormColumnParallelLinear,
+    _DynamicTENorm,
     _DynamicTERowParallelLinear,
 )
 from modelopt.torch.nas.traced_hp import TracedHp
@@ -104,7 +104,7 @@ def _test_mamba_search_space(rank, size):
         if layer.mixer.rmsnorm:
             assert isinstance(layer.mixer.norm, _DynamicExtendedRMSNorm)
     if is_pipeline_last_stage():
-        assert isinstance(model.decoder.final_norm, _DynamicLayerNorm)
+        assert isinstance(model.decoder.final_norm, _DynamicTENorm)
         assert isinstance(model.output_layer, _DynamicColumnParallelLinear)
 
     # NOTE: `search_space_size` does not reduce across TP/PP groups