From 16e4ed634194df70d30763c69b711962cb444a08 Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Sun, 23 Nov 2025 11:11:32 +0000
Subject: [PATCH 01/48] added mtp support

---
 gpt_builders.py                               |  39 +--
 mamba_builders.py                             |  13 +
 .../common/language_module/language_module.py |  37 ++-
 megatron/core/models/gpt/gpt_layer_specs.py   |   2 +-
 megatron/core/models/gpt/gpt_model.py         |  64 +----
 .../core/models/mamba/mamba_layer_specs.py    |  83 +++++++
 megatron/core/models/mamba/mamba_model.py     | 233 +++++++++++++++++-
 megatron/core/pipeline_parallel/schedules.py  |   4 +-
 megatron/core/ssm/mamba_block.py              |   7 +-
 .../transformer/multi_token_prediction.py     | 177 ++++++++++---
 megatron/training/arguments.py                |  28 +++
 megatron/training/checkpointing.py            |   1 +
 megatron/training/training.py                 |  12 +-
 pretrain_mamba.py                             |   2 +-
 .../test_multi_token_prediction.py            |   4 +-
 15 files changed, 590 insertions(+), 116 deletions(-)

diff --git a/gpt_builders.py b/gpt_builders.py
index 4b86e30e597..9ba7ac8f446 100644
--- a/gpt_builders.py
+++ b/gpt_builders.py
@@ -8,6 +8,7 @@
     get_gpt_layer_with_inference_spec,
     get_gpt_mtp_block_spec,
 )
+from megatron.core.models.mamba.mamba_layer_specs import get_mamba_mtp_block_spec
 from megatron.core.models.gpt.heterogeneous.heterogeneous_layer_specs import (
     get_gpt_heterogeneous_layer_spec,
 )
@@ -60,22 +61,29 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_
                 transformer_layer_spec = _get_transformer_layer_spec(use_te, config)
         mtp_block_spec = None
         if args.mtp_num_layers is not None:
-            assert not (config.transformer_impl == "inference_optimized")
-            if (
-                hasattr(transformer_layer_spec, 'layer_specs')
-                and len(transformer_layer_spec.layer_specs) == 0
-            ):
-                # Get the decoder layer spec explicitly if no decoder layer in the last stage,
-                # Only happens with block spec (TransformerBlockSubmodules) when using MoE.
-                transformer_layer_spec_for_mtp = _get_transformer_layer_spec(use_te, config)
+            if args.mtp_hybrid_override_pattern is not None:
+                # Uses the hybrid model for the MTP module.
+                mtp_mamba_stack_spec = import_module(args.mtp_spec)
+                mtp_block_spec = get_mamba_mtp_block_spec(
+                    config, mtp_mamba_stack_spec, use_transformer_engine=use_te, vp_stage=vp_stage
+                )
             else:
-                transformer_layer_spec_for_mtp = transformer_layer_spec
-            mtp_block_spec = get_gpt_mtp_block_spec(
-                config,
-                transformer_layer_spec_for_mtp,
-                use_transformer_engine=use_te,
-                vp_stage=vp_stage,
-            )
+                assert not (config.transformer_impl == "inference_optimized")
+                if (
+                    hasattr(transformer_layer_spec, 'layer_specs')
+                    and len(transformer_layer_spec.layer_specs) == 0
+                ):
+                    # Get the decoder layer spec explicitly if no decoder layer in the last stage,
+                    # Only happens with block spec (TransformerBlockSubmodules) when using MoE.
+                    transformer_layer_spec_for_mtp = _get_transformer_layer_spec(use_te, config)
+                else:
+                    transformer_layer_spec_for_mtp = transformer_layer_spec
+                mtp_block_spec = get_gpt_mtp_block_spec(
+                    config,
+                    transformer_layer_spec_for_mtp,
+                    use_transformer_engine=use_te,
+                    vp_stage=vp_stage,
+                )
 
         model = GPTModel(
             config=config,
@@ -94,6 +102,7 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_
             mtp_block_spec=mtp_block_spec,
             vp_stage=vp_stage,
             pg_collection=pg_collection,
+            mtp_hybrid_override_pattern=args.mtp_hybrid_override_pattern,
         )
 
     return model
diff --git a/mamba_builders.py b/mamba_builders.py
index 6a792ba6ea5..ba769d43ca2 100644
--- a/mamba_builders.py
+++ b/mamba_builders.py
@@ -1,6 +1,7 @@
 # Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 
 from model_provider import count_parameters_in_layer
+from megatron.core.models.mamba.mamba_layer_specs import get_mamba_mtp_block_spec
 from megatron.core.models.mamba import MambaModel
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.spec_utils import import_module
@@ -22,6 +23,15 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
     else:
         raise ValueError("You must provide a valid Mamba layer spec via --spec")
 
+    mtp_block_spec = None
+    if args.mtp_num_layers is not None:
+        use_te = args.transformer_impl == "transformer_engine"
+        assert args.mtp_spec is not None
+        assert args.mtp_hybrid_override_pattern is not None, "We need to set the override hybrid pattern for MTP layers!"
+        mtp_mamba_stack_spec = import_module(args.mtp_spec)
+        mtp_block_spec = get_mamba_mtp_block_spec(
+            config, mtp_mamba_stack_spec, use_transformer_engine=use_te, vp_stage=vp_stage
+        )
     model = MambaModel(
         config=config,
         mamba_stack_spec=mamba_stack_spec,
@@ -31,6 +41,7 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
         hybrid_attention_ratio=args.hybrid_attention_ratio,
         hybrid_mlp_ratio=args.hybrid_mlp_ratio,
         hybrid_override_pattern=args.hybrid_override_pattern,
+        mtp_hybrid_override_pattern=args.mtp_hybrid_override_pattern,
         post_process=post_process,
         fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
         parallel_output=True,
@@ -39,6 +50,8 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
         rotary_percent=args.rotary_percent,
         rotary_base=args.rotary_base,
         pg_collection=pg_collection,
+        vp_stage=vp_stage,
+        mtp_block_spec=mtp_block_spec,
     )
 
     for l in range(model.decoder.num_layers_per_pipeline_rank):
diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
index 19eea55dec3..706f8429a29 100644
--- a/megatron/core/models/common/language_module/language_module.py
+++ b/megatron/core/models/common/language_module/language_module.py
@@ -31,6 +31,10 @@
     make_tp_sharded_tensor_for_checkpoint,
 )
 
+from megatron.core.transformer.multi_token_prediction import (
+    tie_output_layer_state_dict,
+    tie_word_embeddings_state_dict,
+)
 
 class LanguageModule(MegatronModule):
     """Base language module that has common helper functions used across GPT, BERT etc.
@@ -250,12 +254,20 @@ def setup_embeddings_and_output_layer(self) -> None:
             LanguageModule.embedding_warning_printed = True
 
     def shared_embedding_or_output_weight(self) -> Tensor:
-        """Gets the emedding weight or output logit weights when share embedding and output weights set to True.
+        """Gets the emedding weight or output logit weights when share embedding and output weights set to True
+          or when use Multi-Token Prediction (MTP).
 
         Returns:
-            Tensor: During pre processing it returns the input embeddings weight while during post processing it returns the final output layers weight
+            Tensor: During pre processing or MTP process it returns the input embeddings weight while during post processing it returns the final output layers weight
         """
-        if self.pre_process:
+        if self.pre_process or self.mtp_process:
+            # Multi-Token Prediction (MTP) need both embedding layer and output layer.
+            # So there will be both embedding layer and output layer in the mtp process stage.
+            # In this case, if share_embeddings_and_output_weights is True, the shared weights
+            # will be stored in embedding layer, and output layer will not have any weight.
+            assert hasattr(
+                self, 'embedding'
+            ), f"embedding is needed in this pipeline stage, but it is not initialized."
             return self.embedding.word_embeddings.weight
         elif self.post_process:
             return self.output_layer.weight
@@ -288,6 +300,25 @@ def sharded_state_dict(
         output_layer_weight_key = f'{prefix}output_layer.weight'
         output_layer_bias_key = f'{prefix}output_layer.bias'
 
+        # Multi-Token Prediction (MTP) need both embedding layer and output layer in
+        # mtp process stage.
+        # If MTP is not placed in the pre processing stage, we need to maintain a copy of
+        # embedding layer in the mtp process stage and tie it to the embedding in the pre
+        # processing stage.
+        # Also, if MTP is not placed in the post processing stage, we need to maintain a copy
+        # of output layer in the mtp process stage and tie it to the output layer in the post
+        # processing stage.
+        if self.mtp_process and not self.pre_process:
+            emb_weight = self.embedding.word_embeddings.weight
+            tie_word_embeddings_state_dict(sharded_state_dict, emb_weight, first_stage_word_emb_key)
+
+        if self.mtp_process and not self.post_process:
+            # We only need to tie the output layer weight if share_embeddings_and_output_weights
+            # is False. Because if share_embeddings_and_output_weights is True, the shared weight
+            # will be stored in embedding layer, and output layer will not have any weight.
+            if not self.share_embeddings_and_output_weights:
+                output_layer_weight = self.output_layer.weight
+                tie_output_layer_state_dict(sharded_state_dict, output_layer_weight, output_layer_weight_key)
         if self.share_embeddings_and_output_weights:
             self.tie_embeddings_and_output_weights_state_dict(
                 sharded_state_dict, output_layer_weight_key, first_stage_word_emb_key, metadata
diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
index b61a544c5df..900232fa708 100755
--- a/megatron/core/models/gpt/gpt_layer_specs.py
+++ b/megatron/core/models/gpt/gpt_layer_specs.py
@@ -686,7 +686,7 @@ def get_gpt_mtp_block_spec_for_backend(
         raise ValueError(f"Invalid spec: {spec}")
 
     mtp_layer_spec = get_mtp_layer_spec_for_backend(
-        transformer_layer_spec=transformer_layer_spec, backend=backend
+        mtp_model_layer_spec=transformer_layer_spec, backend=backend
     )
     mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0
     mtp_layer_specs = [mtp_layer_spec] * mtp_num_layers
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index 9342e96ce9a..eab4c258427 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -102,6 +102,7 @@ def __init__(
         mtp_block_spec: Optional[ModuleSpec] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
         vp_stage: Optional[int] = None,
+        mtp_hybrid_override_pattern: str = None,
     ) -> None:
         super().__init__(config=config, pg_collection=pg_collection)
 
@@ -210,7 +211,7 @@ def __init__(
 
         if self.mtp_process:
             self.mtp = MultiTokenPredictionBlock(
-                config=self.config, spec=self.mtp_block_spec, vp_stage=vp_stage
+                config=self.config, spec=self.mtp_block_spec, vp_stage=vp_stage, mtp_hybrid_override_pattern=mtp_hybrid_override_pattern
             )
 
         # Output
@@ -532,6 +533,7 @@ def _postprocess(
                 packed_seq_params=packed_seq_params,
                 sequence_len_offset=sequence_len_offset,
                 embedding=self.embedding,
+                packed_seq_params=packed_seq_params,
                 **(extra_block_kwargs or {}),
             )
 
@@ -539,56 +541,16 @@ def _postprocess(
             return hidden_states
 
         if self.mtp_process:
-            mtp_labels = labels.clone()
-            hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0)
-            hidden_states = hidden_states_list[0]
-            if loss_mask is None:
-                # if loss_mask is not provided, use all ones as loss_mask
-                loss_mask = torch.ones_like(mtp_labels)
-            for mtp_layer_number in range(self.config.mtp_num_layers):
-                # output
-                mtp_logits, _ = self.output_layer(
-                    hidden_states_list[mtp_layer_number + 1],
-                    weight=output_weight,
-                    runtime_gather_output=runtime_gather_output,
-                )
-                # Calc loss for the current Multi-Token Prediction (MTP) layers.
-                mtp_labels, _ = roll_tensor(
-                    mtp_labels,
-                    shifts=-1,
-                    dims=-1,
-                    cp_group=self.cp_group,
-                    packed_seq_params=packed_seq_params,
-                )
-                loss_mask, num_tokens = roll_tensor(
-                    loss_mask,
-                    shifts=-1,
-                    dims=-1,
-                    cp_group=self.cp_group,
-                    packed_seq_params=packed_seq_params,
-                )
-                mtp_loss = self.compute_language_model_loss(mtp_labels, mtp_logits)
-                mtp_loss = loss_mask * mtp_loss
-                if self.training:
-                    # TODO(shifangx): remove the use of parallel_state here
-                    # after moving loss logging to loss_func in pretrain_gpt.py
-                    MTPLossLoggingHelper.save_loss_to_tracker(
-                        torch.sum(mtp_loss) / num_tokens,
-                        mtp_layer_number,
-                        self.config.mtp_num_layers,
-                        avg_group=parallel_state.get_data_parallel_group(
-                            with_context_parallel=True
-                        ),
-                    )
-                mtp_loss_scale = self.config.mtp_loss_scaling_factor / self.config.mtp_num_layers
-                if self.config.calculate_per_token_loss:
-                    hidden_states = MTPLossAutoScaler.apply(
-                        hidden_states, mtp_loss_scale * mtp_loss
-                    )
-                else:
-                    hidden_states = MTPLossAutoScaler.apply(
-                        hidden_states, mtp_loss_scale * mtp_loss / num_tokens
-                    )
+            hidden_states = self.mtp.process_loss(
+                hidden_states=hidden_states,
+                labels=labels,
+                loss_mask=loss_mask,
+                output_layer=self.output_layer,
+                output_weight=output_weight,
+                runtime_gather_output=runtime_gather_output,
+                is_training=self.training,
+                compute_language_model_loss=self.compute_language_model_loss,
+            )
         sequence_parallel_override = False
         if in_inference_mode and inference_context.materialize_only_last_token_logits:
             if inference_context.is_static_batching():
diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index f83275ed9c6..286d72b36ff 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -6,6 +6,7 @@
     TENorm,
     TERowParallelLinear,
 )
+from typing import Optional
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec
 from megatron.core.ssm.mamba_block import MambaStack, MambaStackSubmodules
@@ -21,6 +22,15 @@
 from megatron.core.transformer.mlp import MLP, MLPSubmodules
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
+from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionBlockSubmodules
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.multi_token_prediction import (
+    MultiTokenPredictionBlockSubmodules,
+    get_mtp_layer_offset,
+    get_mtp_layer_spec_for_backend,
+    get_mtp_num_layers_to_build,
+)
+from megatron.core.models.backends import BackendSpecProvider, LocalSpecProvider
 
 moe = get_moe_module_spec(
     use_te=True,
@@ -29,6 +39,22 @@
     moe_use_legacy_grouped_gemm=False,
 )
 
+try:
+    from megatron.core.extensions.transformer_engine import TEFusedMLP, TENorm
+    from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
+
+    HAVE_TE = True
+except ImportError:
+    HAVE_TE = False
+
+
+try:
+    from megatron.core.extensions.kitchen import KitchenSpecProvider
+
+    HAVE_KITCHEN = True
+except ImportError:
+    HAVE_KITCHEN = False
+
 mamba_stack_spec = ModuleSpec(
     module=MambaStack,
     submodules=MambaStackSubmodules(
@@ -87,6 +113,7 @@
     ),
 )
 
+<<<<<<< HEAD
 mamba_inference_stack_spec = ModuleSpec(
     module=MambaStack,
     submodules=MambaStackSubmodules(
@@ -146,3 +173,59 @@
         ),
     ),
 )
+=======
+
+def get_mamba_mtp_block_spec(
+    config: TransformerConfig,
+    spec: ModuleSpec,
+    use_transformer_engine: bool, 
+    vp_stage: Optional[int] = None,
+) -> MultiTokenPredictionBlockSubmodules:
+    """Mamba Multi-Token Prediction (MTP) block spec."""
+    if use_transformer_engine:
+        backend: BackendSpecProvider = (
+            KitchenSpecProvider(fallback=TESpecProvider())
+            if config.use_kitchen
+            else TESpecProvider()
+        )
+    else:
+        backend = (
+            KitchenSpecProvider(fallback=LocalSpecProvider())
+            if config.use_kitchen
+            else LocalSpecProvider()
+        )
+    return get_mamba_mtp_block_spec_for_backend(
+        config=config, spec=spec, backend=backend, vp_stage=vp_stage
+    )
+    
+def get_mamba_mtp_block_spec_for_backend(
+    config: TransformerConfig,
+    spec: ModuleSpec, 
+    backend: BackendSpecProvider,
+    vp_stage: Optional[int] = None,
+) -> MultiTokenPredictionBlockSubmodules:
+    """Mamba Multi-Token Prediction (MTP) block spec."""
+    num_layers_to_build = get_mtp_num_layers_to_build(config, vp_stage=vp_stage)
+    if num_layers_to_build == 0:
+        return None
+
+    mtp_layer_spec = get_mtp_layer_spec_for_backend(
+        mtp_model_layer_spec=spec, backend=backend
+    )
+    mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0
+    mtp_layer_specs = [mtp_layer_spec] * mtp_num_layers
+
+    offset = get_mtp_layer_offset(config)
+    # split the mtp layer specs to only include the layers that are built in this pipeline stage.
+    mtp_layer_specs = mtp_layer_specs[offset : offset + num_layers_to_build]
+    if len(mtp_layer_specs) > 0:
+        assert (
+            len(mtp_layer_specs) == config.mtp_num_layers
+        ), +f"currently all of the mtp layers must stage in the same pipeline stage."
+        mtp_block_spec = MultiTokenPredictionBlockSubmodules(layer_specs=mtp_layer_specs)
+    else:
+        mtp_block_spec = None
+
+    return mtp_block_spec
+
+>>>>>>> a05cbed31 (added mtp support)
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 378cf7e47d6..6521100648b 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -3,8 +3,10 @@
 from typing import Literal, Optional
 
 from torch import Tensor
-
-from megatron.core import tensor_parallel
+import torch
+from typing import Dict
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core import parallel_state, tensor_parallel
 from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
 from megatron.core.inference.contexts import BaseInferenceContext
 from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
@@ -17,6 +19,15 @@
 from megatron.core.transformer.enums import ModelType
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.utils import WrappedTensor, deprecate_inference_params
+from megatron.core.transformer.multi_token_prediction import (
+    MTPLossAutoScaler,
+    MTPLossLoggingHelper,
+    MultiTokenPredictionBlock,
+    roll_tensor,
+    tie_output_layer_state_dict,
+    tie_word_embeddings_state_dict,
+)
+from megatron.core.packed_seq_params import PackedSeqParams
 
 
 class MambaModel(LanguageModule):
@@ -63,6 +74,7 @@ def __init__(
         hybrid_attention_ratio: float = 0.0,
         hybrid_mlp_ratio: float = 0.0,
         hybrid_override_pattern: str = None,
+        mtp_hybrid_override_pattern: str = None,
         post_process: bool = True,
         fp16_lm_cross_entropy: bool = False,
         parallel_output: bool = True,
@@ -74,6 +86,8 @@ def __init__(
         scatter_embedding_sequence_parallel: bool = True,
         seq_len_interpolation_factor: Optional[float] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
+        vp_stage: Optional[int] = None,
+        mtp_block_spec: Optional[ModuleSpec] = None,
     ) -> None:
         super().__init__(config=config, pg_collection=pg_collection)
 
@@ -87,17 +101,24 @@ def __init__(
         self.hybrid_attention_ratio = hybrid_attention_ratio
         self.hybrid_mlp_ratio = hybrid_mlp_ratio
         self.hybrid_override_pattern = hybrid_override_pattern
+        self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern
         self.post_process = post_process
         self.fp16_lm_cross_entropy = fp16_lm_cross_entropy
         self.parallel_output = parallel_output
         self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
         self.position_embedding_type = position_embedding_type
+        self.mtp_block_spec = mtp_block_spec
+        self.mtp_process = mtp_block_spec is not None
+
+        # Cache for RoPE tensors which do not change between iterations.
+        self.rotary_pos_emb_cache = {}
+
 
         # megatron core pipelining currently depends on model type
         # TODO: remove this dependency ?
         self.model_type = ModelType.encoder_or_decoder
 
-        if self.pre_process:
+        if self.pre_process or self.mtp_process:
             self.embedding = LanguageModelEmbedding(
                 config=self.config,
                 vocab_size=self.vocab_size,
@@ -129,8 +150,13 @@ def __init__(
             pg_collection=self.pg_collection,
         )
 
+        if self.mtp_process:
+            self.mtp = MultiTokenPredictionBlock(
+                config=self.config, spec=self.mtp_block_spec, vp_stage=vp_stage, mtp_hybrid_override_pattern=self.mtp_hybrid_override_pattern
+            )
+
         # Output
-        if post_process:
+        if post_process or self.mtp_process:
             self.output_layer = tensor_parallel.ColumnParallelLinear(
                 config.hidden_size,
                 self.vocab_size,
@@ -168,6 +194,94 @@ def set_input_tensor(self, input_tensor: Tensor) -> None:
         assert len(input_tensor) == 1, 'input_tensor should only be length 1 for gpt/bert'
         self.decoder.set_input_tensor(input_tensor[0])
 
+    def _preprocess(
+        self,
+        input_ids: Tensor,
+        position_ids: Tensor,
+        decoder_input: Tensor = None,
+        inference_context: BaseInferenceContext = None,
+        packed_seq_params: PackedSeqParams = None,
+    ):
+        """Preprocesses inputs for the transformer decoder.
+
+        Applies embeddings to input tokens, or uses `decoder_input` from a previous
+        pipeline stage. Also sets up rotary positional embeddings.
+        """
+
+        # If decoder_input is provided (not None), then input_ids and position_ids are ignored.
+        # Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input.
+
+        # Decoder embedding.
+        if decoder_input is not None:
+            pass
+        elif self.pre_process:
+            decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
+        else:
+            # intermediate stage of pipeline
+            # decoder will get hidden_states from encoder.input_tensor
+            decoder_input = None
+
+        # Rotary positional embeddings (embedding is None for PP intermediate devices)
+        rotary_pos_emb = None
+        rotary_pos_cos = None
+        rotary_pos_sin = None
+        if self.position_embedding_type == 'rope' and not self.config.multi_latent_attention:
+            if not self.training and self.config.flash_decode and inference_context:
+                assert (
+                    inference_context.is_static_batching()
+                ), "GPTModel currently only supports static inference batching."
+                # Flash decoding uses precomputed cos and sin for RoPE
+                rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb_cache.setdefault(
+                    inference_context.max_sequence_length,
+                    self.rotary_pos_emb.get_cos_sin(inference_context.max_sequence_length),
+                )
+            else:
+                rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
+                    inference_context, self.decoder, decoder_input, self.config, packed_seq_params
+                )
+                rotary_pos_emb = self.rotary_pos_emb(
+                    rotary_seq_len,
+                    packed_seq=packed_seq_params is not None
+                    and packed_seq_params.qkv_format == 'thd',
+                )
+        elif self.position_embedding_type == 'mrope' and not self.config.multi_latent_attention:
+            if self.training or not self.config.flash_decode:
+                rotary_pos_emb = self.rotary_pos_emb(position_ids, self.mrope_section)
+            else:
+                # Flash decoding uses precomputed cos and sin for RoPE
+                raise NotImplementedError(
+                    "Flash decoding uses precomputed cos and sin for RoPE, not implmented in "
+                    "MultimodalRotaryEmbedding yet."
+                )
+
+        if (
+            (self.config.enable_cuda_graph or self.config.flash_decode)
+            and rotary_pos_cos is not None
+            and inference_context
+            and inference_context.is_static_batching()
+            and not self.training
+        ):
+            current_batch_size = input_ids.shape[0]
+            sequence_len_offset = torch.tensor(
+                [inference_context.sequence_len_offset] * current_batch_size,
+                dtype=torch.int32,
+                device=rotary_pos_cos.device,  # Co-locate this with the rotary tensors
+            )
+        else:
+            sequence_len_offset = None
+
+        # Wrap decoder_input to allow the decoder (TransformerBlock) to delete the
+        # reference held by this caller function, enabling early garbage collection for
+        # inference. Skip wrapping if decoder_input is logged after decoder completion.
+        if (
+            inference_context is not None
+            and not self.training
+            and not has_config_logger_enabled(self.config)
+        ):
+            decoder_input = WrappedTensor(decoder_input)
+
+        return decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset
+
     def forward(
         self,
         input_ids: Tensor,
@@ -179,6 +293,7 @@ def forward(
         runtime_gather_output: Optional[bool] = None,
         *,
         inference_params: Optional[BaseInferenceContext] = None,
+        loss_mask: Optional[Tensor] = None,
     ) -> Tensor:
         """Forward function of the Mamba model. This function passes the input tensors
         through the embedding layer, and then the decoder and finally into the post
@@ -237,9 +352,53 @@ def forward(
             rotary_pos_emb=rotary_pos_emb,
         )
 
+        # Get output weight for shared embeddings
+        output_weight = None
+        if self.share_embeddings_and_output_weights:
+            output_weight = self.shared_embedding_or_output_weight()            
+
+        if self.mtp_process:
+            extra_block_kwargs = None
+            packed_seq_params = None
+            decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset = self._preprocess(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                decoder_input=decoder_input,
+                inference_context=inference_context,
+                packed_seq_params=packed_seq_params,
+            )
+
+
+            hidden_states = self.mtp(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                inference_params=inference_params,
+                rotary_pos_emb=rotary_pos_emb,
+                rotary_pos_cos=rotary_pos_cos,
+                rotary_pos_sin=rotary_pos_sin,
+                packed_seq_params=packed_seq_params,
+                sequence_len_offset=sequence_len_offset,
+                embedding=self.embedding,
+                **(extra_block_kwargs or {}),
+            )
+
         if not self.post_process:
             return hidden_states
 
+        if self.mtp_process:
+            hidden_states = self.mtp.process_loss(
+                hidden_states=hidden_states,
+                labels=labels,
+                loss_mask=loss_mask,
+                output_layer=self.output_layer,
+                output_weight=output_weight,
+                runtime_gather_output=runtime_gather_output,
+                is_training=self.training,
+                compute_language_model_loss=self.compute_language_model_loss,
+            )
+
         # logits and loss
         output_weight = None
         if self.share_embeddings_and_output_weights:
@@ -260,7 +419,7 @@ def forward(
                     self.output_layer.sequence_parallel = False
                     sequence_parallel_override = True
 
-                # Reshape [B, 1, H] to [1, B, H] → extract each sample’s true last‐token hidden
+                # Reshape [B, 1, H] to [1, B, H] → extract each sample's true last‐token hidden
                 # state ([B, H]) → unsqueeze back to [B, 1, H]
                 # (so that the output layer, which expects S×B×H, receives only the final token)
                 hidden_states = inference_context.last_token_logits(
@@ -287,3 +446,67 @@ def forward(
         loss = self.compute_language_model_loss(labels, logits)
 
         return loss
+
+    def shared_embedding_or_output_weight(self) -> Optional[Tensor]:
+        """Gets the embedding weight or output logit weights when share input embedding and
+        output weights set to True or when use Multi-Token Prediction (MTP) feature.
+
+        Returns:
+            Optional[Tensor]: During pre processing or MTP process it returns the input embeddings weight.
+            Otherwise, during post processing it returns the final output layers weight.
+        """
+        if self.pre_process or self.mtp_process:
+            # Multi-Token Prediction (MTP) need both embedding layer and output layer.
+            # So there will be both embedding layer and output layer in the mtp process stage.
+            # In this case, if share_embeddings_and_output_weights is True, the shared weights
+            # will be stored in embedding layer, and output layer will not have any weight.
+            assert hasattr(
+                self, 'embedding'
+            ), f"embedding is needed in this pipeline stage, but it is not initialized."
+            return self.embedding.word_embeddings.weight
+        elif self.post_process:
+            return self.output_layer.weight
+        return None
+
+    def sharded_state_dict(
+        self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[Dict] = None
+    ) -> ShardedStateDict:
+        """Sharded state dict implementation for MambaModel backward-compatibility.
+
+        Removing extra state.
+        Tie word embeddings and output layer in mtp process stage.
+
+        Args:
+            prefix (str): Module name prefix.
+            sharded_offsets (tuple): PP related offsets, expected to be empty at this module level.
+            metadata (Optional[Dict]): metadata controlling sharded state dict creation.
+
+        Returns:
+            ShardedStateDict: sharded state dict for the MambaModel
+        """
+        sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+        
+        # Multi-Token Prediction (MTP) need both embedding layer and output layer in
+        # mtp process stage.
+        # If MTP is not placed in the pre processing stage, we need to maintain a copy of
+        # embedding layer in the mtp process stage and tie it to the embedding in the pre
+        # processing stage.
+        # Also, if MTP is not placed in the post processing stage, we need to maintain a copy
+        # of output layer in the mtp process stage and tie it to the output layer in the post
+        # processing stage.
+        if self.mtp_process and not self.pre_process:
+            emb_weight_key = f'{prefix}embedding.word_embeddings.weight'
+            emb_weight = self.embedding.word_embeddings.weight
+            tie_word_embeddings_state_dict(sharded_state_dict, emb_weight, emb_weight_key)
+        if self.mtp_process and not self.post_process:
+            # We only need to tie the output layer weight if share_embeddings_and_output_weights
+            # is False. Because if share_embeddings_and_output_weights is True, the shared weight
+            # will be stored in embedding layer, and output layer will not have any weight.
+            if not self.share_embeddings_and_output_weights:
+                output_layer_weight_key = f'{prefix}output_layer.weight'
+                output_layer_weight = self.output_layer.weight
+                tie_output_layer_state_dict(
+                    sharded_state_dict, output_layer_weight, output_layer_weight_key
+                )
+
+        return sharded_state_dict
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index 5ab4cb57545..3c1fb2e3b18 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -193,7 +193,7 @@ def set_current_microbatch(model, microbatch_id):
             layer.current_microbatch = microbatch_id
         if hasattr(model_with_decoder, 'mtp'):
             for layer in model_with_decoder.mtp.layers:
-                layer.transformer_layer.current_microbatch = microbatch_id
+                layer.mtp_model_layer.current_microbatch = microbatch_id
 
 
 def forward_step_calc_loss(
@@ -384,7 +384,7 @@ def forward_step(
 
     if is_first_microbatch and hasattr(model, 'set_is_first_microbatch'):
         model.set_is_first_microbatch()
-    if current_microbatch is not None:
+    if config.external_cuda_graph and current_microbatch is not None:
         set_current_microbatch(model, current_microbatch)
 
     unwrap_output_tensor = False
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index 1abd12186e3..4e8bc27cc54 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -83,12 +83,15 @@ def __init__(
         device=None,
         dtype=None,
         pg_collection: ProcessGroupCollection = None,
+        vp_stage: Optional[int] = None,
+        num_layers: int = None,
     ) -> None:
         super().__init__(config=config)
         self.residual_in_fp32 = residual_in_fp32
         self.pre_process = pre_process
         self.post_layer_norm = post_layer_norm
         self.post_process = post_process
+        self.vp_stage = vp_stage
 
         assert pg_collection is not None, "pg_collection must be provided for MambaStack"
 
@@ -103,7 +106,7 @@ def __init__(
         self.hybrid_override_pattern = hybrid_override_pattern
 
         self.layer_type_list = allocate_layers(
-            self.config.num_layers,
+            self.config.num_layers if num_layers is None else num_layers,
             self.hybrid_attention_ratio,
             self.hybrid_mlp_ratio,
             self.hybrid_override_pattern,
@@ -147,7 +150,7 @@ def __init__(
                 elif layer_type == LayerSymbols.MOE:
                     # Transformer layers apply their own pp_layer_offset
                     layer = build_module(
-                        submodules.moe_layer, config=self.config, layer_number=i + 1
+                        submodules.moe_layer, config=self.config, layer_number=i + 1,
                     )
                 else:
                     assert False, "unexpected layer_type"
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 8135e1c47a3..f86aa9dd401 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -395,19 +395,19 @@ class MultiTokenPredictionLayerSubmodules:
             embedding normalization to be applied.
         eh_proj (Union[ModuleSpec, type]): Specification or instance of the
             linear projection to be applied.
-        transformer_layer (Union[ModuleSpec, type]): Specification
-            or instance of the transformer block to be applied.
+        mtp_model_layer (Union[ModuleSpec, type]): Specification
+            or instance of the transformer or mamba block to be applied.
     """
 
     enorm: Union[ModuleSpec, type] = None
     hnorm: Union[ModuleSpec, type] = None
     eh_proj: Union[ModuleSpec, type] = None
-    transformer_layer: Union[ModuleSpec, type] = None
+    mtp_model_layer: Union[ModuleSpec, type] = None
     layer_norm: Union[ModuleSpec, type] = None
 
 
 def get_mtp_layer_spec(
-    transformer_layer_spec: ModuleSpec, use_transformer_engine: bool
+    mtp_model_layer_spec: ModuleSpec, use_transformer_engine: bool
 ) -> ModuleSpec:
     """Get the MTP layer spec.
 
@@ -415,13 +415,13 @@ def get_mtp_layer_spec(
         ModuleSpec: Module specification with TE modules
     """
     return get_mtp_layer_spec_for_backend(
-        transformer_layer_spec,
+        mtp_model_layer_spec,
         backend=TESpecProvider() if use_transformer_engine else LocalSpecProvider(),
     )
 
 
 def get_mtp_layer_spec_for_backend(
-    transformer_layer_spec: ModuleSpec, backend: BackendSpecProvider
+   mtp_model_layer_spec: ModuleSpec, backend: BackendSpecProvider
 ) -> ModuleSpec:
     """Get the MTP layer spec.
 
@@ -436,7 +436,7 @@ def get_mtp_layer_spec_for_backend(
             enorm=layer_norm_impl,
             hnorm=layer_norm_impl,
             eh_proj=column_parallel_linear_impl,
-            transformer_layer=transformer_layer_spec,
+            mtp_model_layer=mtp_model_layer_spec,
             layer_norm=layer_norm_impl,
         ),
     )
@@ -534,6 +534,7 @@ def __init__(
         self,
         config: TransformerConfig,
         submodules: MultiTokenPredictionLayerSubmodules,
+        mtp_hybrid_override_pattern: str = None,
         layer_number: int = 1,
         vp_stage: Optional[int] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
@@ -544,14 +545,18 @@ def __init__(
         self.layer_number = layer_number
         self.vp_stage = vp_stage
         self.cp_group = pg_collection.cp
-
-        self_attention_spec = self.submodules.transformer_layer.submodules.self_attention
-        attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
-        assert attn_mask_type in SUPPORTED_ATTN_MASK, (
-            f"Multi-Token Prediction (MTP) is not jet supported with "
-            + f"{attn_mask_type} attention mask type."
-            + f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
-        )
+        self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern
+
+        if hasattr(self.submodules.mtp_model_layer.submodules, 'attention_layer'):
+            self_attention_spec = self.submodules.mtp_model_layer.submodules.attention_layer
+            if self_attention_spec.submodules.self_attention is not None:
+                    self_attention_spec = self_attention_spec.submodules.self_attention
+                    attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
+                    assert attn_mask_type in SUPPORTED_ATTN_MASK, (
+                        f"Multi-Token Prediction (MTP) is not jet supported with "
+                        + f"{attn_mask_type} attention mask type."
+                        + f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
+                    )
 
         self.enorm = build_module(
             self.submodules.enorm,
@@ -582,10 +587,29 @@ def __init__(
             bias=False,
             skip_bias_add=False,
             is_expert=False,
+            tp_comm_buffer_name="mtp_eh_proj"
         )
-        self.transformer_layer = build_module(
-            self.submodules.transformer_layer, config=self.config, vp_stage=vp_stage
-        )
+        if self.config.mtp_num_layers is not None:
+            if self.mtp_hybrid_override_pattern is not None:
+                pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+                # We do not need pre and post process stage for MTP layer, given they are handled in the 
+                # MultiTokenPredictionLayer itself.
+                self.mtp_model_layer = build_module(
+                    self.submodules.mtp_model_layer,
+                    self.config,
+                    pre_process=False, 
+                    post_process=False,
+                    hybrid_override_pattern=self.mtp_hybrid_override_pattern,
+                    dtype=self.config.params_dtype,
+                    pg_collection=pg_collection,
+                    vp_stage=self.vp_stage,
+                    num_layers=self.config.mtp_num_layers
+                )
+            else:
+                # Uses the transformer block spec for MTP layer. This option is only implemented for the 
+                # GPT model. In hybrid model, we model transformer block spec for MTP layers with the hybrid
+                # override pattern.
+                self.mtp_model_layer = build_module(self.submodules.mtp_model_layer, config=self.config, vp_stage=self.vp_stage)
 
         self.final_layernorm = build_module(
             self.submodules.layer_norm,
@@ -696,7 +720,6 @@ def _proj_and_transformer_layer(
             transformer_layer_fp8_context = nullcontext()
 
         # TODO: currently ignoring FP4 in MTP layers because we need more numerical validation
-
         with rng_context:
             with fp8_context:
                 hidden_states = self._concat_embeddings(hidden_states, decoder_input)
@@ -705,19 +728,30 @@ def _proj_and_transformer_layer(
             # transformer layer is cudagraphed, the FP8GlobalStateManager.is_first_fp8_module() is
             # True so that the fp8 weight caching can be triggered correctly.
             with transformer_layer_fp8_context:
-                hidden_states, _ = self.transformer_layer(
-                    hidden_states=hidden_states,
-                    attention_mask=attention_mask,
-                    context=context,
-                    context_mask=context_mask,
-                    rotary_pos_emb=rotary_pos_emb,
-                    rotary_pos_cos=rotary_pos_cos,
-                    rotary_pos_sin=rotary_pos_sin,
-                    attention_bias=attention_bias,
-                    inference_params=inference_params,
-                    packed_seq_params=packed_seq_params,
-                    sequence_len_offset=sequence_len_offset,
-                )
+                if self.mtp_hybrid_override_pattern is not None:
+                    # Since pre-process is set to False, we need to set the input tensor manually.
+                    self.mtp_model_layer.set_input_tensor(hidden_states)
+                    hidden_states = self.mtp_model_layer(
+                        hidden_states=hidden_states,
+                        attention_mask=attention_mask,
+                        rotary_pos_emb=rotary_pos_emb,
+                        inference_params=inference_params,
+                        inference_context=context,
+                    )
+                else:
+                    hidden_states, _ = self.mtp_model_layer(
+                        hidden_states=hidden_states,
+                        attention_mask=attention_mask,
+                        context=context,
+                        context_mask=context_mask,
+                        rotary_pos_emb=rotary_pos_emb,
+                        rotary_pos_cos=rotary_pos_cos,
+                        rotary_pos_sin=rotary_pos_sin,
+                        attention_bias=attention_bias,
+                        inference_params=inference_params,
+                        packed_seq_params=packed_seq_params,
+                        sequence_len_offset=sequence_len_offset,
+                    )
 
         hidden_states = self._postprocess(hidden_states)
 
@@ -948,9 +982,11 @@ def __init__(
         config: TransformerConfig,
         spec: Union[TransformerBlockSubmodules, ModuleSpec],
         vp_stage: Optional[int] = None,
+        mtp_hybrid_override_pattern: str = None,
         pg_collection: ProcessGroupCollection = None,
     ):
         super().__init__(config=config)
+        self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern
         self.submodules = _get_mtp_block_submodules(config, spec)
         self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor
         self.vp_stage = vp_stage
@@ -981,6 +1017,7 @@ def build_layer(layer_spec, layer_number):
                     layer_number=layer_number,
                     vp_stage=self.vp_stage,
                     pg_collection=pg_collection,
+                    mtp_hybrid_override_pattern=self.mtp_hybrid_override_pattern
                 )
             return module
 
@@ -1049,6 +1086,82 @@ def forward(
         hidden_states = torch.cat(hidden_states_list, dim=0)
         return hidden_states
 
+    def process_loss(
+        self,
+        hidden_states: Tensor,
+        labels: Tensor,
+        loss_mask: Optional[Tensor],
+        output_layer: Callable,
+        output_weight: Optional[Tensor],
+        runtime_gather_output: Optional[bool],
+        is_training: bool,
+        compute_language_model_loss: Callable,
+        packed_seq_params: Optional[PackedSeqParams] = None,
+    ) -> Tensor:
+        """Process Multi-Token Prediction (MTP) loss computation.
+        
+        This method handles the MTP loss computation for multiple prediction layers.
+        It chunks the hidden states, computes logits and losses for each MTP layer,
+        and applies loss scaling.
+        
+        Args:
+            hidden_states (Tensor): Hidden states tensor from the model.
+            labels (Tensor): Ground truth labels.
+            loss_mask (Optional[Tensor]): Mask for loss computation. If None, uses all ones.
+            output_layer (Callable): Output layer method to compute logits.
+            output_weight (Optional[Tensor]): Optional output weight for shared embeddings.
+            runtime_gather_output (Optional[bool]): Whether to gather output at runtime.
+            is_training (bool): Whether the model is in training mode.
+            compute_language_model_loss (Callable): Method to compute language model loss.
+        
+        Returns:
+            Tensor: Updated hidden states after MTP loss processing.
+        """
+        mtp_labels = labels.clone()
+        hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0)
+        hidden_states = hidden_states_list[0]
+        
+        if loss_mask is None:
+            # if loss_mask is not provided, use all ones as loss_mask
+            loss_mask = torch.ones_like(mtp_labels)
+        
+        for mtp_layer_number in range(self.config.mtp_num_layers):
+            # output
+            mtp_logits, _ = output_layer(
+                hidden_states_list[mtp_layer_number + 1],
+                weight=output_weight,
+                runtime_gather_output=runtime_gather_output,
+            )
+            # Calc loss for the current Multi-Token Prediction (MTP) layers.
+            mtp_labels, _ = roll_tensor(mtp_labels, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params)
+            loss_mask, num_tokens = roll_tensor(
+                loss_mask, shifts=-1, dims=-1, cp_group=self.cp_group,  packed_seq_params=packed_seq_params
+            )
+            mtp_loss = compute_language_model_loss(mtp_labels, mtp_logits)
+            mtp_loss = loss_mask * mtp_loss
+            if is_training:
+                # TODO(shifangx): remove the use of parallel_state here
+                # after moving loss logging to loss_func in pretrain_gpt.py
+                MTPLossLoggingHelper.save_loss_to_tracker(
+                    torch.sum(mtp_loss) / num_tokens,
+                    mtp_layer_number,
+                    self.config.mtp_num_layers,
+                    avg_group=parallel_state.get_data_parallel_group(
+                        with_context_parallel=True
+                    ),
+                )
+            mtp_loss_scale = self.config.mtp_loss_scaling_factor / self.config.mtp_num_layers
+            if self.config.calculate_per_token_loss:
+                hidden_states = MTPLossAutoScaler.apply(
+                    hidden_states, mtp_loss_scale * mtp_loss
+                )
+            else:
+                hidden_states = MTPLossAutoScaler.apply(
+                    hidden_states, mtp_loss_scale * mtp_loss / num_tokens
+                )
+        
+        return hidden_states
+
     def sharded_state_dict(
         self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None
     ) -> ShardedStateDict:
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index f40f5ce0055..41781992309 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -529,6 +529,9 @@ def validate_args(args, defaults={}):
                 args.global_batch_size), flush=True)
     assert args.global_batch_size > 0
 
+    if args.mtp_hybrid_override_pattern is not None:
+        assert args.mtp_spec is not None, "--mtp-spec must be set for mtp hybrid MTP models"
+    
     # Uneven virtual pipeline parallelism
     assert (
         int(args.num_layers_per_virtual_pipeline_stage is not None)
@@ -600,6 +603,18 @@ def validate_args(args, defaults={}):
                 assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \
                     'Number of layers should be divisible by the pipeline-model-parallel size'
 
+    # Validate MTP hybrid pattern segment count
+    if args.mtp_hybrid_override_pattern is not None:
+        mtp_segments = len(args.mtp_hybrid_override_pattern.split('|'))
+        expected_segments = args.pipeline_model_parallel_size * (
+            args.virtual_pipeline_model_parallel_size or 1
+        )
+        assert mtp_segments == expected_segments, (
+            f'MTP hybrid pattern has {mtp_segments} segments but expected {expected_segments} '
+            f'(PP={args.pipeline_model_parallel_size} * VP={args.virtual_pipeline_model_parallel_size or 1}). '
+            f'Use: {"|||" * (expected_segments - 1)}*E'
+        )
+    
     if args.virtual_pipeline_model_parallel_size is not None:
         if args.overlap_p2p_comm:
             assert args.pipeline_model_parallel_size > 1, \
@@ -3369,6 +3384,19 @@ def _add_experimental_args(parser):
                             'hybrid ratio arguments, then the number of each type'
                             'of layer in the override pattern must match number in'
                             'the overidden pattern')
+    group.add_argument("--mtp-hybrid-override-pattern", type=str, default=None,
+                       help='Force a specific hybrid layer pattern for MTP layers. '
+                       'If a value greater than 0.0 is supplied to any of the hybrid ratio '
+                       'arguments, then the number of each type of layer in the '
+                       'override pattern must match number in the overidden '
+                       'pattern')
+    group.add_argument('--mtp-spec', type=str, default=None, nargs='*',
+                       help='Specify the <module_location function_name> pair '
+                       'that returns a spec for mtp layer to customize a model,' 
+                       'transformer block, or transformer layer, depending on '
+                       'the use case. To use local spec specify local as the argument.'
+                       'For more details, see the model class, '
+                       'TransformerBlock class, or TransformerLayer class.')
     group.add_argument('--mamba-state-dim', type=int, default=128,
                        help='State dimension for Mamba layers.')
     group.add_argument('--mamba-head-dim', type=int, default=64,
diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index 325158780a4..779f27fd581 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -1347,6 +1347,7 @@ def _set_arg(arg_name, old_arg_name=None, force=False):
     _set_arg('hidden_dropout', force=True)
 
     _set_arg('hybrid_override_pattern', force=True)
+    _set_arg('mtp_hybrid_override_pattern', force=True)
     _set_arg('spec', force=True)
     _set_arg('hybrid_attention_ratio', force=True)
     _set_arg('hybrid_mlp_ratio', force=True)
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 3f23ca59f4a..28e3d2b6e94 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -178,6 +178,11 @@ def calculate_layer_counts():
             for layer_type in args.hybrid_override_pattern:
                 if layer_type in counts:
                     counts[layer_type] += 1
+            if args.mtp_hybrid_override_pattern:
+                assert args.mtp_num_layers is not None
+                for layer_type in args.mtp_hybrid_override_pattern:
+                    if layer_type in counts:
+                        counts[layer_type] += args.mtp_num_layers
             return counts['*'], counts['M'], counts['-'], counts['E']
         else:
             num_attn_layers = round(args.num_layers * args.hybrid_attention_ratio)
@@ -254,8 +259,10 @@ def hybrid_flops(batch_size, seq_len, hidden_size,
                      mlp_expansion=4.0, swiglu=False,
                      moe_latent_size=None,
                      moe_ffn_hidden_size=2048, shared_expert_ffn_hidden_size=2048, num_experts_routed_to=1,
-                     vocab_size=256000):
+                     vocab_size=256000, mtp_num_layers=None):
         """Calculate total FLOPs for the hybrid model."""
+        if mtp_num_layers is None:
+            mtp_num_layers = 0
         flops_fwd = (
                 num_attn_layers * attn_layer_flops(batch_size, seq_len, hidden_size,
                                                    num_attn_heads, gqa, gqa_groups, kv_channels) +
@@ -267,7 +274,7 @@ def hybrid_flops(batch_size, seq_len, hidden_size,
                 num_moe_layers * moe_layer_flops(batch_size, seq_len, hidden_size, moe_ffn_hidden_size,
                                                  shared_expert_ffn_hidden_size, num_experts_routed_to,
                                                  moe_latent_size, swiglu) +
-                (2 * batch_size * seq_len * hidden_size * vocab_size)  # logits computation
+                (2 * batch_size * seq_len * hidden_size * vocab_size * (1 + mtp_num_layers))  # logits computation
         )
         return flops_fwd * 3
 
@@ -478,6 +485,7 @@ def transformer_flops():
                                            else args.moe_shared_expert_intermediate_size),
             num_experts_routed_to=args.moe_router_topk,
             vocab_size=args.padded_vocab_size,
+            mtp_num_layers=args.mtp_num_layers,
         )
     else:
         # Compute standard Transformer model FLOPs.
diff --git a/pretrain_mamba.py b/pretrain_mamba.py
index 45b646a6cc0..1eaade5faa3 100644
--- a/pretrain_mamba.py
+++ b/pretrain_mamba.py
@@ -131,7 +131,7 @@ def forward_step(data_iterator, model: MambaModel):
 
     with stimer:
         output_tensor = model(tokens, position_ids, attention_mask,
-                              labels=labels)
+                              labels=labels, loss_mask=loss_mask)
 
     # [ModelOpt]: model is needed to access ModelOpt distillation losses
     return output_tensor, partial(loss_func, loss_mask, model=model)
diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index 05fb2c4fe63..09c08a64f22 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -94,7 +94,7 @@ def test_constructor_local(self, tp):
             assert mtp.layers[i].hnorm.weight.shape[0] == config.hidden_size
             assert mtp.layers[i].eh_proj.weight.shape[0] == config.hidden_size / tp
             assert mtp.layers[i].eh_proj.weight.shape[1] == config.hidden_size * 2
-            assert mtp.layers[i].transformer_layer is not None
+            assert mtp.layers[i].mtp_model_layer is not None
         num_weights = sum([p.numel() for p in mtp.parameters()])
         if tp == 1:
             assert num_weights == 58560 * config.mtp_num_layers
@@ -120,7 +120,7 @@ def test_constructor_ues_te(self, tp, cp):
             assert mtp.layers[i].hnorm.weight.shape[0] == config.hidden_size
             assert mtp.layers[i].eh_proj.weight.shape[0] == config.hidden_size / tp
             assert mtp.layers[i].eh_proj.weight.shape[1] == config.hidden_size * 2
-            assert mtp.layers[i].transformer_layer is not None
+            assert mtp.layers[i].mtp_model_layer is not None
         num_weights = sum([p.numel() for p in mtp.parameters()])
         if tp == 1:
             assert num_weights == 58560 * config.mtp_num_layers

From 3db97419daeda29c28c20144d277ee717e12e477 Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Mon, 1 Dec 2025 17:56:50 +0000
Subject: [PATCH 02/48] addressed comments

---
 .../common/language_module/language_module.py |   2 +-
 .../core/models/mamba/mamba_layer_specs.py    |   7 +-
 megatron/core/ssm/mamba_block.py              |   2 +-
 .../transformer/multi_token_prediction.py     |  11 +-
 .../test_multi_token_prediction.py            | 270 ++++++++++++++++++
 5 files changed, 286 insertions(+), 6 deletions(-)

diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
index 706f8429a29..6f145f7f32c 100644
--- a/megatron/core/models/common/language_module/language_module.py
+++ b/megatron/core/models/common/language_module/language_module.py
@@ -254,7 +254,7 @@ def setup_embeddings_and_output_layer(self) -> None:
             LanguageModule.embedding_warning_printed = True
 
     def shared_embedding_or_output_weight(self) -> Tensor:
-        """Gets the emedding weight or output logit weights when share embedding and output weights set to True
+        """Gets the embedding weight or output logit weights when share embedding and output weights set to True
           or when use Multi-Token Prediction (MTP).
 
         Returns:
diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index 286d72b36ff..0a6e39052c9 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -22,7 +22,6 @@
 from megatron.core.transformer.mlp import MLP, MLPSubmodules
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
-from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionBlockSubmodules
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.transformer.multi_token_prediction import (
     MultiTokenPredictionBlockSubmodules,
@@ -40,7 +39,6 @@
 )
 
 try:
-    from megatron.core.extensions.transformer_engine import TEFusedMLP, TENorm
     from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
 
     HAVE_TE = True
@@ -182,7 +180,10 @@ def get_mamba_mtp_block_spec(
     vp_stage: Optional[int] = None,
 ) -> MultiTokenPredictionBlockSubmodules:
     """Mamba Multi-Token Prediction (MTP) block spec."""
+    assert HAVE_KITCHEN, "Kitchen needs to be installed."
+    
     if use_transformer_engine:
+        assert HAVE_TE, "TransformerEngine needs to be installed."
         backend: BackendSpecProvider = (
             KitchenSpecProvider(fallback=TESpecProvider())
             if config.use_kitchen
@@ -221,7 +222,7 @@ def get_mamba_mtp_block_spec_for_backend(
     if len(mtp_layer_specs) > 0:
         assert (
             len(mtp_layer_specs) == config.mtp_num_layers
-        ), +f"currently all of the mtp layers must stage in the same pipeline stage."
+        ), f"currently all of the mtp layers must stage in the same pipeline stage."
         mtp_block_spec = MultiTokenPredictionBlockSubmodules(layer_specs=mtp_layer_specs)
     else:
         mtp_block_spec = None
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index 4e8bc27cc54..ed9ce49fae7 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -150,7 +150,7 @@ def __init__(
                 elif layer_type == LayerSymbols.MOE:
                     # Transformer layers apply their own pp_layer_offset
                     layer = build_module(
-                        submodules.moe_layer, config=self.config, layer_number=i + 1,
+                        submodules.moe_layer, config=self.config, layer_number=i+1,
                     )
                 else:
                     assert False, "unexpected layer_type"
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index f86aa9dd401..2373b694528 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -553,10 +553,19 @@ def __init__(
                     self_attention_spec = self_attention_spec.submodules.self_attention
                     attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
                     assert attn_mask_type in SUPPORTED_ATTN_MASK, (
-                        f"Multi-Token Prediction (MTP) is not jet supported with "
+                        f"Multi-Token Prediction (MTP) is not yet supported with "
                         + f"{attn_mask_type} attention mask type."
                         + f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
                     )
+        elif hasattr(self.submodules.mtp_model_layer.submodules, 'self_attention'):
+            self_attention_spec = self.submodules.mtp_model_layer.submodules.self_attention
+            if self_attention_spec is not None:
+                attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
+                assert attn_mask_type in SUPPORTED_ATTN_MASK, (
+                    f"Multi-Token Prediction (MTP) is not yet supported with "
+                    + f"{attn_mask_type} attention mask type."
+                    + f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
+                )
 
         self.enorm = build_module(
             self.submodules.enorm,
diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index 09c08a64f22..e818cde16c4 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -13,6 +13,11 @@
     get_gpt_mtp_block_spec,
 )
 from megatron.core.models.gpt.gpt_model import GPTModel
+from megatron.core.models.mamba.mamba_layer_specs import (
+    get_mamba_mtp_block_spec,
+    mamba_stack_spec,
+)
+from megatron.core.models.mamba.mamba_model import MambaModel
 from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.parallel_state import get_context_parallel_group
@@ -684,3 +689,268 @@ def log(self, metrics, iteration):
         assert torch.all(MTPLossLoggingHelper.tracker["values"] == 0)
         assert MTPLossLoggingHelper.tracker["reduce_group"] is None
         assert MTPLossLoggingHelper.tracker["avg_group"] is None
+
+
+class TestMultiTokenPredictionMamba:
+    """Test Multi-Token Prediction with Mamba hybrid models."""
+
+    def setup_method(self, method):
+        self.seq_length = 32
+        self.micro_batch_size = 2
+        os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1'
+
+    def teardown_method(self, method):
+        Utils.destroy_model_parallel()
+        destroy_global_vars()
+        destroy_num_microbatches_calculator()
+        MTPLossLoggingHelper.tracker = {}
+
+    def model_provider(
+        self,
+        pre_process=True,
+        post_process=True,
+        **config_kwargs,
+    ):
+        """Model provider for Mamba hybrid models with MTP."""
+        model_parallel_cuda_manual_seed(_SEED)
+        args = get_args()
+        config = core_transformer_config_from_args(args)
+        
+        mtp_block_spec = get_mamba_mtp_block_spec(
+            config=config, spec=mamba_stack_spec, use_transformer_engine=True
+        )
+        
+        model = MambaModel(
+            config=config,
+            mamba_stack_spec=mamba_stack_spec,
+            vocab_size=args.vocal_size,
+            max_sequence_length=args.max_position_embeddings,
+            pre_process=pre_process,
+            post_process=post_process,
+            hybrid_attention_ratio=args.hybrid_attention_ratio,
+            hybrid_mlp_ratio=args.hybrid_mlp_ratio,
+            hybrid_override_pattern=args.hybrid_override_pattern,
+            mtp_hybrid_override_pattern=args.mtp_hybrid_override_pattern,
+            fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
+            parallel_output=True,
+            share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
+            position_embedding_type=args.position_embedding_type,
+            rotary_percent=args.rotary_percent,
+            mtp_block_spec=mtp_block_spec,
+        )
+        return model
+
+    def create_test_args(
+        self, tp, cp, sequence_length, micro_batch_size, fp8=None, full_recompute=False
+    ):
+        destroy_global_vars()
+        destroy_num_microbatches_calculator()
+
+        sys.argv = ['test_multi_token_prediction_mamba.py']
+        args = parse_args()
+        args.num_layers = 4
+        args.mtp_num_layers = 2
+        args.mtp_loss_scaling_factor = 0.1
+        args.vocal_size = 128800
+        args.hidden_size = 128
+        args.num_attention_heads = 8
+        args.num_query_groups = 8
+        args.mamba_num_groups = 4  
+        args.max_position_embeddings = 256
+        args.micro_batch_size = micro_batch_size
+        args.create_attention_mask_in_dataloader = True
+        args.seq_length = sequence_length
+        args.tensor_model_parallel_size = tp
+        args.sequence_parallel = True if tp > 1 else False
+        args.context_parallel_size = cp
+        args.position_embedding_type = 'rope'
+        args.train_iters = 1
+        args.ckpt_format = 'torch_dist'
+        args.lr = 3e-5
+        args.attention_dropout = 0.0
+        args.hidden_dropout = 0.0
+        args.async_tensor_model_parallel_allreduce = False
+        args.no_save_optim = True
+        args.no_load_optim = True
+        args.no_load_rng = True
+        args.bf16 = True
+        
+        args.hybrid_attention_ratio = 0.5  
+        args.hybrid_mlp_ratio = 0.0  
+        args.hybrid_override_pattern = "M*M*"  
+        args.mtp_hybrid_override_pattern = "M*"  
+        args.spec = "megatron.core.models.mamba.mamba_layer_specs.mamba_stack_spec"
+        args.mtp_spec = "megatron.core.models.mamba.mamba_layer_specs.mamba_stack_spec"
+        
+        if fp8 is not None:
+            args.fp8 = 'e4m3'
+        if full_recompute:
+            args.recompute_granularity = 'full'
+            args.recompute_method = 'uniform'
+            args.recompute_num_layers = 1
+        else:
+            args.recompute_granularity = None
+        args.add_bias_linear = False
+        args.swiglu = True
+
+        validate_args(args)
+        set_global_variables(args, False)
+        return args
+
+    def get_batch(self, seq_length, micro_batch_size):
+        data = list(range(seq_length))
+        input_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        labels = 1 + torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        position_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        attention_mask = torch.ones(
+            (micro_batch_size, 1, seq_length, seq_length), dtype=bool
+        ).cuda()
+        loss_mask = torch.ones(seq_length).repeat((micro_batch_size, 1)).cuda()
+        batch = {
+            'tokens': input_ids,
+            'labels': labels,
+            'loss_mask': loss_mask,
+            'attention_mask': attention_mask,
+            'position_ids': position_ids,
+        }
+        return batch
+
+    @pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available")
+    @pytest.mark.parametrize(("tp", "cp"), [(1, 1), (2, 1)])
+    def test_sharded_state_dict_mamba(self, tp, cp):
+        """Test MTP with Mamba hybrid model - sharded state dict."""
+        args = self.create_test_args(tp, cp, self.seq_length, self.micro_batch_size)
+        set_args(args)
+        torch.manual_seed(_SEED)
+        Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+        mamba_model = get_model(self.model_provider, ModelType.encoder_or_decoder)
+        mamba_model = unwrap_model(mamba_model)
+        sharded_state_dict = mamba_model[0].sharded_state_dict()
+        
+        # Verify MTP layers are in the state dict
+        for i in range(args.mtp_num_layers):
+            assert f"mtp.layers.{i}.enorm.weight" in sharded_state_dict.keys()
+            assert f"mtp.layers.{i}.hnorm.weight" in sharded_state_dict.keys()
+            assert f"mtp.layers.{i}.eh_proj.weight" in sharded_state_dict.keys()
+
+    @pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available")
+    @pytest.mark.parametrize(("tp", "cp"), [(1, 1), (2, 1)])
+    def test_forward_backward_mamba(self, tmp_path_dist_ckpt, tp, cp):
+        """Test MTP forward and backward with Mamba hybrid model."""
+        tp_ref = 1
+        cp_ref = 1
+        args = self.create_test_args(tp_ref, cp_ref, self.seq_length, self.micro_batch_size)
+        set_args(args)
+        torch.manual_seed(_SEED)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp_ref, context_parallel_size=cp_ref
+        )
+        batch = self.get_batch(self.seq_length, self.micro_batch_size)
+        tokens, labels, loss_mask, attention_mask, position_ids = batch.values()
+        
+        mamba_model_ref, optimizer, opt_param_scheduler = setup_model_and_optimizer(
+            self.model_provider, ModelType.encoder_or_decoder
+        )
+        
+        output_ref = mamba_model_ref[0].forward(
+            input_ids=tokens,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            loss_mask=loss_mask,
+        )
+        
+        tracker = MTPLossLoggingHelper.tracker
+        mtp_loss_ref = None
+        assert "values" in tracker
+        mtp_loss_ref = tracker['values'].clone()
+        MTPLossLoggingHelper.clean_loss_in_tracker()
+
+        iteration = 123
+        num_floating_point_operations_so_far = 456
+
+        def set_ckpt_path(ckpt_path):
+            args.save = ckpt_path
+            args.load = ckpt_path
+
+        with TempNamedDir(
+            tmp_path_dist_ckpt / 'test_mtp_mamba_model_reconfiguration'
+        ) as ckpt_dir:
+            set_ckpt_path(ckpt_dir)
+            save_checkpoint(
+                iteration,
+                mamba_model_ref,
+                optimizer,
+                opt_param_scheduler,
+                num_floating_point_operations_so_far,
+            )
+
+            expected_ckpt_path = args.save / "iter_0000123" / ".metadata"
+            assert os.path.exists(expected_ckpt_path)
+
+            Utils.destroy_model_parallel()
+            args = self.create_test_args(tp, cp, self.seq_length, self.micro_batch_size)
+            set_args(args)
+            set_ckpt_path(ckpt_dir)
+            torch.manual_seed(_SEED)
+            Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+            mamba_model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
+                self.model_provider, ModelType.encoder_or_decoder
+            )
+            load_checkpoint(mamba_model, optimizer, opt_param_scheduler, strict=False)
+            
+            batch["output_ref"] = output_ref
+            batch = get_batch_on_this_cp_rank(batch)
+            tokens, labels, loss_mask, attention_mask, position_ids, output_ref = batch.values()
+            
+            output = mamba_model[0].forward(
+                input_ids=tokens,
+                position_ids=position_ids,
+                attention_mask=attention_mask,
+                labels=labels,
+                loss_mask=loss_mask,
+            )
+            
+            tracker = MTPLossLoggingHelper.tracker
+            assert "values" in tracker
+            mtp_loss = tracker['values'].clone()
+            pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['cp'])
+            torch.distributed.all_reduce(
+                mtp_loss, group=pg_collection.cp, op=torch.distributed.ReduceOp.AVG
+            )
+            MTPLossLoggingHelper.clean_loss_in_tracker()
+            
+            assert torch.allclose(output_ref, output, rtol=1e-03, atol=1e-03)
+            assert torch.allclose(mtp_loss, mtp_loss_ref, rtol=1e-02, atol=1e-02)
+
+            assert output.shape[0] == self.micro_batch_size
+            assert output.shape[1] == self.seq_length / cp
+
+            loss = output.mean()
+            loss.backward()
+            for name, param in mamba_model[0].named_parameters():
+                assert param.main_grad is not None
+
+    @pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available")
+    def test_attention_mask_validation_mamba(self):
+        """Test that attention mask type validation works for Mamba hybrid models.
+        """
+        tp = 1
+        cp = 1
+        args = self.create_test_args(tp, cp, self.seq_length, self.micro_batch_size)
+        set_args(args)
+        torch.manual_seed(_SEED)
+        Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+        
+        try:
+            mamba_model = get_model(self.model_provider, ModelType.encoder_or_decoder)
+            mamba_model = unwrap_model(mamba_model)
+            assert isinstance(mamba_model[0], MambaModel)
+            assert mamba_model[0].mtp is not None
+        except AssertionError as e:
+            if "Multi-Token Prediction (MTP) is not yet supported" in str(e):
+                pytest.fail(
+                    f"Attention mask validation failed for Mamba hybrid model: {e}"
+                )
+            else:
+                raise

From aeecff9d02441dfa5308d8f25ce61b09e4732e01 Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Wed, 3 Dec 2025 12:54:46 +0000
Subject: [PATCH 03/48] fixed the duplicate

---
 megatron/core/models/mamba/mamba_model.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 6521100648b..023d337e3f2 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -399,11 +399,7 @@ def forward(
                 compute_language_model_loss=self.compute_language_model_loss,
             )
 
-        # logits and loss
-        output_weight = None
-        if self.share_embeddings_and_output_weights:
-            output_weight = self.shared_embedding_or_output_weight()
-
+        
         sequence_parallel_override = False
         if in_inference_mode and inference_context.materialize_only_last_token_logits:
             if inference_context.is_static_batching():

From 9b4e2b5578863628e81e1a20b481f3dfca1cba12 Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Fri, 5 Dec 2025 11:56:07 +0000
Subject: [PATCH 04/48] added repeated mtp layers

---
 .../core/models/mamba/mamba_layer_specs.py    | 11 +++++--
 megatron/core/ssm/mamba_block.py              | 12 +++++++
 .../transformer/multi_token_prediction.py     | 31 ++++++++++++++-----
 .../core/transformer/transformer_config.py    |  3 ++
 megatron/training/arguments.py                |  4 +++
 5 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index 0a6e39052c9..ced9d66ae83 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -214,15 +214,20 @@ def get_mamba_mtp_block_spec_for_backend(
         mtp_model_layer_spec=spec, backend=backend
     )
     mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0
-    mtp_layer_specs = [mtp_layer_spec] * mtp_num_layers
+    
+    if config.mtp_use_repeated_layer:
+        mtp_layer_specs = [mtp_layer_spec]
+    else:
+        mtp_layer_specs = [mtp_layer_spec] * mtp_num_layers
 
     offset = get_mtp_layer_offset(config)
     # split the mtp layer specs to only include the layers that are built in this pipeline stage.
     mtp_layer_specs = mtp_layer_specs[offset : offset + num_layers_to_build]
     if len(mtp_layer_specs) > 0:
+        expected_num_layers = 1 if config.mtp_use_repeated_layer else config.mtp_num_layers
         assert (
-            len(mtp_layer_specs) == config.mtp_num_layers
-        ), f"currently all of the mtp layers must stage in the same pipeline stage."
+            len(mtp_layer_specs) == expected_num_layers
+        ), f"currently all of the mtp layers must stage in the same pipeline stage. Expected {expected_num_layers}, got {len(mtp_layer_specs)}"
         mtp_block_spec = MultiTokenPredictionBlockSubmodules(layer_specs=mtp_layer_specs)
     else:
         mtp_block_spec = None
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index ed9ce49fae7..6752ed7da09 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -27,6 +27,9 @@
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.transformer.transformer_layer import TransformerLayer
 from megatron.core.transformer.utils import sharded_state_dict_default
+
+from megatron.core.ssm.mlp_layer import MLPLayer
+
 from megatron.core.utils import WrappedTensor, deprecate_inference_params, make_viewless_tensor
 
 
@@ -290,6 +293,15 @@ def forward(
                             rotary_pos_emb=rotary_pos_emb,
                             sequence_len_offset=sequence_len_offset,
                         )
+                    elif isinstance(layer, MLPLayer):
+                        # MLPLayer (standalone MLP without attention)
+                        hidden_states = layer(
+                            hidden_states=hidden_states,
+                            attention_mask=attention_mask,
+                            inference_context=inference_context,
+                            rotary_pos_emb=rotary_pos_emb,
+                            sequence_len_offset=sequence_len_offset,
+                        )
                     else:  # MambaLayer
                         hidden_states = layer(
                             hidden_states=hidden_states,
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 2373b694528..d0a32acfb60 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -982,6 +982,9 @@ class MultiTokenPredictionBlock(MegatronModule):
     the linear projection. The combined serves as the input of the Transformer block at
     the k-th depth to produce the output representation.
 
+    When `mtp_use_repeated_layer=True` in config, instead of creating N separate MTP layers,
+    only 1 layer is created and applied mtp_num_layers times. 
+
     for more information, please refer to DeepSeek-V3 Technical Report
     https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf
     """
@@ -999,6 +1002,8 @@ def __init__(
         self.submodules = _get_mtp_block_submodules(config, spec)
         self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor
         self.vp_stage = vp_stage
+        
+        self.mtp_use_repeated_layer = self.config.mtp_use_repeated_layer
 
         # Initialize Context Parallelism (CP) support for MTP
         # This enables MTP to work with CP > 1 by providing the CP process group
@@ -1030,12 +1035,21 @@ def build_layer(layer_spec, layer_number):
                 )
             return module
 
-        self.layers = torch.nn.ModuleList(
-            [
-                build_layer(layer_spec, i + 1)
-                for i, layer_spec in enumerate(self.submodules.layer_specs)
-            ]
-        )
+        if self.mtp_use_repeated_layer:
+            assert len(self.submodules.layer_specs) == 1, (
+                f"Repeated MTP mode requires exactly 1 layer spec, got {len(self.submodules.layer_specs)}. "
+                f"The layer will be applied {self.config.mtp_num_layers} times."
+            )
+            self.layers = torch.nn.ModuleList([
+                build_layer(self.submodules.layer_specs[0], layer_number=1)
+            ])
+        else:
+            self.layers = torch.nn.ModuleList(
+                [
+                    build_layer(layer_spec, i + 1)
+                    for i, layer_spec in enumerate(self.submodules.layer_specs)
+                ]
+            )
 
     def forward(
         self,
@@ -1071,8 +1085,9 @@ def forward(
         offset = get_mtp_layer_offset(self.config)
         hidden_states_list = list(torch.chunk(hidden_states, 1 + offset, dim=0))
         hidden_states = hidden_states_list[offset]
-        for layer_number in range(len(self.layers)):
-            (hidden_states, input_ids, position_ids) = self.layers[layer_number](
+        for iteration in range(self.config.mtp_num_layers):
+            layer_idx = 0 if self.mtp_use_repeated_layer else iteration
+            (hidden_states, input_ids, position_ids) = self.layers[layer_idx](
                 input_ids=input_ids,
                 position_ids=position_ids,
                 hidden_states=hidden_states,
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 17063217058..71885a8a382 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -51,6 +51,9 @@ class TransformerConfig(ModelParallelConfig):
     mtp_loss_scaling_factor: Optional[float] = None
     """Weighting factor of Multi-Token Prediction (MTP) loss."""
 
+    mtp_use_repeated_layer: bool = False
+    """Use a single MTP layer repeatedly instead of multiple separate layers."""
+
     num_layers_in_first_pipeline_stage: Optional[int] = None
     """Number of transformer layers on first pipeline stage.
     None implies equal layer division across PP ranks."""
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 41781992309..abddad1be4f 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1799,6 +1799,10 @@ def _add_network_size_args(parser):
                        'which serves as an additional training objective.')
     group.add_argument('--moe-latent-size', type=int, default=None,
                        help='Latent projection dimension for MoE. If None, MoE latent projections are not used.')
+    group.add_argument('--mtp-use-repeated-layer', action='store_true',
+                       help='Use a single MTP layer repeatedly instead of multiple separate layers. '
+                       'This is more parameter-efficient. When enabled, only 1 MTP layer is created '
+                       'and applied --mtp-num-layers times.')
     return parser
 
 

From 225c2dccfd01370085479d66b1c681038420a87a Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Fri, 5 Dec 2025 13:37:24 +0000
Subject: [PATCH 05/48] adds mtp-num-per-layer to fix a bug

---
 megatron/core/transformer/multi_token_prediction.py | 5 ++++-
 megatron/core/transformer/transformer_config.py     | 3 +++
 megatron/training/arguments.py                      | 2 ++
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index d0a32acfb60..0f37c62299f 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -601,8 +601,11 @@ def __init__(
         if self.config.mtp_num_layers is not None:
             if self.mtp_hybrid_override_pattern is not None:
                 pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+
                 # We do not need pre and post process stage for MTP layer, given they are handled in the 
                 # MultiTokenPredictionLayer itself.
+                assert self.config.mtp_num_layers_per_layer is not None, \
+                    "mtp_num_layers_per_layer must be set when using mtp_hybrid_override_pattern"
                 self.mtp_model_layer = build_module(
                     self.submodules.mtp_model_layer,
                     self.config,
@@ -612,7 +615,7 @@ def __init__(
                     dtype=self.config.params_dtype,
                     pg_collection=pg_collection,
                     vp_stage=self.vp_stage,
-                    num_layers=self.config.mtp_num_layers
+                    num_layers=self.config.mtp_num_layers_per_layer
                 )
             else:
                 # Uses the transformer block spec for MTP layer. This option is only implemented for the 
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 71885a8a382..7c0450bddfe 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -54,6 +54,9 @@ class TransformerConfig(ModelParallelConfig):
     mtp_use_repeated_layer: bool = False
     """Use a single MTP layer repeatedly instead of multiple separate layers."""
 
+    mtp_num_layers_per_layer: Optional[int] = None
+    """Number of layers inside each MTP layer (e.g., for hybrid models)."""
+
     num_layers_in_first_pipeline_stage: Optional[int] = None
     """Number of transformer layers on first pipeline stage.
     None implies equal layer division across PP ranks."""
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index abddad1be4f..e4359e78184 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1803,6 +1803,8 @@ def _add_network_size_args(parser):
                        help='Use a single MTP layer repeatedly instead of multiple separate layers. '
                        'This is more parameter-efficient. When enabled, only 1 MTP layer is created '
                        'and applied --mtp-num-layers times.')
+    group.add_argument('--mtp-num-layers-per-layer', type=int, default=None,
+                       help='Number of layers inside each MTP layer (e.g., for hybrid models).')
     return parser
 
 

From 61eb031be00a0b26a70cc33aeb90684e8c465874 Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Fri, 5 Dec 2025 14:32:39 +0000
Subject: [PATCH 06/48] added the assert

---
 megatron/core/pipeline_parallel/schedules.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index 3c1fb2e3b18..c8c520b2e51 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -193,6 +193,9 @@ def set_current_microbatch(model, microbatch_id):
             layer.current_microbatch = microbatch_id
         if hasattr(model_with_decoder, 'mtp'):
             for layer in model_with_decoder.mtp.layers:
+                assert hasattr(layer, 'mtp_model_layer'), (
+                    f"MTP layer {layer} must have 'mtp_model_layer' attribute"
+                )
                 layer.mtp_model_layer.current_microbatch = microbatch_id
 
 

From e8304aaca8f12708675fecbe17f02494c9e30d14 Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Tue, 9 Dec 2025 13:36:54 +0000
Subject: [PATCH 07/48] added the assert for vp-stage of MTP layers

---
 megatron/core/transformer/multi_token_prediction.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 0f37c62299f..f1028799fe4 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -1008,6 +1008,13 @@ def __init__(
         
         self.mtp_use_repeated_layer = self.config.mtp_use_repeated_layer
 
+        vp_size = config.virtual_pipeline_model_parallel_size
+        assert is_vp_last_stage(vp_stage=vp_stage, vp_size=vp_size), (
+            f"MTP layers must be placed on the last virtual pipeline stage. "
+            f"Got vp_stage={vp_stage} with vp_size={vp_size}. "
+            f"Placing MTP layers on different VPP stages is not currently supported."
+        )
+
         # Initialize Context Parallelism (CP) support for MTP
         # This enables MTP to work with CP > 1 by providing the CP process group
         # to the roll_tensor function for proper boundary communication

From 7398d87e275e4b6be68119e7b3fa3120ad7a3d77 Mon Sep 17 00:00:00 2001
From: Rabeeh Mahabadi <rkarimimahab@nb-hel-cs-001-vscode-02.cm.cluster>
Date: Thu, 18 Dec 2025 09:57:24 -0800
Subject: [PATCH 08/48] added the LBL fixes

---
 .../core/models/mamba/mamba_layer_specs.py    |  6 +--
 megatron/core/ssm/mamba_block.py              |  9 ++++-
 megatron/core/transformer/moe/moe_layer.py    |  8 +++-
 megatron/core/transformer/moe/router.py       | 37 ++++++++++++++++---
 .../transformer/multi_token_prediction.py     |  7 ++--
 .../core/transformer/transformer_layer.py     |  5 +++
 megatron/training/arguments.py                |  7 +++-
 7 files changed, 63 insertions(+), 16 deletions(-)

diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index ced9d66ae83..431a20d2174 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -111,7 +111,7 @@
     ),
 )
 
-<<<<<<< HEAD
+
 mamba_inference_stack_spec = ModuleSpec(
     module=MambaStack,
     submodules=MambaStackSubmodules(
@@ -171,7 +171,7 @@
         ),
     ),
 )
-=======
+
 
 def get_mamba_mtp_block_spec(
     config: TransformerConfig,
@@ -233,5 +233,3 @@ def get_mamba_mtp_block_spec_for_backend(
         mtp_block_spec = None
 
     return mtp_block_spec
-
->>>>>>> a05cbed31 (added mtp support)
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index 6752ed7da09..ce7cd3a2bcf 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -88,6 +88,7 @@ def __init__(
         pg_collection: ProcessGroupCollection = None,
         vp_stage: Optional[int] = None,
         num_layers: int = None,
+        is_mtp_layer: bool = False,
     ) -> None:
         super().__init__(config=config)
         self.residual_in_fp32 = residual_in_fp32
@@ -95,6 +96,7 @@ def __init__(
         self.post_layer_norm = post_layer_norm
         self.post_process = post_process
         self.vp_stage = vp_stage
+        self.is_mtp_layer = is_mtp_layer
 
         assert pg_collection is not None, "pg_collection must be provided for MambaStack"
 
@@ -141,6 +143,7 @@ def __init__(
                         config=self.config,
                         layer_number=i + 1,
                         pg_collection=pg_collection,
+                        is_mtp_layer=is_mtp_layer,
                     )
                 elif layer_type == LayerSymbols.MLP:
                     # Transformer layers apply their own pp_layer_offset
@@ -153,7 +156,11 @@ def __init__(
                 elif layer_type == LayerSymbols.MOE:
                     # Transformer layers apply their own pp_layer_offset
                     layer = build_module(
-                        submodules.moe_layer, config=self.config, layer_number=i+1,
+                        submodules.moe_layer,
+                        config=self.config,
+                        layer_number=i+1,
+                        pg_collection=pg_collection,
+                        is_mtp_layer=is_mtp_layer,
                     )
                 else:
                     assert False, "unexpected layer_type"
diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index 0b0ee9b6850..e06b9faf7da 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -50,10 +50,12 @@ def __init__(
         config: TransformerConfig,
         layer_number: Optional[int] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
+        is_mtp_layer: bool = False,
     ):
         super(BaseMoELayer, self).__init__(config)
         self.config = config
         self.layer_number = layer_number
+        self.is_mtp_layer = is_mtp_layer
         self.ep_group = pg_collection.ep
         # use pg_collection.expt_tp_group as tensor parallel group in this module.
         self.attn_tp_group = pg_collection.tp
@@ -103,6 +105,7 @@ def __init__(
         submodules: Optional[MoESubmodules] = None,
         layer_number: Optional[int] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
+        is_mtp_layer: bool = False,
     ):
         self.submodules = submodules
         # TODO(Hepteract): delete the usage of the global parallel_state.
@@ -110,7 +113,8 @@ def __init__(
         if pg_collection is None:
             pg_collection = get_default_pg_collection()
         super(MoELayer, self).__init__(
-            config=config, layer_number=layer_number, pg_collection=pg_collection
+            config=config, layer_number=layer_number, pg_collection=pg_collection,
+            is_mtp_layer=is_mtp_layer
         )
         self.moe_layer_recompute = (
             config.recompute_granularity == 'selective' and "moe" in config.recompute_modules
@@ -123,7 +127,7 @@ def __init__(
         self.tp_group = pg_collection.tp
 
         # Initialize router.
-        self.router = TopKRouter(config=self.config, pg_collection=pg_collection)
+        self.router = TopKRouter(config=self.config, pg_collection=pg_collection, is_mtp_layer=is_mtp_layer)
 
         # Initialize latent projections.
         if self.config.moe_latent_size:
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 34d81a21ffa..dc94d394656 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -28,7 +28,8 @@ class Router(ABC, MegatronModule):
     """Base Router class"""
 
     def __init__(
-        self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None
+        self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None,
+        is_mtp_layer: bool = False
     ) -> None:
         """
         Initialize the Router module.
@@ -36,12 +37,14 @@ def __init__(
         Args:
             config (TransformerConfig): Configuration object for the Transformer model.
             pg_collection (ProcessGroupCollection, optional): Process groups for MoE operations.
+            is_mtp_layer (bool): Flag indicating if this router is part of an MTP layer.
         """
         super().__init__(config)
         self.config = config
         self.num_experts = self.config.num_moe_experts
         self.moe_aux_loss_func = None
         self.layer_number = None
+        self.is_mtp_layer = is_mtp_layer
         self.tp_group = pg_collection.tp
         self.cp_group = pg_collection.cp
         self.tp_cp_group = pg_collection.tp_cp
@@ -144,15 +147,17 @@ class TopKRouter(Router):
     """
 
     def __init__(
-        self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None
+        self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None,
+        is_mtp_layer: bool = False
     ) -> None:
         """Initialize the zero token dropping router.
 
         Args:
             config (TransformerConfig): The configuration for the transformer model.
             pg_collection (ProcessGroupCollection, optional): Process groups for MoE operations.
+            is_mtp_layer (bool): Flag indicating if this router is part of an MTP layer.
         """
-        super().__init__(config=config, pg_collection=pg_collection)
+        super().__init__(config=config, pg_collection=pg_collection, is_mtp_layer=is_mtp_layer)
         self.topk = self.config.moe_router_topk
         self.routing_type = self.config.moe_router_load_balancing_type
         self.score_function = self.config.moe_router_score_function
@@ -399,6 +404,11 @@ def attach_and_log_load_balancing_loss(
                 Set this to True if the reduce group has data parallel ranks. This flag is used to
                 ensure the correct reduction in aux loss tracking.
         """
+        # When using repeated MTP layers, the loss is counted "mtp_num_layers" times. To avoid accumulating the
+        # load balancing loss multiple times, we scale it by 1/mtp_num_layers so the total loss is correct.
+        if self.is_mtp_layer and self.config.mtp_use_repeated_layer and self.config.mtp_num_layers is not None:
+            aux_loss = aux_loss / self.config.mtp_num_layers
+
         # TODO (zijiey): fix the per_layer_logging for MTP, currently it will incorrectly
         # add the aux loss logging value to other layer's since it is difficult to get the
         # correct layer_number for MTP. It does not affect the correctness of the calculation
@@ -406,10 +416,16 @@ def attach_and_log_load_balancing_loss(
         num_layers = self.config.num_layers
         if self.config.mtp_num_layers is not None:
             num_layers += self.config.mtp_num_layers
+
+        if self.is_mtp_layer:
+            layer_number = self.layer_number + self.config.num_layers
+        else:
+            layer_number = self.layer_number
+
         save_to_aux_losses_tracker(
             aux_loss_name,
             aux_loss / aux_loss_coeff,
-            self.layer_number,
+            layer_number,
             num_layers,
             reduce_group=reduce_group,
             reduce_group_has_dp=reduce_group_has_dp,
@@ -454,11 +470,22 @@ def apply_z_loss(self, logits):
             else:
                 logits = MoEAuxLossAutoScaler.apply(logits, z_loss)
 
+            # When using repeated MTP layers, the same MTP layer is called mtp_num_layers times. To avoid
+            # accumulating the z_loss multiple times, we scale it by 1/mtp_num_layers so the total loss is correct.
+            if self.is_mtp_layer and self.config.mtp_use_repeated_layer and self.config.mtp_num_layers is not None:
+                z_loss = z_loss / self.config.mtp_num_layers
+
             num_layers = self.config.num_layers
             if self.config.mtp_num_layers is not None:
                 num_layers += self.config.mtp_num_layers
+
+            if self.is_mtp_layer:
+                layer_number = self.layer_number + self.config.num_layers
+            else:
+                layer_number = self.layer_number
+
             save_to_aux_losses_tracker(
-                "z_loss", z_loss / moe_z_loss_coeff, self.layer_number, num_layers
+                "z_loss", z_loss / moe_z_loss_coeff, layer_number, num_layers
             )
         return logits
 
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index f1028799fe4..a9342c9c039 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -615,14 +615,16 @@ def __init__(
                     dtype=self.config.params_dtype,
                     pg_collection=pg_collection,
                     vp_stage=self.vp_stage,
+                    num_layers=self.config.mtp_num_layers_per_layer,
+                    is_mtp_layer=True
                     num_layers=self.config.mtp_num_layers_per_layer
                 )
             else:
                 # Uses the transformer block spec for MTP layer. This option is only implemented for the 
                 # GPT model. In hybrid model, we model transformer block spec for MTP layers with the hybrid
                 # override pattern.
-                self.mtp_model_layer = build_module(self.submodules.mtp_model_layer, config=self.config, vp_stage=self.vp_stage)
-
+                self.mtp_model_layer = build_module(self.submodules.mtp_model_layer, config=self.config, vp_stage=self.vp_stage, is_mtp_layer=True)
+                
         self.final_layernorm = build_module(
             self.submodules.layer_norm,
             config=self.config,
@@ -861,7 +863,6 @@ def forward(
             [s, b, h], and optionally the updated context tensor if cross-attention is used.
         """
         assert context is None, f"multi token prediction + cross attention is not yet supported."
-
         input_ids, position_ids, decoder_input, hidden_states = self._get_embeddings(
             input_ids=input_ids,
             position_ids=position_ids,
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index 24ab98786b6..1335084ad7f 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -266,6 +266,7 @@ def __init__(
         hidden_dropout: Optional[float] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
         vp_stage: Optional[int] = None,
+        is_mtp_layer: bool = False,
     ):
         super().__init__(config=config, vp_stage=vp_stage)
 
@@ -279,6 +280,7 @@ def __init__(
             self.config, vp_stage, get_pg_rank(pg_collection.pp)
         )
         self.hidden_dropout = config.hidden_dropout if hidden_dropout is None else hidden_dropout
+        self.is_mtp_layer = is_mtp_layer
 
         # [Module 1: Input Layernorm] Optional Layernorm on the input data
         # TODO: add pytorch only layernorm
@@ -349,6 +351,9 @@ def __init__(
         if isinstance(submodules.mlp, ModuleSpec):
             if submodules.mlp.module in (MoELayer, GroupedMLP, TEGroupedMLP, SequentialMLP):
                 additional_mlp_kwargs["pg_collection"] = pg_collection
+                # Pass is_mtp_layer flag to MoELayer to distinguish MTP MoE layers.
+                if submodules.mlp.module == MoELayer:
+                    additional_mlp_kwargs["is_mtp_layer"] = self.is_mtp_layer
             elif submodules.mlp.module == MLP:
                 assert hasattr(
                     pg_collection, 'tp'
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index e4359e78184..47d4c0dcc48 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -602,7 +602,6 @@ def validate_args(args, defaults={}):
 
                 assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \
                     'Number of layers should be divisible by the pipeline-model-parallel size'
-
     # Validate MTP hybrid pattern segment count
     if args.mtp_hybrid_override_pattern is not None:
         mtp_segments = len(args.mtp_hybrid_override_pattern.split('|'))
@@ -1797,6 +1796,12 @@ def _add_network_size_args(parser):
                        'We compute the average of the MTP losses across all depths, '
                        'and multiply it the scaling factor to obtain the overall MTP loss, '
                        'which serves as an additional training objective.')
+    group.add_argument('--mtp-use-repeated-layer', action='store_true',
+                       help='Use a single MTP layer repeatedly instead of multiple separate layers. '
+                       'This is more parameter-efficient. When enabled, only 1 MTP layer is created '
+                       'and applied --mtp-num-layers times.')
+    group.add_argument('--mtp-num-layers-per-layer', type=int, default=None,
+                       help='Number of layers inside each MTP layer (e.g., for hybrid models).')
     group.add_argument('--moe-latent-size', type=int, default=None,
                        help='Latent projection dimension for MoE. If None, MoE latent projections are not used.')
     group.add_argument('--mtp-use-repeated-layer', action='store_true',

From 47063da4e03c2dc43854055b0f71aacc667a0ef6 Mon Sep 17 00:00:00 2001
From: Rabeeh Mahabadi <rkarimimahab@nb-hel-cs-001-vscode-02.cm.cluster>
Date: Thu, 18 Dec 2025 10:40:37 -0800
Subject: [PATCH 09/48] added some fixes

---
 megatron/core/models/gpt/gpt_model.py               | 1 -
 megatron/core/transformer/multi_token_prediction.py | 1 -
 2 files changed, 2 deletions(-)

diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index eab4c258427..aa468b0ed9b 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -533,7 +533,6 @@ def _postprocess(
                 packed_seq_params=packed_seq_params,
                 sequence_len_offset=sequence_len_offset,
                 embedding=self.embedding,
-                packed_seq_params=packed_seq_params,
                 **(extra_block_kwargs or {}),
             )
 
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index a9342c9c039..ff43ba71f93 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -617,7 +617,6 @@ def __init__(
                     vp_stage=self.vp_stage,
                     num_layers=self.config.mtp_num_layers_per_layer,
                     is_mtp_layer=True
-                    num_layers=self.config.mtp_num_layers_per_layer
                 )
             else:
                 # Uses the transformer block spec for MTP layer. This option is only implemented for the 

From 4896a80db806dc00c727f0af9ae0ddbbc312949a Mon Sep 17 00:00:00 2001
From: Rabeeh Mahabadi <rkarimimahab@nb-hel-cs-001-vscode-02.cm.cluster>
Date: Thu, 18 Dec 2025 10:45:12 -0800
Subject: [PATCH 10/48] fixed more issues

---
 megatron/training/arguments.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 47d4c0dcc48..9df483d9b80 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1804,12 +1804,6 @@ def _add_network_size_args(parser):
                        help='Number of layers inside each MTP layer (e.g., for hybrid models).')
     group.add_argument('--moe-latent-size', type=int, default=None,
                        help='Latent projection dimension for MoE. If None, MoE latent projections are not used.')
-    group.add_argument('--mtp-use-repeated-layer', action='store_true',
-                       help='Use a single MTP layer repeatedly instead of multiple separate layers. '
-                       'This is more parameter-efficient. When enabled, only 1 MTP layer is created '
-                       'and applied --mtp-num-layers times.')
-    group.add_argument('--mtp-num-layers-per-layer', type=int, default=None,
-                       help='Number of layers inside each MTP layer (e.g., for hybrid models).')
     return parser
 
 

From d5a8dd2feaab4832706146d12d913a701a16d4e3 Mon Sep 17 00:00:00 2001
From: Rabeeh Mahabadi <rkarimimahab@nb-hel-cs-001-vscode-02.cm.cluster>
Date: Fri, 19 Dec 2025 04:55:22 -0800
Subject: [PATCH 11/48] updated

---
 tests/unit_tests/transformer/test_multi_token_prediction.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index e818cde16c4..7b49e4fd852 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -779,6 +779,7 @@ def create_test_args(
         args.hybrid_mlp_ratio = 0.0  
         args.hybrid_override_pattern = "M*M*"  
         args.mtp_hybrid_override_pattern = "M*"  
+        args.mtp_num_layers_per_layer = 2 
         args.spec = "megatron.core.models.mamba.mamba_layer_specs.mamba_stack_spec"
         args.mtp_spec = "megatron.core.models.mamba.mamba_layer_specs.mamba_stack_spec"
         

From 6f5e06ace5971d38e414994e0437a40acaf0e3ab Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Sun, 23 Nov 2025 11:11:32 +0000
Subject: [PATCH 12/48] added mtp support

---
 gpt_builders.py                               |  39 +--
 mamba_builders.py                             |  13 +
 .../common/language_module/language_module.py |  37 ++-
 megatron/core/models/gpt/gpt_layer_specs.py   |   2 +-
 megatron/core/models/gpt/gpt_model.py         |  64 +----
 .../core/models/mamba/mamba_layer_specs.py    |  83 +++++++
 megatron/core/models/mamba/mamba_model.py     | 233 +++++++++++++++++-
 megatron/core/pipeline_parallel/schedules.py  |   4 +-
 megatron/core/ssm/mamba_block.py              |   7 +-
 .../transformer/multi_token_prediction.py     | 177 ++++++++++---
 megatron/training/arguments.py                |  28 +++
 megatron/training/checkpointing.py            |   1 +
 megatron/training/training.py                 |  12 +-
 pretrain_mamba.py                             |   2 +-
 .../test_multi_token_prediction.py            |   4 +-
 15 files changed, 590 insertions(+), 116 deletions(-)

diff --git a/gpt_builders.py b/gpt_builders.py
index 4b86e30e597..9ba7ac8f446 100644
--- a/gpt_builders.py
+++ b/gpt_builders.py
@@ -8,6 +8,7 @@
     get_gpt_layer_with_inference_spec,
     get_gpt_mtp_block_spec,
 )
+from megatron.core.models.mamba.mamba_layer_specs import get_mamba_mtp_block_spec
 from megatron.core.models.gpt.heterogeneous.heterogeneous_layer_specs import (
     get_gpt_heterogeneous_layer_spec,
 )
@@ -60,22 +61,29 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_
                 transformer_layer_spec = _get_transformer_layer_spec(use_te, config)
         mtp_block_spec = None
         if args.mtp_num_layers is not None:
-            assert not (config.transformer_impl == "inference_optimized")
-            if (
-                hasattr(transformer_layer_spec, 'layer_specs')
-                and len(transformer_layer_spec.layer_specs) == 0
-            ):
-                # Get the decoder layer spec explicitly if no decoder layer in the last stage,
-                # Only happens with block spec (TransformerBlockSubmodules) when using MoE.
-                transformer_layer_spec_for_mtp = _get_transformer_layer_spec(use_te, config)
+            if args.mtp_hybrid_override_pattern is not None:
+                # Uses the hybrid model for the MTP module.
+                mtp_mamba_stack_spec = import_module(args.mtp_spec)
+                mtp_block_spec = get_mamba_mtp_block_spec(
+                    config, mtp_mamba_stack_spec, use_transformer_engine=use_te, vp_stage=vp_stage
+                )
             else:
-                transformer_layer_spec_for_mtp = transformer_layer_spec
-            mtp_block_spec = get_gpt_mtp_block_spec(
-                config,
-                transformer_layer_spec_for_mtp,
-                use_transformer_engine=use_te,
-                vp_stage=vp_stage,
-            )
+                assert not (config.transformer_impl == "inference_optimized")
+                if (
+                    hasattr(transformer_layer_spec, 'layer_specs')
+                    and len(transformer_layer_spec.layer_specs) == 0
+                ):
+                    # Get the decoder layer spec explicitly if no decoder layer in the last stage,
+                    # Only happens with block spec (TransformerBlockSubmodules) when using MoE.
+                    transformer_layer_spec_for_mtp = _get_transformer_layer_spec(use_te, config)
+                else:
+                    transformer_layer_spec_for_mtp = transformer_layer_spec
+                mtp_block_spec = get_gpt_mtp_block_spec(
+                    config,
+                    transformer_layer_spec_for_mtp,
+                    use_transformer_engine=use_te,
+                    vp_stage=vp_stage,
+                )
 
         model = GPTModel(
             config=config,
@@ -94,6 +102,7 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_
             mtp_block_spec=mtp_block_spec,
             vp_stage=vp_stage,
             pg_collection=pg_collection,
+            mtp_hybrid_override_pattern=args.mtp_hybrid_override_pattern,
         )
 
     return model
diff --git a/mamba_builders.py b/mamba_builders.py
index 6a792ba6ea5..ba769d43ca2 100644
--- a/mamba_builders.py
+++ b/mamba_builders.py
@@ -1,6 +1,7 @@
 # Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 
 from model_provider import count_parameters_in_layer
+from megatron.core.models.mamba.mamba_layer_specs import get_mamba_mtp_block_spec
 from megatron.core.models.mamba import MambaModel
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.spec_utils import import_module
@@ -22,6 +23,15 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
     else:
         raise ValueError("You must provide a valid Mamba layer spec via --spec")
 
+    mtp_block_spec = None
+    if args.mtp_num_layers is not None:
+        use_te = args.transformer_impl == "transformer_engine"
+        assert args.mtp_spec is not None
+        assert args.mtp_hybrid_override_pattern is not None, "We need to set the override hybrid pattern for MTP layers!"
+        mtp_mamba_stack_spec = import_module(args.mtp_spec)
+        mtp_block_spec = get_mamba_mtp_block_spec(
+            config, mtp_mamba_stack_spec, use_transformer_engine=use_te, vp_stage=vp_stage
+        )
     model = MambaModel(
         config=config,
         mamba_stack_spec=mamba_stack_spec,
@@ -31,6 +41,7 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
         hybrid_attention_ratio=args.hybrid_attention_ratio,
         hybrid_mlp_ratio=args.hybrid_mlp_ratio,
         hybrid_override_pattern=args.hybrid_override_pattern,
+        mtp_hybrid_override_pattern=args.mtp_hybrid_override_pattern,
         post_process=post_process,
         fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
         parallel_output=True,
@@ -39,6 +50,8 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
         rotary_percent=args.rotary_percent,
         rotary_base=args.rotary_base,
         pg_collection=pg_collection,
+        vp_stage=vp_stage,
+        mtp_block_spec=mtp_block_spec,
     )
 
     for l in range(model.decoder.num_layers_per_pipeline_rank):
diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
index 19eea55dec3..706f8429a29 100644
--- a/megatron/core/models/common/language_module/language_module.py
+++ b/megatron/core/models/common/language_module/language_module.py
@@ -31,6 +31,10 @@
     make_tp_sharded_tensor_for_checkpoint,
 )
 
+from megatron.core.transformer.multi_token_prediction import (
+    tie_output_layer_state_dict,
+    tie_word_embeddings_state_dict,
+)
 
 class LanguageModule(MegatronModule):
     """Base language module that has common helper functions used across GPT, BERT etc.
@@ -250,12 +254,20 @@ def setup_embeddings_and_output_layer(self) -> None:
             LanguageModule.embedding_warning_printed = True
 
     def shared_embedding_or_output_weight(self) -> Tensor:
-        """Gets the emedding weight or output logit weights when share embedding and output weights set to True.
+        """Gets the emedding weight or output logit weights when share embedding and output weights set to True
+          or when use Multi-Token Prediction (MTP).
 
         Returns:
-            Tensor: During pre processing it returns the input embeddings weight while during post processing it returns the final output layers weight
+            Tensor: During pre processing or MTP process it returns the input embeddings weight while during post processing it returns the final output layers weight
         """
-        if self.pre_process:
+        if self.pre_process or self.mtp_process:
+            # Multi-Token Prediction (MTP) need both embedding layer and output layer.
+            # So there will be both embedding layer and output layer in the mtp process stage.
+            # In this case, if share_embeddings_and_output_weights is True, the shared weights
+            # will be stored in embedding layer, and output layer will not have any weight.
+            assert hasattr(
+                self, 'embedding'
+            ), f"embedding is needed in this pipeline stage, but it is not initialized."
             return self.embedding.word_embeddings.weight
         elif self.post_process:
             return self.output_layer.weight
@@ -288,6 +300,25 @@ def sharded_state_dict(
         output_layer_weight_key = f'{prefix}output_layer.weight'
         output_layer_bias_key = f'{prefix}output_layer.bias'
 
+        # Multi-Token Prediction (MTP) need both embedding layer and output layer in
+        # mtp process stage.
+        # If MTP is not placed in the pre processing stage, we need to maintain a copy of
+        # embedding layer in the mtp process stage and tie it to the embedding in the pre
+        # processing stage.
+        # Also, if MTP is not placed in the post processing stage, we need to maintain a copy
+        # of output layer in the mtp process stage and tie it to the output layer in the post
+        # processing stage.
+        if self.mtp_process and not self.pre_process:
+            emb_weight = self.embedding.word_embeddings.weight
+            tie_word_embeddings_state_dict(sharded_state_dict, emb_weight, first_stage_word_emb_key)
+
+        if self.mtp_process and not self.post_process:
+            # We only need to tie the output layer weight if share_embeddings_and_output_weights
+            # is False. Because if share_embeddings_and_output_weights is True, the shared weight
+            # will be stored in embedding layer, and output layer will not have any weight.
+            if not self.share_embeddings_and_output_weights:
+                output_layer_weight = self.output_layer.weight
+                tie_output_layer_state_dict(sharded_state_dict, output_layer_weight, output_layer_weight_key)
         if self.share_embeddings_and_output_weights:
             self.tie_embeddings_and_output_weights_state_dict(
                 sharded_state_dict, output_layer_weight_key, first_stage_word_emb_key, metadata
diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
index b61a544c5df..900232fa708 100755
--- a/megatron/core/models/gpt/gpt_layer_specs.py
+++ b/megatron/core/models/gpt/gpt_layer_specs.py
@@ -686,7 +686,7 @@ def get_gpt_mtp_block_spec_for_backend(
         raise ValueError(f"Invalid spec: {spec}")
 
     mtp_layer_spec = get_mtp_layer_spec_for_backend(
-        transformer_layer_spec=transformer_layer_spec, backend=backend
+        mtp_model_layer_spec=transformer_layer_spec, backend=backend
     )
     mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0
     mtp_layer_specs = [mtp_layer_spec] * mtp_num_layers
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index 9342e96ce9a..eab4c258427 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -102,6 +102,7 @@ def __init__(
         mtp_block_spec: Optional[ModuleSpec] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
         vp_stage: Optional[int] = None,
+        mtp_hybrid_override_pattern: str = None,
     ) -> None:
         super().__init__(config=config, pg_collection=pg_collection)
 
@@ -210,7 +211,7 @@ def __init__(
 
         if self.mtp_process:
             self.mtp = MultiTokenPredictionBlock(
-                config=self.config, spec=self.mtp_block_spec, vp_stage=vp_stage
+                config=self.config, spec=self.mtp_block_spec, vp_stage=vp_stage, mtp_hybrid_override_pattern=mtp_hybrid_override_pattern
             )
 
         # Output
@@ -532,6 +533,7 @@ def _postprocess(
                 packed_seq_params=packed_seq_params,
                 sequence_len_offset=sequence_len_offset,
                 embedding=self.embedding,
+                packed_seq_params=packed_seq_params,
                 **(extra_block_kwargs or {}),
             )
 
@@ -539,56 +541,16 @@ def _postprocess(
             return hidden_states
 
         if self.mtp_process:
-            mtp_labels = labels.clone()
-            hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0)
-            hidden_states = hidden_states_list[0]
-            if loss_mask is None:
-                # if loss_mask is not provided, use all ones as loss_mask
-                loss_mask = torch.ones_like(mtp_labels)
-            for mtp_layer_number in range(self.config.mtp_num_layers):
-                # output
-                mtp_logits, _ = self.output_layer(
-                    hidden_states_list[mtp_layer_number + 1],
-                    weight=output_weight,
-                    runtime_gather_output=runtime_gather_output,
-                )
-                # Calc loss for the current Multi-Token Prediction (MTP) layers.
-                mtp_labels, _ = roll_tensor(
-                    mtp_labels,
-                    shifts=-1,
-                    dims=-1,
-                    cp_group=self.cp_group,
-                    packed_seq_params=packed_seq_params,
-                )
-                loss_mask, num_tokens = roll_tensor(
-                    loss_mask,
-                    shifts=-1,
-                    dims=-1,
-                    cp_group=self.cp_group,
-                    packed_seq_params=packed_seq_params,
-                )
-                mtp_loss = self.compute_language_model_loss(mtp_labels, mtp_logits)
-                mtp_loss = loss_mask * mtp_loss
-                if self.training:
-                    # TODO(shifangx): remove the use of parallel_state here
-                    # after moving loss logging to loss_func in pretrain_gpt.py
-                    MTPLossLoggingHelper.save_loss_to_tracker(
-                        torch.sum(mtp_loss) / num_tokens,
-                        mtp_layer_number,
-                        self.config.mtp_num_layers,
-                        avg_group=parallel_state.get_data_parallel_group(
-                            with_context_parallel=True
-                        ),
-                    )
-                mtp_loss_scale = self.config.mtp_loss_scaling_factor / self.config.mtp_num_layers
-                if self.config.calculate_per_token_loss:
-                    hidden_states = MTPLossAutoScaler.apply(
-                        hidden_states, mtp_loss_scale * mtp_loss
-                    )
-                else:
-                    hidden_states = MTPLossAutoScaler.apply(
-                        hidden_states, mtp_loss_scale * mtp_loss / num_tokens
-                    )
+            hidden_states = self.mtp.process_loss(
+                hidden_states=hidden_states,
+                labels=labels,
+                loss_mask=loss_mask,
+                output_layer=self.output_layer,
+                output_weight=output_weight,
+                runtime_gather_output=runtime_gather_output,
+                is_training=self.training,
+                compute_language_model_loss=self.compute_language_model_loss,
+            )
         sequence_parallel_override = False
         if in_inference_mode and inference_context.materialize_only_last_token_logits:
             if inference_context.is_static_batching():
diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index f83275ed9c6..286d72b36ff 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -6,6 +6,7 @@
     TENorm,
     TERowParallelLinear,
 )
+from typing import Optional
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec
 from megatron.core.ssm.mamba_block import MambaStack, MambaStackSubmodules
@@ -21,6 +22,15 @@
 from megatron.core.transformer.mlp import MLP, MLPSubmodules
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
+from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionBlockSubmodules
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.multi_token_prediction import (
+    MultiTokenPredictionBlockSubmodules,
+    get_mtp_layer_offset,
+    get_mtp_layer_spec_for_backend,
+    get_mtp_num_layers_to_build,
+)
+from megatron.core.models.backends import BackendSpecProvider, LocalSpecProvider
 
 moe = get_moe_module_spec(
     use_te=True,
@@ -29,6 +39,22 @@
     moe_use_legacy_grouped_gemm=False,
 )
 
+try:
+    from megatron.core.extensions.transformer_engine import TEFusedMLP, TENorm
+    from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
+
+    HAVE_TE = True
+except ImportError:
+    HAVE_TE = False
+
+
+try:
+    from megatron.core.extensions.kitchen import KitchenSpecProvider
+
+    HAVE_KITCHEN = True
+except ImportError:
+    HAVE_KITCHEN = False
+
 mamba_stack_spec = ModuleSpec(
     module=MambaStack,
     submodules=MambaStackSubmodules(
@@ -87,6 +113,7 @@
     ),
 )
 
+<<<<<<< HEAD
 mamba_inference_stack_spec = ModuleSpec(
     module=MambaStack,
     submodules=MambaStackSubmodules(
@@ -146,3 +173,59 @@
         ),
     ),
 )
+=======
+
+def get_mamba_mtp_block_spec(
+    config: TransformerConfig,
+    spec: ModuleSpec,
+    use_transformer_engine: bool, 
+    vp_stage: Optional[int] = None,
+) -> MultiTokenPredictionBlockSubmodules:
+    """Mamba Multi-Token Prediction (MTP) block spec."""
+    if use_transformer_engine:
+        backend: BackendSpecProvider = (
+            KitchenSpecProvider(fallback=TESpecProvider())
+            if config.use_kitchen
+            else TESpecProvider()
+        )
+    else:
+        backend = (
+            KitchenSpecProvider(fallback=LocalSpecProvider())
+            if config.use_kitchen
+            else LocalSpecProvider()
+        )
+    return get_mamba_mtp_block_spec_for_backend(
+        config=config, spec=spec, backend=backend, vp_stage=vp_stage
+    )
+    
+def get_mamba_mtp_block_spec_for_backend(
+    config: TransformerConfig,
+    spec: ModuleSpec, 
+    backend: BackendSpecProvider,
+    vp_stage: Optional[int] = None,
+) -> MultiTokenPredictionBlockSubmodules:
+    """Mamba Multi-Token Prediction (MTP) block spec."""
+    num_layers_to_build = get_mtp_num_layers_to_build(config, vp_stage=vp_stage)
+    if num_layers_to_build == 0:
+        return None
+
+    mtp_layer_spec = get_mtp_layer_spec_for_backend(
+        mtp_model_layer_spec=spec, backend=backend
+    )
+    mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0
+    mtp_layer_specs = [mtp_layer_spec] * mtp_num_layers
+
+    offset = get_mtp_layer_offset(config)
+    # split the mtp layer specs to only include the layers that are built in this pipeline stage.
+    mtp_layer_specs = mtp_layer_specs[offset : offset + num_layers_to_build]
+    if len(mtp_layer_specs) > 0:
+        assert (
+            len(mtp_layer_specs) == config.mtp_num_layers
+        ), +f"currently all of the mtp layers must stage in the same pipeline stage."
+        mtp_block_spec = MultiTokenPredictionBlockSubmodules(layer_specs=mtp_layer_specs)
+    else:
+        mtp_block_spec = None
+
+    return mtp_block_spec
+
+>>>>>>> a05cbed31 (added mtp support)
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 378cf7e47d6..6521100648b 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -3,8 +3,10 @@
 from typing import Literal, Optional
 
 from torch import Tensor
-
-from megatron.core import tensor_parallel
+import torch
+from typing import Dict
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core import parallel_state, tensor_parallel
 from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
 from megatron.core.inference.contexts import BaseInferenceContext
 from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
@@ -17,6 +19,15 @@
 from megatron.core.transformer.enums import ModelType
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.utils import WrappedTensor, deprecate_inference_params
+from megatron.core.transformer.multi_token_prediction import (
+    MTPLossAutoScaler,
+    MTPLossLoggingHelper,
+    MultiTokenPredictionBlock,
+    roll_tensor,
+    tie_output_layer_state_dict,
+    tie_word_embeddings_state_dict,
+)
+from megatron.core.packed_seq_params import PackedSeqParams
 
 
 class MambaModel(LanguageModule):
@@ -63,6 +74,7 @@ def __init__(
         hybrid_attention_ratio: float = 0.0,
         hybrid_mlp_ratio: float = 0.0,
         hybrid_override_pattern: str = None,
+        mtp_hybrid_override_pattern: str = None,
         post_process: bool = True,
         fp16_lm_cross_entropy: bool = False,
         parallel_output: bool = True,
@@ -74,6 +86,8 @@ def __init__(
         scatter_embedding_sequence_parallel: bool = True,
         seq_len_interpolation_factor: Optional[float] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
+        vp_stage: Optional[int] = None,
+        mtp_block_spec: Optional[ModuleSpec] = None,
     ) -> None:
         super().__init__(config=config, pg_collection=pg_collection)
 
@@ -87,17 +101,24 @@ def __init__(
         self.hybrid_attention_ratio = hybrid_attention_ratio
         self.hybrid_mlp_ratio = hybrid_mlp_ratio
         self.hybrid_override_pattern = hybrid_override_pattern
+        self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern
         self.post_process = post_process
         self.fp16_lm_cross_entropy = fp16_lm_cross_entropy
         self.parallel_output = parallel_output
         self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
         self.position_embedding_type = position_embedding_type
+        self.mtp_block_spec = mtp_block_spec
+        self.mtp_process = mtp_block_spec is not None
+
+        # Cache for RoPE tensors which do not change between iterations.
+        self.rotary_pos_emb_cache = {}
+
 
         # megatron core pipelining currently depends on model type
         # TODO: remove this dependency ?
         self.model_type = ModelType.encoder_or_decoder
 
-        if self.pre_process:
+        if self.pre_process or self.mtp_process:
             self.embedding = LanguageModelEmbedding(
                 config=self.config,
                 vocab_size=self.vocab_size,
@@ -129,8 +150,13 @@ def __init__(
             pg_collection=self.pg_collection,
         )
 
+        if self.mtp_process:
+            self.mtp = MultiTokenPredictionBlock(
+                config=self.config, spec=self.mtp_block_spec, vp_stage=vp_stage, mtp_hybrid_override_pattern=self.mtp_hybrid_override_pattern
+            )
+
         # Output
-        if post_process:
+        if post_process or self.mtp_process:
             self.output_layer = tensor_parallel.ColumnParallelLinear(
                 config.hidden_size,
                 self.vocab_size,
@@ -168,6 +194,94 @@ def set_input_tensor(self, input_tensor: Tensor) -> None:
         assert len(input_tensor) == 1, 'input_tensor should only be length 1 for gpt/bert'
         self.decoder.set_input_tensor(input_tensor[0])
 
+    def _preprocess(
+        self,
+        input_ids: Tensor,
+        position_ids: Tensor,
+        decoder_input: Tensor = None,
+        inference_context: BaseInferenceContext = None,
+        packed_seq_params: PackedSeqParams = None,
+    ):
+        """Preprocesses inputs for the transformer decoder.
+
+        Applies embeddings to input tokens, or uses `decoder_input` from a previous
+        pipeline stage. Also sets up rotary positional embeddings.
+        """
+
+        # If decoder_input is provided (not None), then input_ids and position_ids are ignored.
+        # Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input.
+
+        # Decoder embedding.
+        if decoder_input is not None:
+            pass
+        elif self.pre_process:
+            decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
+        else:
+            # intermediate stage of pipeline
+            # decoder will get hidden_states from encoder.input_tensor
+            decoder_input = None
+
+        # Rotary positional embeddings (embedding is None for PP intermediate devices)
+        rotary_pos_emb = None
+        rotary_pos_cos = None
+        rotary_pos_sin = None
+        if self.position_embedding_type == 'rope' and not self.config.multi_latent_attention:
+            if not self.training and self.config.flash_decode and inference_context:
+                assert (
+                    inference_context.is_static_batching()
+                ), "GPTModel currently only supports static inference batching."
+                # Flash decoding uses precomputed cos and sin for RoPE
+                rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb_cache.setdefault(
+                    inference_context.max_sequence_length,
+                    self.rotary_pos_emb.get_cos_sin(inference_context.max_sequence_length),
+                )
+            else:
+                rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
+                    inference_context, self.decoder, decoder_input, self.config, packed_seq_params
+                )
+                rotary_pos_emb = self.rotary_pos_emb(
+                    rotary_seq_len,
+                    packed_seq=packed_seq_params is not None
+                    and packed_seq_params.qkv_format == 'thd',
+                )
+        elif self.position_embedding_type == 'mrope' and not self.config.multi_latent_attention:
+            if self.training or not self.config.flash_decode:
+                rotary_pos_emb = self.rotary_pos_emb(position_ids, self.mrope_section)
+            else:
+                # Flash decoding uses precomputed cos and sin for RoPE
+                raise NotImplementedError(
+                    "Flash decoding uses precomputed cos and sin for RoPE, not implmented in "
+                    "MultimodalRotaryEmbedding yet."
+                )
+
+        if (
+            (self.config.enable_cuda_graph or self.config.flash_decode)
+            and rotary_pos_cos is not None
+            and inference_context
+            and inference_context.is_static_batching()
+            and not self.training
+        ):
+            current_batch_size = input_ids.shape[0]
+            sequence_len_offset = torch.tensor(
+                [inference_context.sequence_len_offset] * current_batch_size,
+                dtype=torch.int32,
+                device=rotary_pos_cos.device,  # Co-locate this with the rotary tensors
+            )
+        else:
+            sequence_len_offset = None
+
+        # Wrap decoder_input to allow the decoder (TransformerBlock) to delete the
+        # reference held by this caller function, enabling early garbage collection for
+        # inference. Skip wrapping if decoder_input is logged after decoder completion.
+        if (
+            inference_context is not None
+            and not self.training
+            and not has_config_logger_enabled(self.config)
+        ):
+            decoder_input = WrappedTensor(decoder_input)
+
+        return decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset
+
     def forward(
         self,
         input_ids: Tensor,
@@ -179,6 +293,7 @@ def forward(
         runtime_gather_output: Optional[bool] = None,
         *,
         inference_params: Optional[BaseInferenceContext] = None,
+        loss_mask: Optional[Tensor] = None,
     ) -> Tensor:
         """Forward function of the Mamba model. This function passes the input tensors
         through the embedding layer, and then the decoder and finally into the post
@@ -237,9 +352,53 @@ def forward(
             rotary_pos_emb=rotary_pos_emb,
         )
 
+        # Get output weight for shared embeddings
+        output_weight = None
+        if self.share_embeddings_and_output_weights:
+            output_weight = self.shared_embedding_or_output_weight()            
+
+        if self.mtp_process:
+            extra_block_kwargs = None
+            packed_seq_params = None
+            decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset = self._preprocess(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                decoder_input=decoder_input,
+                inference_context=inference_context,
+                packed_seq_params=packed_seq_params,
+            )
+
+
+            hidden_states = self.mtp(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                inference_params=inference_params,
+                rotary_pos_emb=rotary_pos_emb,
+                rotary_pos_cos=rotary_pos_cos,
+                rotary_pos_sin=rotary_pos_sin,
+                packed_seq_params=packed_seq_params,
+                sequence_len_offset=sequence_len_offset,
+                embedding=self.embedding,
+                **(extra_block_kwargs or {}),
+            )
+
         if not self.post_process:
             return hidden_states
 
+        if self.mtp_process:
+            hidden_states = self.mtp.process_loss(
+                hidden_states=hidden_states,
+                labels=labels,
+                loss_mask=loss_mask,
+                output_layer=self.output_layer,
+                output_weight=output_weight,
+                runtime_gather_output=runtime_gather_output,
+                is_training=self.training,
+                compute_language_model_loss=self.compute_language_model_loss,
+            )
+
         # logits and loss
         output_weight = None
         if self.share_embeddings_and_output_weights:
@@ -260,7 +419,7 @@ def forward(
                     self.output_layer.sequence_parallel = False
                     sequence_parallel_override = True
 
-                # Reshape [B, 1, H] to [1, B, H] → extract each sample’s true last‐token hidden
+                # Reshape [B, 1, H] to [1, B, H] → extract each sample's true last‐token hidden
                 # state ([B, H]) → unsqueeze back to [B, 1, H]
                 # (so that the output layer, which expects S×B×H, receives only the final token)
                 hidden_states = inference_context.last_token_logits(
@@ -287,3 +446,67 @@ def forward(
         loss = self.compute_language_model_loss(labels, logits)
 
         return loss
+
+    def shared_embedding_or_output_weight(self) -> Optional[Tensor]:
+        """Gets the embedding weight or output logit weights when share input embedding and
+        output weights set to True or when use Multi-Token Prediction (MTP) feature.
+
+        Returns:
+            Optional[Tensor]: During pre processing or MTP process it returns the input embeddings weight.
+            Otherwise, during post processing it returns the final output layers weight.
+        """
+        if self.pre_process or self.mtp_process:
+            # Multi-Token Prediction (MTP) need both embedding layer and output layer.
+            # So there will be both embedding layer and output layer in the mtp process stage.
+            # In this case, if share_embeddings_and_output_weights is True, the shared weights
+            # will be stored in embedding layer, and output layer will not have any weight.
+            assert hasattr(
+                self, 'embedding'
+            ), f"embedding is needed in this pipeline stage, but it is not initialized."
+            return self.embedding.word_embeddings.weight
+        elif self.post_process:
+            return self.output_layer.weight
+        return None
+
+    def sharded_state_dict(
+        self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[Dict] = None
+    ) -> ShardedStateDict:
+        """Sharded state dict implementation for MambaModel backward-compatibility.
+
+        Removing extra state.
+        Tie word embeddings and output layer in mtp process stage.
+
+        Args:
+            prefix (str): Module name prefix.
+            sharded_offsets (tuple): PP related offsets, expected to be empty at this module level.
+            metadata (Optional[Dict]): metadata controlling sharded state dict creation.
+
+        Returns:
+            ShardedStateDict: sharded state dict for the MambaModel
+        """
+        sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+        
+        # Multi-Token Prediction (MTP) need both embedding layer and output layer in
+        # mtp process stage.
+        # If MTP is not placed in the pre processing stage, we need to maintain a copy of
+        # embedding layer in the mtp process stage and tie it to the embedding in the pre
+        # processing stage.
+        # Also, if MTP is not placed in the post processing stage, we need to maintain a copy
+        # of output layer in the mtp process stage and tie it to the output layer in the post
+        # processing stage.
+        if self.mtp_process and not self.pre_process:
+            emb_weight_key = f'{prefix}embedding.word_embeddings.weight'
+            emb_weight = self.embedding.word_embeddings.weight
+            tie_word_embeddings_state_dict(sharded_state_dict, emb_weight, emb_weight_key)
+        if self.mtp_process and not self.post_process:
+            # We only need to tie the output layer weight if share_embeddings_and_output_weights
+            # is False. Because if share_embeddings_and_output_weights is True, the shared weight
+            # will be stored in embedding layer, and output layer will not have any weight.
+            if not self.share_embeddings_and_output_weights:
+                output_layer_weight_key = f'{prefix}output_layer.weight'
+                output_layer_weight = self.output_layer.weight
+                tie_output_layer_state_dict(
+                    sharded_state_dict, output_layer_weight, output_layer_weight_key
+                )
+
+        return sharded_state_dict
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index 5ab4cb57545..3c1fb2e3b18 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -193,7 +193,7 @@ def set_current_microbatch(model, microbatch_id):
             layer.current_microbatch = microbatch_id
         if hasattr(model_with_decoder, 'mtp'):
             for layer in model_with_decoder.mtp.layers:
-                layer.transformer_layer.current_microbatch = microbatch_id
+                layer.mtp_model_layer.current_microbatch = microbatch_id
 
 
 def forward_step_calc_loss(
@@ -384,7 +384,7 @@ def forward_step(
 
     if is_first_microbatch and hasattr(model, 'set_is_first_microbatch'):
         model.set_is_first_microbatch()
-    if current_microbatch is not None:
+    if config.external_cuda_graph and current_microbatch is not None:
         set_current_microbatch(model, current_microbatch)
 
     unwrap_output_tensor = False
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index 1abd12186e3..4e8bc27cc54 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -83,12 +83,15 @@ def __init__(
         device=None,
         dtype=None,
         pg_collection: ProcessGroupCollection = None,
+        vp_stage: Optional[int] = None,
+        num_layers: int = None,
     ) -> None:
         super().__init__(config=config)
         self.residual_in_fp32 = residual_in_fp32
         self.pre_process = pre_process
         self.post_layer_norm = post_layer_norm
         self.post_process = post_process
+        self.vp_stage = vp_stage
 
         assert pg_collection is not None, "pg_collection must be provided for MambaStack"
 
@@ -103,7 +106,7 @@ def __init__(
         self.hybrid_override_pattern = hybrid_override_pattern
 
         self.layer_type_list = allocate_layers(
-            self.config.num_layers,
+            self.config.num_layers if num_layers is None else num_layers,
             self.hybrid_attention_ratio,
             self.hybrid_mlp_ratio,
             self.hybrid_override_pattern,
@@ -147,7 +150,7 @@ def __init__(
                 elif layer_type == LayerSymbols.MOE:
                     # Transformer layers apply their own pp_layer_offset
                     layer = build_module(
-                        submodules.moe_layer, config=self.config, layer_number=i + 1
+                        submodules.moe_layer, config=self.config, layer_number=i + 1,
                     )
                 else:
                     assert False, "unexpected layer_type"
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 8135e1c47a3..f86aa9dd401 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -395,19 +395,19 @@ class MultiTokenPredictionLayerSubmodules:
             embedding normalization to be applied.
         eh_proj (Union[ModuleSpec, type]): Specification or instance of the
             linear projection to be applied.
-        transformer_layer (Union[ModuleSpec, type]): Specification
-            or instance of the transformer block to be applied.
+        mtp_model_layer (Union[ModuleSpec, type]): Specification
+            or instance of the transformer or mamba block to be applied.
     """
 
     enorm: Union[ModuleSpec, type] = None
     hnorm: Union[ModuleSpec, type] = None
     eh_proj: Union[ModuleSpec, type] = None
-    transformer_layer: Union[ModuleSpec, type] = None
+    mtp_model_layer: Union[ModuleSpec, type] = None
     layer_norm: Union[ModuleSpec, type] = None
 
 
 def get_mtp_layer_spec(
-    transformer_layer_spec: ModuleSpec, use_transformer_engine: bool
+    mtp_model_layer_spec: ModuleSpec, use_transformer_engine: bool
 ) -> ModuleSpec:
     """Get the MTP layer spec.
 
@@ -415,13 +415,13 @@ def get_mtp_layer_spec(
         ModuleSpec: Module specification with TE modules
     """
     return get_mtp_layer_spec_for_backend(
-        transformer_layer_spec,
+        mtp_model_layer_spec,
         backend=TESpecProvider() if use_transformer_engine else LocalSpecProvider(),
     )
 
 
 def get_mtp_layer_spec_for_backend(
-    transformer_layer_spec: ModuleSpec, backend: BackendSpecProvider
+   mtp_model_layer_spec: ModuleSpec, backend: BackendSpecProvider
 ) -> ModuleSpec:
     """Get the MTP layer spec.
 
@@ -436,7 +436,7 @@ def get_mtp_layer_spec_for_backend(
             enorm=layer_norm_impl,
             hnorm=layer_norm_impl,
             eh_proj=column_parallel_linear_impl,
-            transformer_layer=transformer_layer_spec,
+            mtp_model_layer=mtp_model_layer_spec,
             layer_norm=layer_norm_impl,
         ),
     )
@@ -534,6 +534,7 @@ def __init__(
         self,
         config: TransformerConfig,
         submodules: MultiTokenPredictionLayerSubmodules,
+        mtp_hybrid_override_pattern: str = None,
         layer_number: int = 1,
         vp_stage: Optional[int] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
@@ -544,14 +545,18 @@ def __init__(
         self.layer_number = layer_number
         self.vp_stage = vp_stage
         self.cp_group = pg_collection.cp
-
-        self_attention_spec = self.submodules.transformer_layer.submodules.self_attention
-        attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
-        assert attn_mask_type in SUPPORTED_ATTN_MASK, (
-            f"Multi-Token Prediction (MTP) is not jet supported with "
-            + f"{attn_mask_type} attention mask type."
-            + f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
-        )
+        self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern
+
+        if hasattr(self.submodules.mtp_model_layer.submodules, 'attention_layer'):
+            self_attention_spec = self.submodules.mtp_model_layer.submodules.attention_layer
+            if self_attention_spec.submodules.self_attention is not None:
+                    self_attention_spec = self_attention_spec.submodules.self_attention
+                    attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
+                    assert attn_mask_type in SUPPORTED_ATTN_MASK, (
+                        f"Multi-Token Prediction (MTP) is not jet supported with "
+                        + f"{attn_mask_type} attention mask type."
+                        + f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
+                    )
 
         self.enorm = build_module(
             self.submodules.enorm,
@@ -582,10 +587,29 @@ def __init__(
             bias=False,
             skip_bias_add=False,
             is_expert=False,
+            tp_comm_buffer_name="mtp_eh_proj"
         )
-        self.transformer_layer = build_module(
-            self.submodules.transformer_layer, config=self.config, vp_stage=vp_stage
-        )
+        if self.config.mtp_num_layers is not None:
+            if self.mtp_hybrid_override_pattern is not None:
+                pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+                # We do not need pre and post process stage for MTP layer, given they are handled in the 
+                # MultiTokenPredictionLayer itself.
+                self.mtp_model_layer = build_module(
+                    self.submodules.mtp_model_layer,
+                    self.config,
+                    pre_process=False, 
+                    post_process=False,
+                    hybrid_override_pattern=self.mtp_hybrid_override_pattern,
+                    dtype=self.config.params_dtype,
+                    pg_collection=pg_collection,
+                    vp_stage=self.vp_stage,
+                    num_layers=self.config.mtp_num_layers
+                )
+            else:
+                # Uses the transformer block spec for MTP layer. This option is only implemented for the 
+                # GPT model. In hybrid model, we model transformer block spec for MTP layers with the hybrid
+                # override pattern.
+                self.mtp_model_layer = build_module(self.submodules.mtp_model_layer, config=self.config, vp_stage=self.vp_stage)
 
         self.final_layernorm = build_module(
             self.submodules.layer_norm,
@@ -696,7 +720,6 @@ def _proj_and_transformer_layer(
             transformer_layer_fp8_context = nullcontext()
 
         # TODO: currently ignoring FP4 in MTP layers because we need more numerical validation
-
         with rng_context:
             with fp8_context:
                 hidden_states = self._concat_embeddings(hidden_states, decoder_input)
@@ -705,19 +728,30 @@ def _proj_and_transformer_layer(
             # transformer layer is cudagraphed, the FP8GlobalStateManager.is_first_fp8_module() is
             # True so that the fp8 weight caching can be triggered correctly.
             with transformer_layer_fp8_context:
-                hidden_states, _ = self.transformer_layer(
-                    hidden_states=hidden_states,
-                    attention_mask=attention_mask,
-                    context=context,
-                    context_mask=context_mask,
-                    rotary_pos_emb=rotary_pos_emb,
-                    rotary_pos_cos=rotary_pos_cos,
-                    rotary_pos_sin=rotary_pos_sin,
-                    attention_bias=attention_bias,
-                    inference_params=inference_params,
-                    packed_seq_params=packed_seq_params,
-                    sequence_len_offset=sequence_len_offset,
-                )
+                if self.mtp_hybrid_override_pattern is not None:
+                    # Since pre-process is set to False, we need to set the input tensor manually.
+                    self.mtp_model_layer.set_input_tensor(hidden_states)
+                    hidden_states = self.mtp_model_layer(
+                        hidden_states=hidden_states,
+                        attention_mask=attention_mask,
+                        rotary_pos_emb=rotary_pos_emb,
+                        inference_params=inference_params,
+                        inference_context=context,
+                    )
+                else:
+                    hidden_states, _ = self.mtp_model_layer(
+                        hidden_states=hidden_states,
+                        attention_mask=attention_mask,
+                        context=context,
+                        context_mask=context_mask,
+                        rotary_pos_emb=rotary_pos_emb,
+                        rotary_pos_cos=rotary_pos_cos,
+                        rotary_pos_sin=rotary_pos_sin,
+                        attention_bias=attention_bias,
+                        inference_params=inference_params,
+                        packed_seq_params=packed_seq_params,
+                        sequence_len_offset=sequence_len_offset,
+                    )
 
         hidden_states = self._postprocess(hidden_states)
 
@@ -948,9 +982,11 @@ def __init__(
         config: TransformerConfig,
         spec: Union[TransformerBlockSubmodules, ModuleSpec],
         vp_stage: Optional[int] = None,
+        mtp_hybrid_override_pattern: str = None,
         pg_collection: ProcessGroupCollection = None,
     ):
         super().__init__(config=config)
+        self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern
         self.submodules = _get_mtp_block_submodules(config, spec)
         self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor
         self.vp_stage = vp_stage
@@ -981,6 +1017,7 @@ def build_layer(layer_spec, layer_number):
                     layer_number=layer_number,
                     vp_stage=self.vp_stage,
                     pg_collection=pg_collection,
+                    mtp_hybrid_override_pattern=self.mtp_hybrid_override_pattern
                 )
             return module
 
@@ -1049,6 +1086,82 @@ def forward(
         hidden_states = torch.cat(hidden_states_list, dim=0)
         return hidden_states
 
+    def process_loss(
+        self,
+        hidden_states: Tensor,
+        labels: Tensor,
+        loss_mask: Optional[Tensor],
+        output_layer: Callable,
+        output_weight: Optional[Tensor],
+        runtime_gather_output: Optional[bool],
+        is_training: bool,
+        compute_language_model_loss: Callable,
+        packed_seq_params: Optional[PackedSeqParams] = None,
+    ) -> Tensor:
+        """Process Multi-Token Prediction (MTP) loss computation.
+        
+        This method handles the MTP loss computation for multiple prediction layers.
+        It chunks the hidden states, computes logits and losses for each MTP layer,
+        and applies loss scaling.
+        
+        Args:
+            hidden_states (Tensor): Hidden states tensor from the model.
+            labels (Tensor): Ground truth labels.
+            loss_mask (Optional[Tensor]): Mask for loss computation. If None, uses all ones.
+            output_layer (Callable): Output layer method to compute logits.
+            output_weight (Optional[Tensor]): Optional output weight for shared embeddings.
+            runtime_gather_output (Optional[bool]): Whether to gather output at runtime.
+            is_training (bool): Whether the model is in training mode.
+            compute_language_model_loss (Callable): Method to compute language model loss.
+        
+        Returns:
+            Tensor: Updated hidden states after MTP loss processing.
+        """
+        mtp_labels = labels.clone()
+        hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0)
+        hidden_states = hidden_states_list[0]
+        
+        if loss_mask is None:
+            # if loss_mask is not provided, use all ones as loss_mask
+            loss_mask = torch.ones_like(mtp_labels)
+        
+        for mtp_layer_number in range(self.config.mtp_num_layers):
+            # output
+            mtp_logits, _ = output_layer(
+                hidden_states_list[mtp_layer_number + 1],
+                weight=output_weight,
+                runtime_gather_output=runtime_gather_output,
+            )
+            # Calc loss for the current Multi-Token Prediction (MTP) layers.
+            mtp_labels, _ = roll_tensor(mtp_labels, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params)
+            loss_mask, num_tokens = roll_tensor(
+                loss_mask, shifts=-1, dims=-1, cp_group=self.cp_group,  packed_seq_params=packed_seq_params
+            )
+            mtp_loss = compute_language_model_loss(mtp_labels, mtp_logits)
+            mtp_loss = loss_mask * mtp_loss
+            if is_training:
+                # TODO(shifangx): remove the use of parallel_state here
+                # after moving loss logging to loss_func in pretrain_gpt.py
+                MTPLossLoggingHelper.save_loss_to_tracker(
+                    torch.sum(mtp_loss) / num_tokens,
+                    mtp_layer_number,
+                    self.config.mtp_num_layers,
+                    avg_group=parallel_state.get_data_parallel_group(
+                        with_context_parallel=True
+                    ),
+                )
+            mtp_loss_scale = self.config.mtp_loss_scaling_factor / self.config.mtp_num_layers
+            if self.config.calculate_per_token_loss:
+                hidden_states = MTPLossAutoScaler.apply(
+                    hidden_states, mtp_loss_scale * mtp_loss
+                )
+            else:
+                hidden_states = MTPLossAutoScaler.apply(
+                    hidden_states, mtp_loss_scale * mtp_loss / num_tokens
+                )
+        
+        return hidden_states
+
     def sharded_state_dict(
         self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None
     ) -> ShardedStateDict:
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index d4802bfcdae..dd3e01aee7e 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -529,6 +529,9 @@ def validate_args(args, defaults={}):
                 args.global_batch_size), flush=True)
     assert args.global_batch_size > 0
 
+    if args.mtp_hybrid_override_pattern is not None:
+        assert args.mtp_spec is not None, "--mtp-spec must be set for mtp hybrid MTP models"
+    
     # Uneven virtual pipeline parallelism
     assert (
         int(args.num_layers_per_virtual_pipeline_stage is not None)
@@ -600,6 +603,18 @@ def validate_args(args, defaults={}):
                 assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \
                     'Number of layers should be divisible by the pipeline-model-parallel size'
 
+    # Validate MTP hybrid pattern segment count
+    if args.mtp_hybrid_override_pattern is not None:
+        mtp_segments = len(args.mtp_hybrid_override_pattern.split('|'))
+        expected_segments = args.pipeline_model_parallel_size * (
+            args.virtual_pipeline_model_parallel_size or 1
+        )
+        assert mtp_segments == expected_segments, (
+            f'MTP hybrid pattern has {mtp_segments} segments but expected {expected_segments} '
+            f'(PP={args.pipeline_model_parallel_size} * VP={args.virtual_pipeline_model_parallel_size or 1}). '
+            f'Use: {"|||" * (expected_segments - 1)}*E'
+        )
+    
     if args.virtual_pipeline_model_parallel_size is not None:
         if args.overlap_p2p_comm:
             assert args.pipeline_model_parallel_size > 1, \
@@ -3371,6 +3386,19 @@ def _add_experimental_args(parser):
                             'hybrid ratio arguments, then the number of each type'
                             'of layer in the override pattern must match number in'
                             'the overidden pattern')
+    group.add_argument("--mtp-hybrid-override-pattern", type=str, default=None,
+                       help='Force a specific hybrid layer pattern for MTP layers. '
+                       'If a value greater than 0.0 is supplied to any of the hybrid ratio '
+                       'arguments, then the number of each type of layer in the '
+                       'override pattern must match number in the overidden '
+                       'pattern')
+    group.add_argument('--mtp-spec', type=str, default=None, nargs='*',
+                       help='Specify the <module_location function_name> pair '
+                       'that returns a spec for mtp layer to customize a model,' 
+                       'transformer block, or transformer layer, depending on '
+                       'the use case. To use local spec specify local as the argument.'
+                       'For more details, see the model class, '
+                       'TransformerBlock class, or TransformerLayer class.')
     group.add_argument('--mamba-state-dim', type=int, default=128,
                        help='State dimension for Mamba layers.')
     group.add_argument('--mamba-head-dim', type=int, default=64,
diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index 325158780a4..779f27fd581 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -1347,6 +1347,7 @@ def _set_arg(arg_name, old_arg_name=None, force=False):
     _set_arg('hidden_dropout', force=True)
 
     _set_arg('hybrid_override_pattern', force=True)
+    _set_arg('mtp_hybrid_override_pattern', force=True)
     _set_arg('spec', force=True)
     _set_arg('hybrid_attention_ratio', force=True)
     _set_arg('hybrid_mlp_ratio', force=True)
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 3f23ca59f4a..28e3d2b6e94 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -178,6 +178,11 @@ def calculate_layer_counts():
             for layer_type in args.hybrid_override_pattern:
                 if layer_type in counts:
                     counts[layer_type] += 1
+            if args.mtp_hybrid_override_pattern:
+                assert args.mtp_num_layers is not None
+                for layer_type in args.mtp_hybrid_override_pattern:
+                    if layer_type in counts:
+                        counts[layer_type] += args.mtp_num_layers
             return counts['*'], counts['M'], counts['-'], counts['E']
         else:
             num_attn_layers = round(args.num_layers * args.hybrid_attention_ratio)
@@ -254,8 +259,10 @@ def hybrid_flops(batch_size, seq_len, hidden_size,
                      mlp_expansion=4.0, swiglu=False,
                      moe_latent_size=None,
                      moe_ffn_hidden_size=2048, shared_expert_ffn_hidden_size=2048, num_experts_routed_to=1,
-                     vocab_size=256000):
+                     vocab_size=256000, mtp_num_layers=None):
         """Calculate total FLOPs for the hybrid model."""
+        if mtp_num_layers is None:
+            mtp_num_layers = 0
         flops_fwd = (
                 num_attn_layers * attn_layer_flops(batch_size, seq_len, hidden_size,
                                                    num_attn_heads, gqa, gqa_groups, kv_channels) +
@@ -267,7 +274,7 @@ def hybrid_flops(batch_size, seq_len, hidden_size,
                 num_moe_layers * moe_layer_flops(batch_size, seq_len, hidden_size, moe_ffn_hidden_size,
                                                  shared_expert_ffn_hidden_size, num_experts_routed_to,
                                                  moe_latent_size, swiglu) +
-                (2 * batch_size * seq_len * hidden_size * vocab_size)  # logits computation
+                (2 * batch_size * seq_len * hidden_size * vocab_size * (1 + mtp_num_layers))  # logits computation
         )
         return flops_fwd * 3
 
@@ -478,6 +485,7 @@ def transformer_flops():
                                            else args.moe_shared_expert_intermediate_size),
             num_experts_routed_to=args.moe_router_topk,
             vocab_size=args.padded_vocab_size,
+            mtp_num_layers=args.mtp_num_layers,
         )
     else:
         # Compute standard Transformer model FLOPs.
diff --git a/pretrain_mamba.py b/pretrain_mamba.py
index 45b646a6cc0..1eaade5faa3 100644
--- a/pretrain_mamba.py
+++ b/pretrain_mamba.py
@@ -131,7 +131,7 @@ def forward_step(data_iterator, model: MambaModel):
 
     with stimer:
         output_tensor = model(tokens, position_ids, attention_mask,
-                              labels=labels)
+                              labels=labels, loss_mask=loss_mask)
 
     # [ModelOpt]: model is needed to access ModelOpt distillation losses
     return output_tensor, partial(loss_func, loss_mask, model=model)
diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index 05fb2c4fe63..09c08a64f22 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -94,7 +94,7 @@ def test_constructor_local(self, tp):
             assert mtp.layers[i].hnorm.weight.shape[0] == config.hidden_size
             assert mtp.layers[i].eh_proj.weight.shape[0] == config.hidden_size / tp
             assert mtp.layers[i].eh_proj.weight.shape[1] == config.hidden_size * 2
-            assert mtp.layers[i].transformer_layer is not None
+            assert mtp.layers[i].mtp_model_layer is not None
         num_weights = sum([p.numel() for p in mtp.parameters()])
         if tp == 1:
             assert num_weights == 58560 * config.mtp_num_layers
@@ -120,7 +120,7 @@ def test_constructor_ues_te(self, tp, cp):
             assert mtp.layers[i].hnorm.weight.shape[0] == config.hidden_size
             assert mtp.layers[i].eh_proj.weight.shape[0] == config.hidden_size / tp
             assert mtp.layers[i].eh_proj.weight.shape[1] == config.hidden_size * 2
-            assert mtp.layers[i].transformer_layer is not None
+            assert mtp.layers[i].mtp_model_layer is not None
         num_weights = sum([p.numel() for p in mtp.parameters()])
         if tp == 1:
             assert num_weights == 58560 * config.mtp_num_layers

From 3b057cac3b50443e8da49f3e6206a8562b25e9f7 Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Mon, 1 Dec 2025 17:56:50 +0000
Subject: [PATCH 13/48] addressed comments

---
 .../common/language_module/language_module.py |   2 +-
 .../core/models/mamba/mamba_layer_specs.py    |   7 +-
 megatron/core/ssm/mamba_block.py              |   2 +-
 .../transformer/multi_token_prediction.py     |  11 +-
 .../test_multi_token_prediction.py            | 270 ++++++++++++++++++
 5 files changed, 286 insertions(+), 6 deletions(-)

diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
index 706f8429a29..6f145f7f32c 100644
--- a/megatron/core/models/common/language_module/language_module.py
+++ b/megatron/core/models/common/language_module/language_module.py
@@ -254,7 +254,7 @@ def setup_embeddings_and_output_layer(self) -> None:
             LanguageModule.embedding_warning_printed = True
 
     def shared_embedding_or_output_weight(self) -> Tensor:
-        """Gets the emedding weight or output logit weights when share embedding and output weights set to True
+        """Gets the embedding weight or output logit weights when share embedding and output weights set to True
           or when use Multi-Token Prediction (MTP).
 
         Returns:
diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index 286d72b36ff..0a6e39052c9 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -22,7 +22,6 @@
 from megatron.core.transformer.mlp import MLP, MLPSubmodules
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
-from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionBlockSubmodules
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.transformer.multi_token_prediction import (
     MultiTokenPredictionBlockSubmodules,
@@ -40,7 +39,6 @@
 )
 
 try:
-    from megatron.core.extensions.transformer_engine import TEFusedMLP, TENorm
     from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
 
     HAVE_TE = True
@@ -182,7 +180,10 @@ def get_mamba_mtp_block_spec(
     vp_stage: Optional[int] = None,
 ) -> MultiTokenPredictionBlockSubmodules:
     """Mamba Multi-Token Prediction (MTP) block spec."""
+    assert HAVE_KITCHEN, "Kitchen needs to be installed."
+    
     if use_transformer_engine:
+        assert HAVE_TE, "TransformerEngine needs to be installed."
         backend: BackendSpecProvider = (
             KitchenSpecProvider(fallback=TESpecProvider())
             if config.use_kitchen
@@ -221,7 +222,7 @@ def get_mamba_mtp_block_spec_for_backend(
     if len(mtp_layer_specs) > 0:
         assert (
             len(mtp_layer_specs) == config.mtp_num_layers
-        ), +f"currently all of the mtp layers must stage in the same pipeline stage."
+        ), f"currently all of the mtp layers must stage in the same pipeline stage."
         mtp_block_spec = MultiTokenPredictionBlockSubmodules(layer_specs=mtp_layer_specs)
     else:
         mtp_block_spec = None
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index 4e8bc27cc54..ed9ce49fae7 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -150,7 +150,7 @@ def __init__(
                 elif layer_type == LayerSymbols.MOE:
                     # Transformer layers apply their own pp_layer_offset
                     layer = build_module(
-                        submodules.moe_layer, config=self.config, layer_number=i + 1,
+                        submodules.moe_layer, config=self.config, layer_number=i+1,
                     )
                 else:
                     assert False, "unexpected layer_type"
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index f86aa9dd401..2373b694528 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -553,10 +553,19 @@ def __init__(
                     self_attention_spec = self_attention_spec.submodules.self_attention
                     attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
                     assert attn_mask_type in SUPPORTED_ATTN_MASK, (
-                        f"Multi-Token Prediction (MTP) is not jet supported with "
+                        f"Multi-Token Prediction (MTP) is not yet supported with "
                         + f"{attn_mask_type} attention mask type."
                         + f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
                     )
+        elif hasattr(self.submodules.mtp_model_layer.submodules, 'self_attention'):
+            self_attention_spec = self.submodules.mtp_model_layer.submodules.self_attention
+            if self_attention_spec is not None:
+                attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
+                assert attn_mask_type in SUPPORTED_ATTN_MASK, (
+                    f"Multi-Token Prediction (MTP) is not yet supported with "
+                    + f"{attn_mask_type} attention mask type."
+                    + f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
+                )
 
         self.enorm = build_module(
             self.submodules.enorm,
diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index 09c08a64f22..e818cde16c4 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -13,6 +13,11 @@
     get_gpt_mtp_block_spec,
 )
 from megatron.core.models.gpt.gpt_model import GPTModel
+from megatron.core.models.mamba.mamba_layer_specs import (
+    get_mamba_mtp_block_spec,
+    mamba_stack_spec,
+)
+from megatron.core.models.mamba.mamba_model import MambaModel
 from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.parallel_state import get_context_parallel_group
@@ -684,3 +689,268 @@ def log(self, metrics, iteration):
         assert torch.all(MTPLossLoggingHelper.tracker["values"] == 0)
         assert MTPLossLoggingHelper.tracker["reduce_group"] is None
         assert MTPLossLoggingHelper.tracker["avg_group"] is None
+
+
+class TestMultiTokenPredictionMamba:
+    """Test Multi-Token Prediction with Mamba hybrid models."""
+
+    def setup_method(self, method):
+        self.seq_length = 32
+        self.micro_batch_size = 2
+        os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1'
+
+    def teardown_method(self, method):
+        Utils.destroy_model_parallel()
+        destroy_global_vars()
+        destroy_num_microbatches_calculator()
+        MTPLossLoggingHelper.tracker = {}
+
+    def model_provider(
+        self,
+        pre_process=True,
+        post_process=True,
+        **config_kwargs,
+    ):
+        """Model provider for Mamba hybrid models with MTP."""
+        model_parallel_cuda_manual_seed(_SEED)
+        args = get_args()
+        config = core_transformer_config_from_args(args)
+        
+        mtp_block_spec = get_mamba_mtp_block_spec(
+            config=config, spec=mamba_stack_spec, use_transformer_engine=True
+        )
+        
+        model = MambaModel(
+            config=config,
+            mamba_stack_spec=mamba_stack_spec,
+            vocab_size=args.vocal_size,
+            max_sequence_length=args.max_position_embeddings,
+            pre_process=pre_process,
+            post_process=post_process,
+            hybrid_attention_ratio=args.hybrid_attention_ratio,
+            hybrid_mlp_ratio=args.hybrid_mlp_ratio,
+            hybrid_override_pattern=args.hybrid_override_pattern,
+            mtp_hybrid_override_pattern=args.mtp_hybrid_override_pattern,
+            fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
+            parallel_output=True,
+            share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
+            position_embedding_type=args.position_embedding_type,
+            rotary_percent=args.rotary_percent,
+            mtp_block_spec=mtp_block_spec,
+        )
+        return model
+
+    def create_test_args(
+        self, tp, cp, sequence_length, micro_batch_size, fp8=None, full_recompute=False
+    ):
+        destroy_global_vars()
+        destroy_num_microbatches_calculator()
+
+        sys.argv = ['test_multi_token_prediction_mamba.py']
+        args = parse_args()
+        args.num_layers = 4
+        args.mtp_num_layers = 2
+        args.mtp_loss_scaling_factor = 0.1
+        args.vocal_size = 128800
+        args.hidden_size = 128
+        args.num_attention_heads = 8
+        args.num_query_groups = 8
+        args.mamba_num_groups = 4  
+        args.max_position_embeddings = 256
+        args.micro_batch_size = micro_batch_size
+        args.create_attention_mask_in_dataloader = True
+        args.seq_length = sequence_length
+        args.tensor_model_parallel_size = tp
+        args.sequence_parallel = True if tp > 1 else False
+        args.context_parallel_size = cp
+        args.position_embedding_type = 'rope'
+        args.train_iters = 1
+        args.ckpt_format = 'torch_dist'
+        args.lr = 3e-5
+        args.attention_dropout = 0.0
+        args.hidden_dropout = 0.0
+        args.async_tensor_model_parallel_allreduce = False
+        args.no_save_optim = True
+        args.no_load_optim = True
+        args.no_load_rng = True
+        args.bf16 = True
+        
+        args.hybrid_attention_ratio = 0.5  
+        args.hybrid_mlp_ratio = 0.0  
+        args.hybrid_override_pattern = "M*M*"  
+        args.mtp_hybrid_override_pattern = "M*"  
+        args.spec = "megatron.core.models.mamba.mamba_layer_specs.mamba_stack_spec"
+        args.mtp_spec = "megatron.core.models.mamba.mamba_layer_specs.mamba_stack_spec"
+        
+        if fp8 is not None:
+            args.fp8 = 'e4m3'
+        if full_recompute:
+            args.recompute_granularity = 'full'
+            args.recompute_method = 'uniform'
+            args.recompute_num_layers = 1
+        else:
+            args.recompute_granularity = None
+        args.add_bias_linear = False
+        args.swiglu = True
+
+        validate_args(args)
+        set_global_variables(args, False)
+        return args
+
+    def get_batch(self, seq_length, micro_batch_size):
+        data = list(range(seq_length))
+        input_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        labels = 1 + torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        position_ids = torch.tensor(data, dtype=torch.int64).repeat((micro_batch_size, 1)).cuda()
+        attention_mask = torch.ones(
+            (micro_batch_size, 1, seq_length, seq_length), dtype=bool
+        ).cuda()
+        loss_mask = torch.ones(seq_length).repeat((micro_batch_size, 1)).cuda()
+        batch = {
+            'tokens': input_ids,
+            'labels': labels,
+            'loss_mask': loss_mask,
+            'attention_mask': attention_mask,
+            'position_ids': position_ids,
+        }
+        return batch
+
+    @pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available")
+    @pytest.mark.parametrize(("tp", "cp"), [(1, 1), (2, 1)])
+    def test_sharded_state_dict_mamba(self, tp, cp):
+        """Test MTP with Mamba hybrid model - sharded state dict."""
+        args = self.create_test_args(tp, cp, self.seq_length, self.micro_batch_size)
+        set_args(args)
+        torch.manual_seed(_SEED)
+        Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+        mamba_model = get_model(self.model_provider, ModelType.encoder_or_decoder)
+        mamba_model = unwrap_model(mamba_model)
+        sharded_state_dict = mamba_model[0].sharded_state_dict()
+        
+        # Verify MTP layers are in the state dict
+        for i in range(args.mtp_num_layers):
+            assert f"mtp.layers.{i}.enorm.weight" in sharded_state_dict.keys()
+            assert f"mtp.layers.{i}.hnorm.weight" in sharded_state_dict.keys()
+            assert f"mtp.layers.{i}.eh_proj.weight" in sharded_state_dict.keys()
+
+    @pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available")
+    @pytest.mark.parametrize(("tp", "cp"), [(1, 1), (2, 1)])
+    def test_forward_backward_mamba(self, tmp_path_dist_ckpt, tp, cp):
+        """Test MTP forward and backward with Mamba hybrid model."""
+        tp_ref = 1
+        cp_ref = 1
+        args = self.create_test_args(tp_ref, cp_ref, self.seq_length, self.micro_batch_size)
+        set_args(args)
+        torch.manual_seed(_SEED)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp_ref, context_parallel_size=cp_ref
+        )
+        batch = self.get_batch(self.seq_length, self.micro_batch_size)
+        tokens, labels, loss_mask, attention_mask, position_ids = batch.values()
+        
+        mamba_model_ref, optimizer, opt_param_scheduler = setup_model_and_optimizer(
+            self.model_provider, ModelType.encoder_or_decoder
+        )
+        
+        output_ref = mamba_model_ref[0].forward(
+            input_ids=tokens,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            loss_mask=loss_mask,
+        )
+        
+        tracker = MTPLossLoggingHelper.tracker
+        mtp_loss_ref = None
+        assert "values" in tracker
+        mtp_loss_ref = tracker['values'].clone()
+        MTPLossLoggingHelper.clean_loss_in_tracker()
+
+        iteration = 123
+        num_floating_point_operations_so_far = 456
+
+        def set_ckpt_path(ckpt_path):
+            args.save = ckpt_path
+            args.load = ckpt_path
+
+        with TempNamedDir(
+            tmp_path_dist_ckpt / 'test_mtp_mamba_model_reconfiguration'
+        ) as ckpt_dir:
+            set_ckpt_path(ckpt_dir)
+            save_checkpoint(
+                iteration,
+                mamba_model_ref,
+                optimizer,
+                opt_param_scheduler,
+                num_floating_point_operations_so_far,
+            )
+
+            expected_ckpt_path = args.save / "iter_0000123" / ".metadata"
+            assert os.path.exists(expected_ckpt_path)
+
+            Utils.destroy_model_parallel()
+            args = self.create_test_args(tp, cp, self.seq_length, self.micro_batch_size)
+            set_args(args)
+            set_ckpt_path(ckpt_dir)
+            torch.manual_seed(_SEED)
+            Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+            mamba_model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
+                self.model_provider, ModelType.encoder_or_decoder
+            )
+            load_checkpoint(mamba_model, optimizer, opt_param_scheduler, strict=False)
+            
+            batch["output_ref"] = output_ref
+            batch = get_batch_on_this_cp_rank(batch)
+            tokens, labels, loss_mask, attention_mask, position_ids, output_ref = batch.values()
+            
+            output = mamba_model[0].forward(
+                input_ids=tokens,
+                position_ids=position_ids,
+                attention_mask=attention_mask,
+                labels=labels,
+                loss_mask=loss_mask,
+            )
+            
+            tracker = MTPLossLoggingHelper.tracker
+            assert "values" in tracker
+            mtp_loss = tracker['values'].clone()
+            pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['cp'])
+            torch.distributed.all_reduce(
+                mtp_loss, group=pg_collection.cp, op=torch.distributed.ReduceOp.AVG
+            )
+            MTPLossLoggingHelper.clean_loss_in_tracker()
+            
+            assert torch.allclose(output_ref, output, rtol=1e-03, atol=1e-03)
+            assert torch.allclose(mtp_loss, mtp_loss_ref, rtol=1e-02, atol=1e-02)
+
+            assert output.shape[0] == self.micro_batch_size
+            assert output.shape[1] == self.seq_length / cp
+
+            loss = output.mean()
+            loss.backward()
+            for name, param in mamba_model[0].named_parameters():
+                assert param.main_grad is not None
+
+    @pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available")
+    def test_attention_mask_validation_mamba(self):
+        """Test that attention mask type validation works for Mamba hybrid models.
+        """
+        tp = 1
+        cp = 1
+        args = self.create_test_args(tp, cp, self.seq_length, self.micro_batch_size)
+        set_args(args)
+        torch.manual_seed(_SEED)
+        Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
+        
+        try:
+            mamba_model = get_model(self.model_provider, ModelType.encoder_or_decoder)
+            mamba_model = unwrap_model(mamba_model)
+            assert isinstance(mamba_model[0], MambaModel)
+            assert mamba_model[0].mtp is not None
+        except AssertionError as e:
+            if "Multi-Token Prediction (MTP) is not yet supported" in str(e):
+                pytest.fail(
+                    f"Attention mask validation failed for Mamba hybrid model: {e}"
+                )
+            else:
+                raise

From 7649f6e361b6d69203b458cc0ee742cd4ad90d5a Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Wed, 3 Dec 2025 12:54:46 +0000
Subject: [PATCH 14/48] fixed the duplicate

---
 megatron/core/models/mamba/mamba_model.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 6521100648b..023d337e3f2 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -399,11 +399,7 @@ def forward(
                 compute_language_model_loss=self.compute_language_model_loss,
             )
 
-        # logits and loss
-        output_weight = None
-        if self.share_embeddings_and_output_weights:
-            output_weight = self.shared_embedding_or_output_weight()
-
+        
         sequence_parallel_override = False
         if in_inference_mode and inference_context.materialize_only_last_token_logits:
             if inference_context.is_static_batching():

From 07d7e6cb182235b969e31d01218087330d066d1b Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Fri, 5 Dec 2025 11:56:07 +0000
Subject: [PATCH 15/48] added repeated mtp layers

---
 .../core/models/mamba/mamba_layer_specs.py    | 11 +++++--
 megatron/core/ssm/mamba_block.py              | 12 +++++++
 .../transformer/multi_token_prediction.py     | 31 ++++++++++++++-----
 .../core/transformer/transformer_config.py    |  3 ++
 megatron/training/arguments.py                |  4 +++
 5 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index 0a6e39052c9..ced9d66ae83 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -214,15 +214,20 @@ def get_mamba_mtp_block_spec_for_backend(
         mtp_model_layer_spec=spec, backend=backend
     )
     mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0
-    mtp_layer_specs = [mtp_layer_spec] * mtp_num_layers
+    
+    if config.mtp_use_repeated_layer:
+        mtp_layer_specs = [mtp_layer_spec]
+    else:
+        mtp_layer_specs = [mtp_layer_spec] * mtp_num_layers
 
     offset = get_mtp_layer_offset(config)
     # split the mtp layer specs to only include the layers that are built in this pipeline stage.
     mtp_layer_specs = mtp_layer_specs[offset : offset + num_layers_to_build]
     if len(mtp_layer_specs) > 0:
+        expected_num_layers = 1 if config.mtp_use_repeated_layer else config.mtp_num_layers
         assert (
-            len(mtp_layer_specs) == config.mtp_num_layers
-        ), f"currently all of the mtp layers must stage in the same pipeline stage."
+            len(mtp_layer_specs) == expected_num_layers
+        ), f"currently all of the mtp layers must stage in the same pipeline stage. Expected {expected_num_layers}, got {len(mtp_layer_specs)}"
         mtp_block_spec = MultiTokenPredictionBlockSubmodules(layer_specs=mtp_layer_specs)
     else:
         mtp_block_spec = None
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index ed9ce49fae7..6752ed7da09 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -27,6 +27,9 @@
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.transformer.transformer_layer import TransformerLayer
 from megatron.core.transformer.utils import sharded_state_dict_default
+
+from megatron.core.ssm.mlp_layer import MLPLayer
+
 from megatron.core.utils import WrappedTensor, deprecate_inference_params, make_viewless_tensor
 
 
@@ -290,6 +293,15 @@ def forward(
                             rotary_pos_emb=rotary_pos_emb,
                             sequence_len_offset=sequence_len_offset,
                         )
+                    elif isinstance(layer, MLPLayer):
+                        # MLPLayer (standalone MLP without attention)
+                        hidden_states = layer(
+                            hidden_states=hidden_states,
+                            attention_mask=attention_mask,
+                            inference_context=inference_context,
+                            rotary_pos_emb=rotary_pos_emb,
+                            sequence_len_offset=sequence_len_offset,
+                        )
                     else:  # MambaLayer
                         hidden_states = layer(
                             hidden_states=hidden_states,
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 2373b694528..d0a32acfb60 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -982,6 +982,9 @@ class MultiTokenPredictionBlock(MegatronModule):
     the linear projection. The combined serves as the input of the Transformer block at
     the k-th depth to produce the output representation.
 
+    When `mtp_use_repeated_layer=True` in config, instead of creating N separate MTP layers,
+    only 1 layer is created and applied mtp_num_layers times. 
+
     for more information, please refer to DeepSeek-V3 Technical Report
     https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf
     """
@@ -999,6 +1002,8 @@ def __init__(
         self.submodules = _get_mtp_block_submodules(config, spec)
         self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor
         self.vp_stage = vp_stage
+        
+        self.mtp_use_repeated_layer = self.config.mtp_use_repeated_layer
 
         # Initialize Context Parallelism (CP) support for MTP
         # This enables MTP to work with CP > 1 by providing the CP process group
@@ -1030,12 +1035,21 @@ def build_layer(layer_spec, layer_number):
                 )
             return module
 
-        self.layers = torch.nn.ModuleList(
-            [
-                build_layer(layer_spec, i + 1)
-                for i, layer_spec in enumerate(self.submodules.layer_specs)
-            ]
-        )
+        if self.mtp_use_repeated_layer:
+            assert len(self.submodules.layer_specs) == 1, (
+                f"Repeated MTP mode requires exactly 1 layer spec, got {len(self.submodules.layer_specs)}. "
+                f"The layer will be applied {self.config.mtp_num_layers} times."
+            )
+            self.layers = torch.nn.ModuleList([
+                build_layer(self.submodules.layer_specs[0], layer_number=1)
+            ])
+        else:
+            self.layers = torch.nn.ModuleList(
+                [
+                    build_layer(layer_spec, i + 1)
+                    for i, layer_spec in enumerate(self.submodules.layer_specs)
+                ]
+            )
 
     def forward(
         self,
@@ -1071,8 +1085,9 @@ def forward(
         offset = get_mtp_layer_offset(self.config)
         hidden_states_list = list(torch.chunk(hidden_states, 1 + offset, dim=0))
         hidden_states = hidden_states_list[offset]
-        for layer_number in range(len(self.layers)):
-            (hidden_states, input_ids, position_ids) = self.layers[layer_number](
+        for iteration in range(self.config.mtp_num_layers):
+            layer_idx = 0 if self.mtp_use_repeated_layer else iteration
+            (hidden_states, input_ids, position_ids) = self.layers[layer_idx](
                 input_ids=input_ids,
                 position_ids=position_ids,
                 hidden_states=hidden_states,
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 17063217058..71885a8a382 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -51,6 +51,9 @@ class TransformerConfig(ModelParallelConfig):
     mtp_loss_scaling_factor: Optional[float] = None
     """Weighting factor of Multi-Token Prediction (MTP) loss."""
 
+    mtp_use_repeated_layer: bool = False
+    """Use a single MTP layer repeatedly instead of multiple separate layers."""
+
     num_layers_in_first_pipeline_stage: Optional[int] = None
     """Number of transformer layers on first pipeline stage.
     None implies equal layer division across PP ranks."""
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index dd3e01aee7e..4ca15e6bdc8 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1799,6 +1799,10 @@ def _add_network_size_args(parser):
                        'which serves as an additional training objective.')
     group.add_argument('--moe-latent-size', type=int, default=None,
                        help='Latent projection dimension for MoE. If None, MoE latent projections are not used.')
+    group.add_argument('--mtp-use-repeated-layer', action='store_true',
+                       help='Use a single MTP layer repeatedly instead of multiple separate layers. '
+                       'This is more parameter-efficient. When enabled, only 1 MTP layer is created '
+                       'and applied --mtp-num-layers times.')
     return parser
 
 

From bd1ac40fbaa025706a9d9a424cb2e5913dfc4a55 Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Fri, 5 Dec 2025 13:37:24 +0000
Subject: [PATCH 16/48] adds mtp-num-per-layer to fix a bug

---
 megatron/core/transformer/multi_token_prediction.py | 5 ++++-
 megatron/core/transformer/transformer_config.py     | 3 +++
 megatron/training/arguments.py                      | 2 ++
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index d0a32acfb60..0f37c62299f 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -601,8 +601,11 @@ def __init__(
         if self.config.mtp_num_layers is not None:
             if self.mtp_hybrid_override_pattern is not None:
                 pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+
                 # We do not need pre and post process stage for MTP layer, given they are handled in the 
                 # MultiTokenPredictionLayer itself.
+                assert self.config.mtp_num_layers_per_layer is not None, \
+                    "mtp_num_layers_per_layer must be set when using mtp_hybrid_override_pattern"
                 self.mtp_model_layer = build_module(
                     self.submodules.mtp_model_layer,
                     self.config,
@@ -612,7 +615,7 @@ def __init__(
                     dtype=self.config.params_dtype,
                     pg_collection=pg_collection,
                     vp_stage=self.vp_stage,
-                    num_layers=self.config.mtp_num_layers
+                    num_layers=self.config.mtp_num_layers_per_layer
                 )
             else:
                 # Uses the transformer block spec for MTP layer. This option is only implemented for the 
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 71885a8a382..7c0450bddfe 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -54,6 +54,9 @@ class TransformerConfig(ModelParallelConfig):
     mtp_use_repeated_layer: bool = False
     """Use a single MTP layer repeatedly instead of multiple separate layers."""
 
+    mtp_num_layers_per_layer: Optional[int] = None
+    """Number of layers inside each MTP layer (e.g., for hybrid models)."""
+
     num_layers_in_first_pipeline_stage: Optional[int] = None
     """Number of transformer layers on first pipeline stage.
     None implies equal layer division across PP ranks."""
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 4ca15e6bdc8..943e56fbefa 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1803,6 +1803,8 @@ def _add_network_size_args(parser):
                        help='Use a single MTP layer repeatedly instead of multiple separate layers. '
                        'This is more parameter-efficient. When enabled, only 1 MTP layer is created '
                        'and applied --mtp-num-layers times.')
+    group.add_argument('--mtp-num-layers-per-layer', type=int, default=None,
+                       help='Number of layers inside each MTP layer (e.g., for hybrid models).')
     return parser
 
 

From cbd19f1651c420aa011f3385d3361a8be5bf5b56 Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Fri, 5 Dec 2025 14:32:39 +0000
Subject: [PATCH 17/48] added the assert

---
 megatron/core/pipeline_parallel/schedules.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index 3c1fb2e3b18..c8c520b2e51 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -193,6 +193,9 @@ def set_current_microbatch(model, microbatch_id):
             layer.current_microbatch = microbatch_id
         if hasattr(model_with_decoder, 'mtp'):
             for layer in model_with_decoder.mtp.layers:
+                assert hasattr(layer, 'mtp_model_layer'), (
+                    f"MTP layer {layer} must have 'mtp_model_layer' attribute"
+                )
                 layer.mtp_model_layer.current_microbatch = microbatch_id
 
 

From d1ba2289b8bae2424f59e6c8f45735d5622f944d Mon Sep 17 00:00:00 2001
From: Rabeeh Karimi Mahabadi <rkarimimahab@nvidia.com>
Date: Tue, 9 Dec 2025 13:36:54 +0000
Subject: [PATCH 18/48] added the assert for vp-stage of MTP layers

---
 megatron/core/transformer/multi_token_prediction.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 0f37c62299f..f1028799fe4 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -1008,6 +1008,13 @@ def __init__(
         
         self.mtp_use_repeated_layer = self.config.mtp_use_repeated_layer
 
+        vp_size = config.virtual_pipeline_model_parallel_size
+        assert is_vp_last_stage(vp_stage=vp_stage, vp_size=vp_size), (
+            f"MTP layers must be placed on the last virtual pipeline stage. "
+            f"Got vp_stage={vp_stage} with vp_size={vp_size}. "
+            f"Placing MTP layers on different VPP stages is not currently supported."
+        )
+
         # Initialize Context Parallelism (CP) support for MTP
         # This enables MTP to work with CP > 1 by providing the CP process group
         # to the roll_tensor function for proper boundary communication

From 4d0cd4d34a18463d410cb3b497b2b1cc355c237f Mon Sep 17 00:00:00 2001
From: Rabeeh Mahabadi <rkarimimahab@nb-hel-cs-001-vscode-02.cm.cluster>
Date: Thu, 18 Dec 2025 09:57:24 -0800
Subject: [PATCH 19/48] added the LBL fixes

---
 .../core/models/mamba/mamba_layer_specs.py    |  6 +--
 megatron/core/ssm/mamba_block.py              |  9 ++++-
 megatron/core/transformer/moe/moe_layer.py    |  8 +++-
 megatron/core/transformer/moe/router.py       | 37 ++++++++++++++++---
 .../transformer/multi_token_prediction.py     |  7 ++--
 .../core/transformer/transformer_layer.py     |  5 +++
 megatron/training/arguments.py                |  7 +++-
 7 files changed, 63 insertions(+), 16 deletions(-)

diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index ced9d66ae83..431a20d2174 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -111,7 +111,7 @@
     ),
 )
 
-<<<<<<< HEAD
+
 mamba_inference_stack_spec = ModuleSpec(
     module=MambaStack,
     submodules=MambaStackSubmodules(
@@ -171,7 +171,7 @@
         ),
     ),
 )
-=======
+
 
 def get_mamba_mtp_block_spec(
     config: TransformerConfig,
@@ -233,5 +233,3 @@ def get_mamba_mtp_block_spec_for_backend(
         mtp_block_spec = None
 
     return mtp_block_spec
-
->>>>>>> a05cbed31 (added mtp support)
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index 6752ed7da09..ce7cd3a2bcf 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -88,6 +88,7 @@ def __init__(
         pg_collection: ProcessGroupCollection = None,
         vp_stage: Optional[int] = None,
         num_layers: int = None,
+        is_mtp_layer: bool = False,
     ) -> None:
         super().__init__(config=config)
         self.residual_in_fp32 = residual_in_fp32
@@ -95,6 +96,7 @@ def __init__(
         self.post_layer_norm = post_layer_norm
         self.post_process = post_process
         self.vp_stage = vp_stage
+        self.is_mtp_layer = is_mtp_layer
 
         assert pg_collection is not None, "pg_collection must be provided for MambaStack"
 
@@ -141,6 +143,7 @@ def __init__(
                         config=self.config,
                         layer_number=i + 1,
                         pg_collection=pg_collection,
+                        is_mtp_layer=is_mtp_layer,
                     )
                 elif layer_type == LayerSymbols.MLP:
                     # Transformer layers apply their own pp_layer_offset
@@ -153,7 +156,11 @@ def __init__(
                 elif layer_type == LayerSymbols.MOE:
                     # Transformer layers apply their own pp_layer_offset
                     layer = build_module(
-                        submodules.moe_layer, config=self.config, layer_number=i+1,
+                        submodules.moe_layer,
+                        config=self.config,
+                        layer_number=i+1,
+                        pg_collection=pg_collection,
+                        is_mtp_layer=is_mtp_layer,
                     )
                 else:
                     assert False, "unexpected layer_type"
diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index 0b0ee9b6850..e06b9faf7da 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -50,10 +50,12 @@ def __init__(
         config: TransformerConfig,
         layer_number: Optional[int] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
+        is_mtp_layer: bool = False,
     ):
         super(BaseMoELayer, self).__init__(config)
         self.config = config
         self.layer_number = layer_number
+        self.is_mtp_layer = is_mtp_layer
         self.ep_group = pg_collection.ep
         # use pg_collection.expt_tp_group as tensor parallel group in this module.
         self.attn_tp_group = pg_collection.tp
@@ -103,6 +105,7 @@ def __init__(
         submodules: Optional[MoESubmodules] = None,
         layer_number: Optional[int] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
+        is_mtp_layer: bool = False,
     ):
         self.submodules = submodules
         # TODO(Hepteract): delete the usage of the global parallel_state.
@@ -110,7 +113,8 @@ def __init__(
         if pg_collection is None:
             pg_collection = get_default_pg_collection()
         super(MoELayer, self).__init__(
-            config=config, layer_number=layer_number, pg_collection=pg_collection
+            config=config, layer_number=layer_number, pg_collection=pg_collection,
+            is_mtp_layer=is_mtp_layer
         )
         self.moe_layer_recompute = (
             config.recompute_granularity == 'selective' and "moe" in config.recompute_modules
@@ -123,7 +127,7 @@ def __init__(
         self.tp_group = pg_collection.tp
 
         # Initialize router.
-        self.router = TopKRouter(config=self.config, pg_collection=pg_collection)
+        self.router = TopKRouter(config=self.config, pg_collection=pg_collection, is_mtp_layer=is_mtp_layer)
 
         # Initialize latent projections.
         if self.config.moe_latent_size:
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 34d81a21ffa..dc94d394656 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -28,7 +28,8 @@ class Router(ABC, MegatronModule):
     """Base Router class"""
 
     def __init__(
-        self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None
+        self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None,
+        is_mtp_layer: bool = False
     ) -> None:
         """
         Initialize the Router module.
@@ -36,12 +37,14 @@ def __init__(
         Args:
             config (TransformerConfig): Configuration object for the Transformer model.
             pg_collection (ProcessGroupCollection, optional): Process groups for MoE operations.
+            is_mtp_layer (bool): Flag indicating if this router is part of an MTP layer.
         """
         super().__init__(config)
         self.config = config
         self.num_experts = self.config.num_moe_experts
         self.moe_aux_loss_func = None
         self.layer_number = None
+        self.is_mtp_layer = is_mtp_layer
         self.tp_group = pg_collection.tp
         self.cp_group = pg_collection.cp
         self.tp_cp_group = pg_collection.tp_cp
@@ -144,15 +147,17 @@ class TopKRouter(Router):
     """
 
     def __init__(
-        self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None
+        self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None,
+        is_mtp_layer: bool = False
     ) -> None:
         """Initialize the zero token dropping router.
 
         Args:
             config (TransformerConfig): The configuration for the transformer model.
             pg_collection (ProcessGroupCollection, optional): Process groups for MoE operations.
+            is_mtp_layer (bool): Flag indicating if this router is part of an MTP layer.
         """
-        super().__init__(config=config, pg_collection=pg_collection)
+        super().__init__(config=config, pg_collection=pg_collection, is_mtp_layer=is_mtp_layer)
         self.topk = self.config.moe_router_topk
         self.routing_type = self.config.moe_router_load_balancing_type
         self.score_function = self.config.moe_router_score_function
@@ -399,6 +404,11 @@ def attach_and_log_load_balancing_loss(
                 Set this to True if the reduce group has data parallel ranks. This flag is used to
                 ensure the correct reduction in aux loss tracking.
         """
+        # When using repeated MTP layers, the loss is counted "mtp_num_layers" times. To avoid accumulating the
+        # load balancing loss multiple times, we scale it by 1/mtp_num_layers so the total loss is correct.
+        if self.is_mtp_layer and self.config.mtp_use_repeated_layer and self.config.mtp_num_layers is not None:
+            aux_loss = aux_loss / self.config.mtp_num_layers
+
         # TODO (zijiey): fix the per_layer_logging for MTP, currently it will incorrectly
         # add the aux loss logging value to other layer's since it is difficult to get the
         # correct layer_number for MTP. It does not affect the correctness of the calculation
@@ -406,10 +416,16 @@ def attach_and_log_load_balancing_loss(
         num_layers = self.config.num_layers
         if self.config.mtp_num_layers is not None:
             num_layers += self.config.mtp_num_layers
+
+        if self.is_mtp_layer:
+            layer_number = self.layer_number + self.config.num_layers
+        else:
+            layer_number = self.layer_number
+
         save_to_aux_losses_tracker(
             aux_loss_name,
             aux_loss / aux_loss_coeff,
-            self.layer_number,
+            layer_number,
             num_layers,
             reduce_group=reduce_group,
             reduce_group_has_dp=reduce_group_has_dp,
@@ -454,11 +470,22 @@ def apply_z_loss(self, logits):
             else:
                 logits = MoEAuxLossAutoScaler.apply(logits, z_loss)
 
+            # When using repeated MTP layers, the same MTP layer is called mtp_num_layers times. To avoid
+            # accumulating the z_loss multiple times, we scale it by 1/mtp_num_layers so the total loss is correct.
+            if self.is_mtp_layer and self.config.mtp_use_repeated_layer and self.config.mtp_num_layers is not None:
+                z_loss = z_loss / self.config.mtp_num_layers
+
             num_layers = self.config.num_layers
             if self.config.mtp_num_layers is not None:
                 num_layers += self.config.mtp_num_layers
+
+            if self.is_mtp_layer:
+                layer_number = self.layer_number + self.config.num_layers
+            else:
+                layer_number = self.layer_number
+
             save_to_aux_losses_tracker(
-                "z_loss", z_loss / moe_z_loss_coeff, self.layer_number, num_layers
+                "z_loss", z_loss / moe_z_loss_coeff, layer_number, num_layers
             )
         return logits
 
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index f1028799fe4..a9342c9c039 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -615,14 +615,16 @@ def __init__(
                     dtype=self.config.params_dtype,
                     pg_collection=pg_collection,
                     vp_stage=self.vp_stage,
+                    num_layers=self.config.mtp_num_layers_per_layer,
+                    is_mtp_layer=True
                     num_layers=self.config.mtp_num_layers_per_layer
                 )
             else:
                 # Uses the transformer block spec for MTP layer. This option is only implemented for the 
                 # GPT model. In hybrid model, we model transformer block spec for MTP layers with the hybrid
                 # override pattern.
-                self.mtp_model_layer = build_module(self.submodules.mtp_model_layer, config=self.config, vp_stage=self.vp_stage)
-
+                self.mtp_model_layer = build_module(self.submodules.mtp_model_layer, config=self.config, vp_stage=self.vp_stage, is_mtp_layer=True)
+                
         self.final_layernorm = build_module(
             self.submodules.layer_norm,
             config=self.config,
@@ -861,7 +863,6 @@ def forward(
             [s, b, h], and optionally the updated context tensor if cross-attention is used.
         """
         assert context is None, f"multi token prediction + cross attention is not yet supported."
-
         input_ids, position_ids, decoder_input, hidden_states = self._get_embeddings(
             input_ids=input_ids,
             position_ids=position_ids,
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index 24ab98786b6..1335084ad7f 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -266,6 +266,7 @@ def __init__(
         hidden_dropout: Optional[float] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
         vp_stage: Optional[int] = None,
+        is_mtp_layer: bool = False,
     ):
         super().__init__(config=config, vp_stage=vp_stage)
 
@@ -279,6 +280,7 @@ def __init__(
             self.config, vp_stage, get_pg_rank(pg_collection.pp)
         )
         self.hidden_dropout = config.hidden_dropout if hidden_dropout is None else hidden_dropout
+        self.is_mtp_layer = is_mtp_layer
 
         # [Module 1: Input Layernorm] Optional Layernorm on the input data
         # TODO: add pytorch only layernorm
@@ -349,6 +351,9 @@ def __init__(
         if isinstance(submodules.mlp, ModuleSpec):
             if submodules.mlp.module in (MoELayer, GroupedMLP, TEGroupedMLP, SequentialMLP):
                 additional_mlp_kwargs["pg_collection"] = pg_collection
+                # Pass is_mtp_layer flag to MoELayer to distinguish MTP MoE layers.
+                if submodules.mlp.module == MoELayer:
+                    additional_mlp_kwargs["is_mtp_layer"] = self.is_mtp_layer
             elif submodules.mlp.module == MLP:
                 assert hasattr(
                     pg_collection, 'tp'
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 943e56fbefa..eb68c075a8f 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -602,7 +602,6 @@ def validate_args(args, defaults={}):
 
                 assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \
                     'Number of layers should be divisible by the pipeline-model-parallel size'
-
     # Validate MTP hybrid pattern segment count
     if args.mtp_hybrid_override_pattern is not None:
         mtp_segments = len(args.mtp_hybrid_override_pattern.split('|'))
@@ -1797,6 +1796,12 @@ def _add_network_size_args(parser):
                        'We compute the average of the MTP losses across all depths, '
                        'and multiply it the scaling factor to obtain the overall MTP loss, '
                        'which serves as an additional training objective.')
+    group.add_argument('--mtp-use-repeated-layer', action='store_true',
+                       help='Use a single MTP layer repeatedly instead of multiple separate layers. '
+                       'This is more parameter-efficient. When enabled, only 1 MTP layer is created '
+                       'and applied --mtp-num-layers times.')
+    group.add_argument('--mtp-num-layers-per-layer', type=int, default=None,
+                       help='Number of layers inside each MTP layer (e.g., for hybrid models).')
     group.add_argument('--moe-latent-size', type=int, default=None,
                        help='Latent projection dimension for MoE. If None, MoE latent projections are not used.')
     group.add_argument('--mtp-use-repeated-layer', action='store_true',

From 305ba71f258e0ebc6f7ffada0411cda10f0d0c13 Mon Sep 17 00:00:00 2001
From: Rabeeh Mahabadi <rkarimimahab@nb-hel-cs-001-vscode-02.cm.cluster>
Date: Thu, 18 Dec 2025 10:40:37 -0800
Subject: [PATCH 20/48] added some fixes

---
 megatron/core/models/gpt/gpt_model.py               | 1 -
 megatron/core/transformer/multi_token_prediction.py | 1 -
 2 files changed, 2 deletions(-)

diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index eab4c258427..aa468b0ed9b 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -533,7 +533,6 @@ def _postprocess(
                 packed_seq_params=packed_seq_params,
                 sequence_len_offset=sequence_len_offset,
                 embedding=self.embedding,
-                packed_seq_params=packed_seq_params,
                 **(extra_block_kwargs or {}),
             )
 
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index a9342c9c039..ff43ba71f93 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -617,7 +617,6 @@ def __init__(
                     vp_stage=self.vp_stage,
                     num_layers=self.config.mtp_num_layers_per_layer,
                     is_mtp_layer=True
-                    num_layers=self.config.mtp_num_layers_per_layer
                 )
             else:
                 # Uses the transformer block spec for MTP layer. This option is only implemented for the 

From c60715aabb2eaea4a87d0bcf758974a5c7a233a8 Mon Sep 17 00:00:00 2001
From: Rabeeh Mahabadi <rkarimimahab@nb-hel-cs-001-vscode-02.cm.cluster>
Date: Thu, 18 Dec 2025 10:45:12 -0800
Subject: [PATCH 21/48] fixed more issues

---
 megatron/training/arguments.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index eb68c075a8f..d4a3337d423 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1804,12 +1804,6 @@ def _add_network_size_args(parser):
                        help='Number of layers inside each MTP layer (e.g., for hybrid models).')
     group.add_argument('--moe-latent-size', type=int, default=None,
                        help='Latent projection dimension for MoE. If None, MoE latent projections are not used.')
-    group.add_argument('--mtp-use-repeated-layer', action='store_true',
-                       help='Use a single MTP layer repeatedly instead of multiple separate layers. '
-                       'This is more parameter-efficient. When enabled, only 1 MTP layer is created '
-                       'and applied --mtp-num-layers times.')
-    group.add_argument('--mtp-num-layers-per-layer', type=int, default=None,
-                       help='Number of layers inside each MTP layer (e.g., for hybrid models).')
     return parser
 
 

From cb77b0ef5663702121176e0fdaf081bd59461192 Mon Sep 17 00:00:00 2001
From: Rabeeh Mahabadi <rkarimimahab@nb-hel-cs-001-vscode-02.cm.cluster>
Date: Fri, 19 Dec 2025 04:55:22 -0800
Subject: [PATCH 22/48] updated

---
 tests/unit_tests/transformer/test_multi_token_prediction.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index e818cde16c4..7b49e4fd852 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -779,6 +779,7 @@ def create_test_args(
         args.hybrid_mlp_ratio = 0.0  
         args.hybrid_override_pattern = "M*M*"  
         args.mtp_hybrid_override_pattern = "M*"  
+        args.mtp_num_layers_per_layer = 2 
         args.spec = "megatron.core.models.mamba.mamba_layer_specs.mamba_stack_spec"
         args.mtp_spec = "megatron.core.models.mamba.mamba_layer_specs.mamba_stack_spec"
         

From d732548e39e23c40a945064dfe5ebfbdcb0ded84 Mon Sep 17 00:00:00 2001
From: Rabeeh Mahabadi <rkarimimahab@nb-hel-cs-001-vscode-02.cm.cluster>
Date: Mon, 22 Dec 2025 01:30:10 -0800
Subject: [PATCH 23/48] fixed some lintingand small bugs

---
 megatron/core/models/gpt/gpt_model.py         |  2 +-
 .../core/models/mamba/mamba_layer_specs.py    | 11 +++--
 megatron/core/ssm/mamba_block.py              |  4 +-
 megatron/core/transformer/moe/router.py       |  1 -
 .../transformer/multi_token_prediction.py     | 46 +++++++++++--------
 .../test_multi_token_prediction.py            | 42 ++++++++---------
 6 files changed, 56 insertions(+), 50 deletions(-)

diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index aa468b0ed9b..465103aafbf 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -620,7 +620,7 @@ def shared_embedding_or_output_weight(self) -> Tensor:
             # will be stored in embedding layer, and output layer will not have any weight.
             assert hasattr(
                 self, 'embedding'
-            ), f"embedding is needed in this pipeline stage, but it is not initialized."
+            ), "embedding is needed in this pipeline stage, but it is not initialized."
             return self.embedding.word_embeddings.weight
         elif self.post_process:
             return self.output_layer.weight
diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index 431a20d2174..6fb520f5e58 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -176,12 +176,12 @@
 def get_mamba_mtp_block_spec(
     config: TransformerConfig,
     spec: ModuleSpec,
-    use_transformer_engine: bool, 
+    use_transformer_engine: bool,
     vp_stage: Optional[int] = None,
 ) -> MultiTokenPredictionBlockSubmodules:
     """Mamba Multi-Token Prediction (MTP) block spec."""
     assert HAVE_KITCHEN, "Kitchen needs to be installed."
-    
+
     if use_transformer_engine:
         assert HAVE_TE, "TransformerEngine needs to be installed."
         backend: BackendSpecProvider = (
@@ -198,10 +198,11 @@ def get_mamba_mtp_block_spec(
     return get_mamba_mtp_block_spec_for_backend(
         config=config, spec=spec, backend=backend, vp_stage=vp_stage
     )
-    
+
+
 def get_mamba_mtp_block_spec_for_backend(
     config: TransformerConfig,
-    spec: ModuleSpec, 
+    spec: ModuleSpec,
     backend: BackendSpecProvider,
     vp_stage: Optional[int] = None,
 ) -> MultiTokenPredictionBlockSubmodules:
@@ -214,7 +215,7 @@ def get_mamba_mtp_block_spec_for_backend(
         mtp_model_layer_spec=spec, backend=backend
     )
     mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0
-    
+
     if config.mtp_use_repeated_layer:
         mtp_layer_specs = [mtp_layer_spec]
     else:
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index ce7cd3a2bcf..19481c61c82 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -158,7 +158,7 @@ def __init__(
                     layer = build_module(
                         submodules.moe_layer,
                         config=self.config,
-                        layer_number=i+1,
+                        layer_number=i + 1,
                         pg_collection=pg_collection,
                         is_mtp_layer=is_mtp_layer,
                     )
@@ -328,7 +328,7 @@ def forward(
 
         # Ensure that the tensor passed between pipeline parallel stages is
         # viewless. See related notes in TransformerBlock and TransformerLayer
-        output = make_viewless_tensor(
+        hidden_states = make_viewless_tensor(
             inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True
         )
 
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index dc94d394656..5448543c84f 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -457,7 +457,6 @@ def apply_z_loss(self, logits):
             # Skip Z loss calculations when using torch.no_grad() or checkpointing.
             moe_z_loss_coeff = self.config.moe_z_loss_coeff / self.tp_cp_group.size()
             z_loss = z_loss_func(logits, moe_z_loss_coeff)
-            scale_up = 1.0
             if self.calculate_per_token_loss:
                 # The expected final scaling for z_loss gradients is
                 # 1/(num_micro_batches * dp_size).
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index ff43ba71f93..8f46a25fbea 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
+import warnings
 from contextlib import nullcontext
 from dataclasses import dataclass
 from typing import Callable, List, Optional, Union
@@ -368,7 +369,7 @@ def track_mtp_metrics(loss_scale, iteration, writer, wandb_writer=None, total_lo
         mtp_losses = tracker["values"] * loss_scale
         mtp_num_layers = mtp_losses.shape[0]
         for i in range(mtp_num_layers):
-            name = f"mtp_{i+1} loss"
+            name = f"mtp_{i + 1} loss"
             loss = mtp_losses[i]
             if total_loss_dict is not None:
                 if name in total_loss_dict:
@@ -421,7 +422,7 @@ def get_mtp_layer_spec(
 
 
 def get_mtp_layer_spec_for_backend(
-   mtp_model_layer_spec: ModuleSpec, backend: BackendSpecProvider
+    mtp_model_layer_spec: ModuleSpec, backend: BackendSpecProvider
 ) -> ModuleSpec:
     """Get the MTP layer spec.
 
@@ -550,21 +551,21 @@ def __init__(
         if hasattr(self.submodules.mtp_model_layer.submodules, 'attention_layer'):
             self_attention_spec = self.submodules.mtp_model_layer.submodules.attention_layer
             if self_attention_spec.submodules.self_attention is not None:
-                    self_attention_spec = self_attention_spec.submodules.self_attention
-                    attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
-                    assert attn_mask_type in SUPPORTED_ATTN_MASK, (
-                        f"Multi-Token Prediction (MTP) is not yet supported with "
-                        + f"{attn_mask_type} attention mask type."
-                        + f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
-                    )
+                self_attention_spec = self_attention_spec.submodules.self_attention
+                attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
+                assert attn_mask_type in SUPPORTED_ATTN_MASK, (
+                    f"Multi-Token Prediction (MTP) is not yet supported with "
+                    f"{attn_mask_type} attention mask type. "
+                    f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
+                )
         elif hasattr(self.submodules.mtp_model_layer.submodules, 'self_attention'):
             self_attention_spec = self.submodules.mtp_model_layer.submodules.self_attention
             if self_attention_spec is not None:
                 attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
                 assert attn_mask_type in SUPPORTED_ATTN_MASK, (
                     f"Multi-Token Prediction (MTP) is not yet supported with "
-                    + f"{attn_mask_type} attention mask type."
-                    + f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
+                    f"{attn_mask_type} attention mask type. "
+                    f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
                 )
 
         self.enorm = build_module(
@@ -602,14 +603,14 @@ def __init__(
             if self.mtp_hybrid_override_pattern is not None:
                 pg_collection = ProcessGroupCollection.use_mpu_process_groups()
 
-                # We do not need pre and post process stage for MTP layer, given they are handled in the 
-                # MultiTokenPredictionLayer itself.
+                # We do not need pre and post process stage for MTP layer, given they are
+                # handled in the MultiTokenPredictionLayer itself.
                 assert self.config.mtp_num_layers_per_layer is not None, \
                     "mtp_num_layers_per_layer must be set when using mtp_hybrid_override_pattern"
                 self.mtp_model_layer = build_module(
                     self.submodules.mtp_model_layer,
                     self.config,
-                    pre_process=False, 
+                    pre_process=False,
                     post_process=False,
                     hybrid_override_pattern=self.mtp_hybrid_override_pattern,
                     dtype=self.config.params_dtype,
@@ -619,11 +620,16 @@ def __init__(
                     is_mtp_layer=True
                 )
             else:
-                # Uses the transformer block spec for MTP layer. This option is only implemented for the 
-                # GPT model. In hybrid model, we model transformer block spec for MTP layers with the hybrid
-                # override pattern.
-                self.mtp_model_layer = build_module(self.submodules.mtp_model_layer, config=self.config, vp_stage=self.vp_stage, is_mtp_layer=True)
-                
+                # Uses the transformer block spec for MTP layer. This option is only
+                # implemented for the GPT model. In hybrid model, we model transformer
+                # block spec for MTP layers with the hybrid override pattern.
+                self.mtp_model_layer = build_module(
+                    self.submodules.mtp_model_layer,
+                    config=self.config,
+                    vp_stage=self.vp_stage,
+                    is_mtp_layer=True
+                )
+
         self.final_layernorm = build_module(
             self.submodules.layer_norm,
             config=self.config,
@@ -861,7 +867,7 @@ def forward(
             Union[Tensor, Tuple[Tensor, Tensor]]: The output hidden states tensor of shape
             [s, b, h], and optionally the updated context tensor if cross-attention is used.
         """
-        assert context is None, f"multi token prediction + cross attention is not yet supported."
+        assert context is None, "multi token prediction + cross attention is not yet supported."
         input_ids, position_ids, decoder_input, hidden_states = self._get_embeddings(
             input_ids=input_ids,
             position_ids=position_ids,
diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index 7b49e4fd852..a1e8c68ea1f 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -682,8 +682,8 @@ def log(self, metrics, iteration):
 
         # Verify total_loss_dict is populated
         for i in range(num_layers):
-            assert f"mtp_{i+1} loss" in total_loss_dict
-            assert total_loss_dict[f"mtp_{i+1} loss"] == loss * loss_scale
+            assert f"mtp_{i + 1} loss" in total_loss_dict
+            assert total_loss_dict[f"mtp_{i + 1} loss"] == loss * loss_scale
 
         # Verify tracker is cleaned
         assert torch.all(MTPLossLoggingHelper.tracker["values"] == 0)
@@ -715,11 +715,11 @@ def model_provider(
         model_parallel_cuda_manual_seed(_SEED)
         args = get_args()
         config = core_transformer_config_from_args(args)
-        
+
         mtp_block_spec = get_mamba_mtp_block_spec(
             config=config, spec=mamba_stack_spec, use_transformer_engine=True
         )
-        
+
         model = MambaModel(
             config=config,
             mamba_stack_spec=mamba_stack_spec,
@@ -755,7 +755,7 @@ def create_test_args(
         args.hidden_size = 128
         args.num_attention_heads = 8
         args.num_query_groups = 8
-        args.mamba_num_groups = 4  
+        args.mamba_num_groups = 4
         args.max_position_embeddings = 256
         args.micro_batch_size = micro_batch_size
         args.create_attention_mask_in_dataloader = True
@@ -774,15 +774,15 @@ def create_test_args(
         args.no_load_optim = True
         args.no_load_rng = True
         args.bf16 = True
-        
-        args.hybrid_attention_ratio = 0.5  
-        args.hybrid_mlp_ratio = 0.0  
-        args.hybrid_override_pattern = "M*M*"  
-        args.mtp_hybrid_override_pattern = "M*"  
-        args.mtp_num_layers_per_layer = 2 
+
+        args.hybrid_attention_ratio = 0.5
+        args.hybrid_mlp_ratio = 0.0
+        args.hybrid_override_pattern = "M*M*"
+        args.mtp_hybrid_override_pattern = "M*"
+        args.mtp_num_layers_per_layer = 2
         args.spec = "megatron.core.models.mamba.mamba_layer_specs.mamba_stack_spec"
         args.mtp_spec = "megatron.core.models.mamba.mamba_layer_specs.mamba_stack_spec"
-        
+
         if fp8 is not None:
             args.fp8 = 'e4m3'
         if full_recompute:
@@ -827,7 +827,7 @@ def test_sharded_state_dict_mamba(self, tp, cp):
         mamba_model = get_model(self.model_provider, ModelType.encoder_or_decoder)
         mamba_model = unwrap_model(mamba_model)
         sharded_state_dict = mamba_model[0].sharded_state_dict()
-        
+
         # Verify MTP layers are in the state dict
         for i in range(args.mtp_num_layers):
             assert f"mtp.layers.{i}.enorm.weight" in sharded_state_dict.keys()
@@ -848,11 +848,11 @@ def test_forward_backward_mamba(self, tmp_path_dist_ckpt, tp, cp):
         )
         batch = self.get_batch(self.seq_length, self.micro_batch_size)
         tokens, labels, loss_mask, attention_mask, position_ids = batch.values()
-        
+
         mamba_model_ref, optimizer, opt_param_scheduler = setup_model_and_optimizer(
             self.model_provider, ModelType.encoder_or_decoder
         )
-        
+
         output_ref = mamba_model_ref[0].forward(
             input_ids=tokens,
             position_ids=position_ids,
@@ -860,7 +860,7 @@ def test_forward_backward_mamba(self, tmp_path_dist_ckpt, tp, cp):
             labels=labels,
             loss_mask=loss_mask,
         )
-        
+
         tracker = MTPLossLoggingHelper.tracker
         mtp_loss_ref = None
         assert "values" in tracker
@@ -899,11 +899,11 @@ def set_ckpt_path(ckpt_path):
                 self.model_provider, ModelType.encoder_or_decoder
             )
             load_checkpoint(mamba_model, optimizer, opt_param_scheduler, strict=False)
-            
+
             batch["output_ref"] = output_ref
             batch = get_batch_on_this_cp_rank(batch)
             tokens, labels, loss_mask, attention_mask, position_ids, output_ref = batch.values()
-            
+
             output = mamba_model[0].forward(
                 input_ids=tokens,
                 position_ids=position_ids,
@@ -911,7 +911,7 @@ def set_ckpt_path(ckpt_path):
                 labels=labels,
                 loss_mask=loss_mask,
             )
-            
+
             tracker = MTPLossLoggingHelper.tracker
             assert "values" in tracker
             mtp_loss = tracker['values'].clone()
@@ -920,7 +920,7 @@ def set_ckpt_path(ckpt_path):
                 mtp_loss, group=pg_collection.cp, op=torch.distributed.ReduceOp.AVG
             )
             MTPLossLoggingHelper.clean_loss_in_tracker()
-            
+
             assert torch.allclose(output_ref, output, rtol=1e-03, atol=1e-03)
             assert torch.allclose(mtp_loss, mtp_loss_ref, rtol=1e-02, atol=1e-02)
 
@@ -942,7 +942,7 @@ def test_attention_mask_validation_mamba(self):
         set_args(args)
         torch.manual_seed(_SEED)
         Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp)
-        
+
         try:
             mamba_model = get_model(self.model_provider, ModelType.encoder_or_decoder)
             mamba_model = unwrap_model(mamba_model)

From 9d07525a47222f20f60aa7d868555cd9e352f3d1 Mon Sep 17 00:00:00 2001
From: Rabeeh Mahabadi <rkarimimahab@nb-hel-cs-001-vscode-02.cm.cluster>
Date: Mon, 22 Dec 2025 01:44:29 -0800
Subject: [PATCH 24/48] fixed linting errors

---
 .../core/models/mamba/mamba_layer_specs.py    |  1 +
 .../transformer/multi_token_prediction.py     | 22 +++++++++----------
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index fb73698babd..0cb23257f9d 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -198,6 +198,7 @@ def get_mamba_mtp_block_spec(
         config=config, spec=spec, backend=backend, vp_stage=vp_stage
     )
 
+
 def get_mamba_mtp_block_spec_for_backend(
     config: TransformerConfig,
     spec: ModuleSpec,
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 96b20626fa8..720aa9031bf 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -629,7 +629,7 @@ def __init__(
                     vp_stage=self.vp_stage,
                     is_mtp_layer=True
                 )
-                
+
         self.final_layernorm = build_module(
             self.submodules.layer_norm,
             config=self.config,
@@ -992,7 +992,7 @@ class MultiTokenPredictionBlock(MegatronModule):
     the k-th depth to produce the output representation.
 
     When `mtp_use_repeated_layer=True` in config, instead of creating N separate MTP layers,
-    only 1 layer is created and applied mtp_num_layers times. 
+    only 1 layer is created and applied mtp_num_layers times.
 
     for more information, please refer to DeepSeek-V3 Technical Report
     https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf
@@ -1011,7 +1011,7 @@ def __init__(
         self.submodules = _get_mtp_block_submodules(config, spec)
         self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor
         self.vp_stage = vp_stage
-        
+
         self.mtp_use_repeated_layer = self.config.mtp_use_repeated_layer
 
         vp_size = config.virtual_pipeline_model_parallel_size
@@ -1139,11 +1139,11 @@ def process_loss(
         packed_seq_params: Optional[PackedSeqParams] = None,
     ) -> Tensor:
         """Process Multi-Token Prediction (MTP) loss computation.
-        
+
         This method handles the MTP loss computation for multiple prediction layers.
         It chunks the hidden states, computes logits and losses for each MTP layer,
         and applies loss scaling.
-        
+
         Args:
             hidden_states (Tensor): Hidden states tensor from the model.
             labels (Tensor): Ground truth labels.
@@ -1153,18 +1153,18 @@ def process_loss(
             runtime_gather_output (Optional[bool]): Whether to gather output at runtime.
             is_training (bool): Whether the model is in training mode.
             compute_language_model_loss (Callable): Method to compute language model loss.
-        
+
         Returns:
             Tensor: Updated hidden states after MTP loss processing.
         """
         mtp_labels = labels.clone()
         hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0)
         hidden_states = hidden_states_list[0]
-        
+
         if loss_mask is None:
             # if loss_mask is not provided, use all ones as loss_mask
             loss_mask = torch.ones_like(mtp_labels)
-        
+
         for mtp_layer_number in range(self.config.mtp_num_layers):
             # output
             mtp_logits, _ = output_layer(
@@ -1175,7 +1175,7 @@ def process_loss(
             # Calc loss for the current Multi-Token Prediction (MTP) layers.
             mtp_labels, _ = roll_tensor(mtp_labels, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params)
             loss_mask, num_tokens = roll_tensor(
-                loss_mask, shifts=-1, dims=-1, cp_group=self.cp_group,  packed_seq_params=packed_seq_params
+                loss_mask, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params
             )
             mtp_loss = compute_language_model_loss(mtp_labels, mtp_logits)
             mtp_loss = loss_mask * mtp_loss
@@ -1199,7 +1199,7 @@ def process_loss(
                 hidden_states = MTPLossAutoScaler.apply(
                     hidden_states, mtp_loss_scale * mtp_loss / num_tokens
                 )
-        
+
         return hidden_states
 
     def sharded_state_dict(
@@ -1221,7 +1221,7 @@ def sharded_state_dict(
         layer_prefix = f'{prefix}layers.'
         for layer in self.layers:
             offset = get_mtp_layer_offset(self.config)
-            sharded_prefix = f'{layer_prefix}{layer.layer_number - 1 }.'
+            sharded_prefix = f'{layer_prefix}{layer.layer_number - 1}.'
 
             state_dict_prefix = f'{layer_prefix}{layer.layer_number - 1 - offset}.'
             sharded_pp_offset = []

From b61e281ca74e21ee8b55ca32f95f00786fb45468 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Mon, 12 Jan 2026 17:10:48 -0800
Subject: [PATCH 25/48] Add layer parsing for MTP

---
 .../core/ssm/mamba_hybrid_layer_allocation.py | 124 +++++++++++++++-
 .../ssm/test_mamba_hybrid_layer_allocation.py | 134 +++++++++++++++++-
 2 files changed, 256 insertions(+), 2 deletions(-)

diff --git a/megatron/core/ssm/mamba_hybrid_layer_allocation.py b/megatron/core/ssm/mamba_hybrid_layer_allocation.py
index fe997e2249a..10aea459d36 100644
--- a/megatron/core/ssm/mamba_hybrid_layer_allocation.py
+++ b/megatron/core/ssm/mamba_hybrid_layer_allocation.py
@@ -1,7 +1,8 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 
 import logging
-from typing import Dict, List, Tuple
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
 
 if __name__ != "__main__":
     from megatron.core.utils import log_single_rank
@@ -29,9 +30,130 @@ class Symbols:
     ATTENTION = "*"
     MLP = "-"
     MOE = 'E'
+    MTP_SEPARATOR = "/"
     VALID = {MAMBA, ATTENTION, MLP, MOE}
 
 
+@dataclass
+class ParsedHybridPattern:
+    """Result of parsing a unified hybrid pattern string.
+
+    A unified pattern encodes both the main decoder pattern and the MTP pattern
+    in a single string using "/" as a separator.
+
+    Format: "<main_pattern>/<mtp_pattern>/<mtp_pattern>/..."
+
+    Examples:
+        - "M*M*" -> main="M*M*", mtp=None, depths=0 (no MTP)
+        - "M*M*/MM/MM" -> main="M*M*", mtp="MM", depths=2
+        - "MMMM/*M/*M/*M" -> main="MMMM", mtp="*M", depths=3
+
+    The "/" symbol introduces MTP patterns. Each repeated pattern after the main
+    decoder represents one MTP prediction depth.
+
+    Attributes:
+        main_pattern: The main decoder layer pattern (e.g., "M*M*")
+        mtp_pattern: The MTP layer pattern per depth (e.g., "MM"), or None if no MTP
+        mtp_num_depths: Number of MTP prediction depths (0 if no MTP)
+    """
+
+    main_pattern: Optional[str]
+    mtp_pattern: Optional[str]
+    mtp_num_depths: int
+
+
+def parse_hybrid_pattern(pattern: Optional[str]) -> ParsedHybridPattern:
+    """Parse a unified hybrid pattern string into main and MTP components.
+
+    The pattern uses "/" as a separator between the main decoder pattern and
+    MTP patterns. Each MTP pattern after the separator represents one prediction
+    depth.
+
+    Format: "<main_pattern>/<mtp_pattern>/<mtp_pattern>/..."
+
+    Args:
+        pattern: Unified pattern string, e.g., "M*M*/MM/MM" or just "M*M*"
+
+    Returns:
+        ParsedHybridPattern with main_pattern, mtp_pattern, and mtp_num_depths
+
+    Raises:
+        ValueError: If MTP patterns are inconsistent (all must be identical)
+        ValueError: If pattern contains invalid layer symbols
+
+    Examples:
+        >>> parse_hybrid_pattern("M*M*")
+        ParsedHybridPattern(main_pattern="M*M*", mtp_pattern=None, mtp_num_depths=0)
+
+        >>> parse_hybrid_pattern("M*M*/MM/MM")
+        ParsedHybridPattern(main_pattern="M*M*", mtp_pattern="MM", mtp_num_depths=2)
+
+        >>> parse_hybrid_pattern("MMMM/*M/*M/*M")
+        ParsedHybridPattern(main_pattern="MMMM", mtp_pattern="*M", mtp_num_depths=3)
+    """
+    if pattern is None:
+        return ParsedHybridPattern(main_pattern=None, mtp_pattern=None, mtp_num_depths=0)
+
+    # Split by MTP separator
+    parts = pattern.split(Symbols.MTP_SEPARATOR)
+
+    if len(parts) == 1:
+        # No MTP separator found - pattern is main decoder only
+        main_pattern = parts[0]
+        _validate_pattern(main_pattern, "main")
+        return ParsedHybridPattern(main_pattern=main_pattern, mtp_pattern=None, mtp_num_depths=0)
+
+    # First part is main decoder pattern
+    main_pattern = parts[0]
+    if main_pattern:
+        _validate_pattern(main_pattern, "main")
+
+    # Remaining parts are MTP patterns (one per depth)
+    mtp_parts = parts[1:]
+
+    if not mtp_parts or all(p == "" for p in mtp_parts):
+        # No MTP patterns after separator
+        return ParsedHybridPattern(
+            main_pattern=main_pattern if main_pattern else None, mtp_pattern=None, mtp_num_depths=0
+        )
+
+    # Validate all MTP patterns are identical
+    mtp_pattern = mtp_parts[0]
+    for i, part in enumerate(mtp_parts[1:], start=2):
+        if part != mtp_pattern:
+            raise ValueError(
+                f"All MTP patterns must be identical. "
+                f"Pattern 1 is '{mtp_pattern}', but pattern {i} is '{part}'. "
+                f"Full pattern: '{pattern}'"
+            )
+
+    _validate_pattern(mtp_pattern, "MTP")
+
+    return ParsedHybridPattern(
+        main_pattern=main_pattern if main_pattern else None,
+        mtp_pattern=mtp_pattern,
+        mtp_num_depths=len(mtp_parts),
+    )
+
+
+def _validate_pattern(pattern: str, pattern_name: str) -> None:
+    """Validate that a pattern contains only valid layer symbols.
+
+    Args:
+        pattern: Layer pattern string to validate
+        pattern_name: Name of pattern for error messages (e.g., "main" or "MTP")
+
+    Raises:
+        ValueError: If pattern contains invalid symbols
+    """
+    for char in pattern:
+        if char not in Symbols.VALID:
+            raise ValueError(
+                f"In {pattern_name} pattern, '{char}' is not a valid layer symbol. "
+                f"Valid symbols are: {Symbols.VALID}"
+            )
+
+
 def _allocate_auto(
     total_layers_count: int, target_attention_ratio: float, target_mlp_ratio: float
 ) -> list:
diff --git a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
index 77d02c69607..7edb46a7e80 100644
--- a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
+++ b/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
@@ -6,7 +6,12 @@
 import pytest
 import torch
 
-from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols, allocate_layers
+from megatron.core.ssm.mamba_hybrid_layer_allocation import (
+    ParsedHybridPattern,
+    Symbols,
+    allocate_layers,
+    parse_hybrid_pattern,
+)
 
 
 @pytest.mark.internal
@@ -75,3 +80,130 @@ def test_wrong_length_override_pattern(self):
     def test_wrong_number_of_layer_types_in_override_pattern(self):
         # This override_pattern has too many mlps and not enough attention
         layer_types = allocate_layers(8, 0.5, 0.25, "M*--M**-")
+
+
+@pytest.mark.internal
+class TestParseHybridPattern:
+    """Tests for parse_hybrid_pattern with unified pattern syntax."""
+
+    def test_none_pattern(self):
+        """Test that None pattern returns all None values."""
+        result = parse_hybrid_pattern(None)
+        assert result.main_pattern is None
+        assert result.mtp_pattern is None
+        assert result.mtp_num_depths == 0
+
+    def test_main_pattern_only(self):
+        """Test patterns without MTP (no / separator)."""
+        test_cases = [
+            ("M*M*", "M*M*"),
+            ("MMMM", "MMMM"),
+            ("*M*M", "*M*M"),
+            ("MM-*", "MM-*"),
+            ("E", "E"),
+        ]
+        for pattern, expected_main in test_cases:
+            result = parse_hybrid_pattern(pattern)
+            assert result.main_pattern == expected_main, f"Failed for pattern: {pattern}"
+            assert result.mtp_pattern is None
+            assert result.mtp_num_depths == 0
+
+    def test_main_with_single_mtp_depth(self):
+        """Test patterns with 1 MTP depth."""
+        test_cases = [
+            ("M*M*/MM", "M*M*", "MM", 1),
+            ("MMMM/*M", "MMMM", "*M", 1),
+            ("M/M", "M", "M", 1),
+        ]
+        for pattern, expected_main, expected_mtp, expected_depths in test_cases:
+            result = parse_hybrid_pattern(pattern)
+            assert result.main_pattern == expected_main, f"Failed for pattern: {pattern}"
+            assert result.mtp_pattern == expected_mtp, f"Failed for pattern: {pattern}"
+            assert result.mtp_num_depths == expected_depths, f"Failed for pattern: {pattern}"
+
+    def test_main_with_multiple_mtp_depths(self):
+        """Test patterns with multiple MTP depths."""
+        test_cases = [
+            ("M*M*/MM/MM", "M*M*", "MM", 2),
+            ("M*M*/MM/MM/MM", "M*M*", "MM", 3),
+            ("MMMM/*M/*M/*M", "MMMM", "*M", 3),
+            ("M*/*/*/*", "M*", "*", 4),
+            ("M/M/M/M/M", "M", "M", 4),
+        ]
+        for pattern, expected_main, expected_mtp, expected_depths in test_cases:
+            result = parse_hybrid_pattern(pattern)
+            assert result.main_pattern == expected_main, f"Failed for pattern: {pattern}"
+            assert result.mtp_pattern == expected_mtp, f"Failed for pattern: {pattern}"
+            assert result.mtp_num_depths == expected_depths, f"Failed for pattern: {pattern}"
+
+    def test_mtp_patterns_must_be_identical(self):
+        """Test that mismatched MTP patterns raise ValueError."""
+        invalid_patterns = [
+            "M*M*/MM/M*",  # MM != M*
+            "M*M*/MM/MM/M",  # MM != M
+            "MMMM/*M/M*",  # *M != M*
+        ]
+        for pattern in invalid_patterns:
+            with pytest.raises(ValueError, match="All MTP patterns must be identical"):
+                parse_hybrid_pattern(pattern)
+
+    def test_invalid_symbols_in_main_pattern(self):
+        """Test that invalid symbols in main pattern raise ValueError."""
+        invalid_patterns = [
+            "M*X*",  # X is not valid
+            "MaMM",  # a is not valid
+            "M*M*1",  # 1 is not valid
+        ]
+        for pattern in invalid_patterns:
+            with pytest.raises(ValueError, match="not a valid layer symbol"):
+                parse_hybrid_pattern(pattern)
+
+    def test_invalid_symbols_in_mtp_pattern(self):
+        """Test that invalid symbols in MTP pattern raise ValueError."""
+        invalid_patterns = [
+            "M*M*/MX",  # X is not valid
+            "M*M*/MM/Ma",  # a is not valid
+        ]
+        for pattern in invalid_patterns:
+            with pytest.raises(ValueError, match="not a valid layer symbol"):
+                parse_hybrid_pattern(pattern)
+
+    def test_empty_main_pattern_with_mtp(self):
+        """Test pattern that starts with / (empty main pattern)."""
+        result = parse_hybrid_pattern("/MM/MM")
+        assert result.main_pattern is None
+        assert result.mtp_pattern == "MM"
+        assert result.mtp_num_depths == 2
+
+    def test_trailing_separator(self):
+        """Test patterns with trailing separator."""
+        # "M*M*/" means main="M*M*", one empty MTP pattern
+        result = parse_hybrid_pattern("M*M*/")
+        assert result.main_pattern == "M*M*"
+        # Empty string after separator means no valid MTP pattern
+        assert result.mtp_pattern is None
+        assert result.mtp_num_depths == 0
+
+    def test_complex_patterns(self):
+        """Test more complex realistic patterns."""
+        test_cases = [
+            # Main decoder with attention, MTP with mamba only
+            ("M*M*M*M*/MMM/MMM", "M*M*M*M*", "MMM", 2),
+            # Main decoder with MLP, MTP with attention+mamba
+            ("MM-MM-/*M/*M", "MM-MM-", "*M", 2),
+            # All attention main, mamba MTP
+            ("*****/M/M/M/M", "*****", "M", 4),
+            # MoE in main pattern
+            ("MEME/MM/MM", "MEME", "MM", 2),
+        ]
+        for pattern, expected_main, expected_mtp, expected_depths in test_cases:
+            result = parse_hybrid_pattern(pattern)
+            assert result.main_pattern == expected_main, f"Failed for pattern: {pattern}"
+            assert result.mtp_pattern == expected_mtp, f"Failed for pattern: {pattern}"
+            assert result.mtp_num_depths == expected_depths, f"Failed for pattern: {pattern}"
+
+    def test_dataclass_equality(self):
+        """Test that ParsedHybridPattern supports equality comparison."""
+        p1 = parse_hybrid_pattern("M*M*/MM/MM")
+        p2 = ParsedHybridPattern(main_pattern="M*M*", mtp_pattern="MM", mtp_num_depths=2)
+        assert p1 == p2

From c61f5bde01eacd24d9ad932b804a872080db6e45 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Mon, 12 Jan 2026 17:22:49 -0800
Subject: [PATCH 26/48] Shared Layer Builder + MTP Layer Construction

---
 .../core/models/mamba/mamba_layer_specs.py    | 128 ++++------
 megatron/core/ssm/layer_builder.py            | 139 +++++++++++
 megatron/core/ssm/mamba_block.py              |  65 ++---
 .../transformer/multi_token_prediction.py     | 225 +++++++++++++++---
 4 files changed, 403 insertions(+), 154 deletions(-)
 create mode 100644 megatron/core/ssm/layer_builder.py

diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index 0cb23257f9d..ef217cb963f 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -6,7 +6,6 @@
     TENorm,
     TERowParallelLinear,
 )
-from typing import Optional
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec
 from megatron.core.ssm.mamba_block import MambaStack, MambaStackSubmodules
@@ -22,14 +21,12 @@
 from megatron.core.transformer.mlp import MLP, MLPSubmodules
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
-from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.transformer.multi_token_prediction import (
+    MultiTokenPredictionBlock,
     MultiTokenPredictionBlockSubmodules,
-    get_mtp_layer_offset,
-    get_mtp_layer_spec_for_backend,
-    get_mtp_num_layers_to_build,
+    MultiTokenPredictionLayer,
+    MultiTokenPredictionLayerSubmodules,
 )
-from megatron.core.models.backends import BackendSpecProvider, LocalSpecProvider
 
 moe = get_moe_module_spec(
     use_te=True,
@@ -38,21 +35,6 @@
     moe_use_legacy_grouped_gemm=False,
 )
 
-try:
-    from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
-
-    HAVE_TE = True
-except ImportError:
-    HAVE_TE = False
-
-
-try:
-    from megatron.core.extensions.kitchen import KitchenSpecProvider
-
-    HAVE_KITCHEN = True
-except ImportError:
-    HAVE_KITCHEN = False
-
 mamba_stack_spec = ModuleSpec(
     module=MambaStack,
     submodules=MambaStackSubmodules(
@@ -173,63 +155,49 @@
 )
 
 
-def get_mamba_mtp_block_spec(
-    config: TransformerConfig,
-    spec: ModuleSpec,
-    use_transformer_engine: bool,
-    vp_stage: Optional[int] = None,
-) -> MultiTokenPredictionBlockSubmodules:
-    """Mamba Multi-Token Prediction (MTP) block spec."""
-    assert HAVE_KITCHEN, "Kitchen needs to be installed."
-    if use_transformer_engine:
-        assert HAVE_TE, "TransformerEngine needs to be installed."
-        backend: BackendSpecProvider = (
-            KitchenSpecProvider(fallback=TESpecProvider())
-            if config.use_kitchen
-            else TESpecProvider()
-        )
-    else:
-        backend = (
-            KitchenSpecProvider(fallback=LocalSpecProvider())
-            if config.use_kitchen
-            else LocalSpecProvider()
-        )
-    return get_mamba_mtp_block_spec_for_backend(
-        config=config, spec=spec, backend=backend, vp_stage=vp_stage
-    )
-
-
-def get_mamba_mtp_block_spec_for_backend(
-    config: TransformerConfig,
-    spec: ModuleSpec,
-    backend: BackendSpecProvider,
-    vp_stage: Optional[int] = None,
-) -> MultiTokenPredictionBlockSubmodules:
-    """Mamba Multi-Token Prediction (MTP) block spec."""
-    num_layers_to_build = get_mtp_num_layers_to_build(config, vp_stage=vp_stage)
-    if num_layers_to_build == 0:
-        return None
-
-    mtp_layer_spec = get_mtp_layer_spec_for_backend(
-        mtp_model_layer_spec=spec, backend=backend
+def get_mamba_mtp_block_spec(use_te: bool = True) -> ModuleSpec:
+    """MTP block spec for Mamba using unified pattern syntax.
+
+    This spec provides norms and projection only - inner layers are built
+    by MultiTokenPredictionLayer using the shared layer_builder with
+    mtp_layer_pattern and mamba_submodules passed from MambaModel.
+
+    The number of MTP depths is determined by the parsed unified pattern
+    (e.g., "M*M*/MM/MM" -> main="M*M*", mtp="MM", 2 depths).
+
+    Args:
+        use_te: Whether to use TransformerEngine modules (default: True)
+
+    Returns:
+        ModuleSpec for MultiTokenPredictionBlock
+
+    Example:
+        >>> mtp_spec = get_mamba_mtp_block_spec()
+        >>> mtp_block = MultiTokenPredictionBlock(
+        ...     config=config,
+        ...     spec=mtp_spec,
+        ...     mtp_layer_pattern="MM",
+        ...     mtp_num_depths=2,
+        ...     mamba_submodules=mamba_stack_spec.submodules,
+        ... )
+    """
+    norm = TENorm if use_te else TENorm  # Fallback to TENorm for now
+    linear = TELayerNormColumnParallelLinear if use_te else TELayerNormColumnParallelLinear
+
+    return ModuleSpec(
+        module=MultiTokenPredictionBlock,
+        submodules=MultiTokenPredictionBlockSubmodules(
+            layer_specs=[
+                ModuleSpec(
+                    module=MultiTokenPredictionLayer,
+                    submodules=MultiTokenPredictionLayerSubmodules(
+                        enorm=norm,
+                        hnorm=norm,
+                        eh_proj=linear,
+                        mtp_model_layer=None,  # Built via pattern + mamba_submodules
+                        layer_norm=norm,
+                    ),
+                )
+            ]
+        ),
     )
-    mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0
-
-    if config.mtp_use_repeated_layer:
-        mtp_layer_specs = [mtp_layer_spec]
-    else:
-        mtp_layer_specs = [mtp_layer_spec] * mtp_num_layers
-
-    offset = get_mtp_layer_offset(config)
-    # split the mtp layer specs to only include the layers that are built in this pipeline stage.
-    mtp_layer_specs = mtp_layer_specs[offset : offset + num_layers_to_build]
-    if len(mtp_layer_specs) > 0:
-        expected_num_layers = 1 if config.mtp_use_repeated_layer else config.mtp_num_layers
-        assert (
-            len(mtp_layer_specs) == expected_num_layers
-        ), f"currently all of the mtp layers must stage in the same pipeline stage. Expected {expected_num_layers}, got {len(mtp_layer_specs)}"
-        mtp_block_spec = MultiTokenPredictionBlockSubmodules(layer_specs=mtp_layer_specs)
-    else:
-        mtp_block_spec = None
-
-    return mtp_block_spec
diff --git a/megatron/core/ssm/layer_builder.py b/megatron/core/ssm/layer_builder.py
new file mode 100644
index 00000000000..52342d0ec9d
--- /dev/null
+++ b/megatron/core/ssm/layer_builder.py
@@ -0,0 +1,139 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+"""
+Shared layer-building utilities for Mamba hybrid architectures.
+
+Used by:
+- MambaStack: for main decoder layers
+- MultiTokenPredictionLayer: for MTP inner layers
+
+This module provides a single source of truth for converting layer pattern strings
+(e.g., "M*M*" or "MM") into actual layer modules.
+"""
+
+from typing import TYPE_CHECKING, List
+
+from torch import nn
+
+from megatron.core.fp8_utils import get_fp8_context
+from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols
+from megatron.core.transformer.spec_utils import build_module
+
+if TYPE_CHECKING:
+    from megatron.core.process_groups_config import ProcessGroupCollection
+    from megatron.core.ssm.mamba_block import MambaStackSubmodules
+    from megatron.core.transformer import TransformerConfig
+
+
+def build_layers_from_pattern(
+    pattern: str,
+    submodules: "MambaStackSubmodules",
+    config: "TransformerConfig",
+    pg_collection: "ProcessGroupCollection",
+    layer_offset: int = 0,
+    residual_in_fp32: bool = False,
+    is_mtp_layer: bool = False,
+) -> nn.ModuleList:
+    """
+    Build layers from a pattern string. Single source of truth for pattern → layers.
+
+    This function is shared between:
+    - MambaStack (for main decoder layers)
+    - MultiTokenPredictionLayer (for MTP inner layers)
+
+    Args:
+        pattern: Layer pattern string, e.g., "M*M*" or "MM"
+            - M = Mamba layer
+            - * = Attention (TransformerLayer)
+            - - = MLP layer
+            - E = MoE layer
+        submodules: MambaStackSubmodules containing specs for each layer type
+        config: TransformerConfig for the model
+        pg_collection: Process group collection for distributed training
+        layer_offset: Starting layer number offset (for layer_number assignment)
+        residual_in_fp32: Whether to use FP32 for residual connections (Mamba layers)
+        is_mtp_layer: Whether these layers are for MTP (affects attention layer building)
+
+    Returns:
+        nn.ModuleList of layers matching the pattern
+
+    Example:
+        >>> submodules = mamba_stack_spec.submodules
+        >>> layers = build_layers_from_pattern("M*M", submodules, config, pg_collection)
+        >>> # Returns [MambaLayer, TransformerLayer, MambaLayer]
+    """
+    layers = nn.ModuleList()
+
+    for i, layer_type in enumerate(pattern):
+        layer_number = i + 1 + layer_offset
+        fp8_init_context = get_fp8_context(config, i + layer_offset, is_init=True)
+
+        with fp8_init_context:
+            if layer_type == Symbols.MAMBA:
+                layer = build_module(
+                    submodules.mamba_layer,
+                    config=config,
+                    residual_in_fp32=residual_in_fp32,
+                    layer_number=layer_number,
+                    pp_layer_offset=layer_offset,
+                    pg_collection=pg_collection,
+                )
+            elif layer_type == Symbols.ATTENTION:
+                # Transformer layers apply their own pp_layer_offset
+                layer = build_module(
+                    submodules.attention_layer,
+                    config=config,
+                    layer_number=i + 1,
+                    pg_collection=pg_collection,
+                    is_mtp_layer=is_mtp_layer,
+                )
+            elif layer_type == Symbols.MLP:
+                # MLP layers apply their own pp_layer_offset
+                layer = build_module(
+                    submodules.mlp_layer,
+                    config=config,
+                    layer_number=i + 1,
+                    pg_collection=pg_collection,
+                )
+            elif layer_type == Symbols.MOE:
+                # MoE layers apply their own pp_layer_offset
+                layer = build_module(
+                    submodules.moe_layer,
+                    config=config,
+                    layer_number=i + 1,
+                    pg_collection=pg_collection,
+                    is_mtp_layer=is_mtp_layer,
+                )
+            else:
+                raise ValueError(
+                    f"Unknown layer type: '{layer_type}'. "
+                    f"Valid types are: {Symbols.VALID}"
+                )
+
+        layers.append(layer)
+
+    return layers
+
+
+def get_layer_types_from_pattern(pattern: str) -> List[str]:
+    """
+    Convert a pattern string to a list of layer type symbols.
+
+    This is a simple utility that validates and returns the pattern as a list.
+
+    Args:
+        pattern: Layer pattern string, e.g., "M*M*"
+
+    Returns:
+        List of layer type symbols, e.g., ["M", "*", "M", "*"]
+
+    Raises:
+        ValueError: If pattern contains invalid symbols
+    """
+    layer_types = list(pattern)
+    for char in layer_types:
+        if char not in Symbols.VALID:
+            raise ValueError(
+                f"Invalid layer type: '{char}'. Valid types are: {Symbols.VALID}"
+            )
+    return layer_types
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index 19481c61c82..f606fdfefd7 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -19,8 +19,9 @@
 from megatron.core.fp8_utils import get_fp8_context
 from megatron.core.inference.contexts import BaseInferenceContext
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.ssm.layer_builder import build_layers_from_pattern
 from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols as LayerSymbols
-from megatron.core.ssm.mamba_hybrid_layer_allocation import allocate_layers
+from megatron.core.ssm.mamba_hybrid_layer_allocation import allocate_layers, parse_hybrid_pattern
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.module import MegatronModule
@@ -109,12 +110,17 @@ def __init__(
         self.hybrid_attention_ratio = hybrid_attention_ratio
         self.hybrid_mlp_ratio = hybrid_mlp_ratio
         self.hybrid_override_pattern = hybrid_override_pattern
+        self.pg_collection = pg_collection
+
+        # Parse pattern - extract only main pattern, MTP handled by MambaModel
+        parsed = parse_hybrid_pattern(hybrid_override_pattern)
+        main_pattern = parsed.main_pattern  # MTP pattern ignored here
 
         self.layer_type_list = allocate_layers(
             self.config.num_layers if num_layers is None else num_layers,
             self.hybrid_attention_ratio,
             self.hybrid_mlp_ratio,
-            self.hybrid_override_pattern,
+            main_pattern,
         )
 
         pp_layer_offset = 0
@@ -122,49 +128,18 @@ def __init__(
             pp_layer_offset, self.layer_type_list = self._select_layers_for_pipeline_parallel(
                 self.layer_type_list
             )
-
-        self.layers = nn.ModuleList()
-        for i, layer_type in enumerate(self.layer_type_list):
-            fp8_init_context = get_fp8_context(self.config, i + pp_layer_offset, is_init=True)
-            with fp8_init_context:
-                if layer_type == LayerSymbols.MAMBA:
-                    layer = build_module(
-                        submodules.mamba_layer,
-                        config=self.config,
-                        residual_in_fp32=residual_in_fp32,
-                        layer_number=i + 1 + pp_layer_offset,
-                        pp_layer_offset=pp_layer_offset,
-                        pg_collection=pg_collection,
-                    )
-                elif layer_type == LayerSymbols.ATTENTION:
-                    # Transformer layers apply their own pp_layer_offset
-                    layer = build_module(
-                        submodules.attention_layer,
-                        config=self.config,
-                        layer_number=i + 1,
-                        pg_collection=pg_collection,
-                        is_mtp_layer=is_mtp_layer,
-                    )
-                elif layer_type == LayerSymbols.MLP:
-                    # Transformer layers apply their own pp_layer_offset
-                    layer = build_module(
-                        submodules.mlp_layer,
-                        config=self.config,
-                        layer_number=i + 1,
-                        pg_collection=pg_collection,
-                    )
-                elif layer_type == LayerSymbols.MOE:
-                    # Transformer layers apply their own pp_layer_offset
-                    layer = build_module(
-                        submodules.moe_layer,
-                        config=self.config,
-                        layer_number=i + 1,
-                        pg_collection=pg_collection,
-                        is_mtp_layer=is_mtp_layer,
-                    )
-                else:
-                    assert False, "unexpected layer_type"
-            self.layers.append(layer)
+        self._pp_layer_offset = pp_layer_offset
+
+        # Build main decoder layers using shared layer builder
+        self.layers = build_layers_from_pattern(
+            pattern=''.join(self.layer_type_list),
+            submodules=submodules,
+            config=self.config,
+            pg_collection=pg_collection,
+            layer_offset=pp_layer_offset,
+            residual_in_fp32=residual_in_fp32,
+            is_mtp_layer=is_mtp_layer,
+        )
 
         # Required for activation recomputation
         self.num_layers_per_pipeline_rank = len(self.layers)
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 720aa9031bf..92583da3c83 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -23,6 +23,12 @@
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+
+# Type hint imports for Mamba support
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from megatron.core.ssm.mamba_block import MambaStackSubmodules
 from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import (
@@ -450,6 +456,27 @@ def get_mtp_layer_offset(config: TransformerConfig) -> int:
     return 0
 
 
+class MTPModelLayerContainer(MegatronModule):
+    """Container for MTP inner layers built from a pattern.
+
+    This container wraps the inner layers (e.g., Mamba, Attention layers) built
+    from a pattern string. It provides proper state dict key structure:
+    `mtp_model_layer.layers.{i}.*` for checkpoint compatibility.
+
+    This is used when MTP builds its own layers using the shared layer_builder,
+    instead of receiving a pre-built MambaStack.
+    """
+
+    def __init__(self, layers: torch.nn.ModuleList, config: TransformerConfig):
+        super().__init__(config=config)
+        self.layers = layers
+
+    def set_input_tensor(self, input_tensor: Tensor):
+        """Set input tensor for pipeline parallelism compatibility."""
+        # Inner layers handle their own input tensors
+        pass
+
+
 def get_mtp_num_layers_to_build(
     config: TransformerConfig, vp_stage: Optional[int] = None, pp_rank: Optional[int] = None
 ) -> int:
@@ -539,6 +566,9 @@ def __init__(
         layer_number: int = 1,
         vp_stage: Optional[int] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
+        # New: For Mamba path - pattern and submodules to build inner layers directly
+        mtp_layer_pattern: Optional[str] = None,
+        mamba_submodules: Optional["MambaStackSubmodules"] = None,
     ):
         super().__init__(config=config)
         self.sequence_parallel = config.sequence_parallel
@@ -547,26 +577,32 @@ def __init__(
         self.vp_stage = vp_stage
         self.cp_group = pg_collection.cp
         self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern
-
-        if hasattr(self.submodules.mtp_model_layer.submodules, 'attention_layer'):
-            self_attention_spec = self.submodules.mtp_model_layer.submodules.attention_layer
-            if self_attention_spec.submodules.self_attention is not None:
-                self_attention_spec = self_attention_spec.submodules.self_attention
-                attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
-                assert attn_mask_type in SUPPORTED_ATTN_MASK, (
-                    f"Multi-Token Prediction (MTP) is not yet supported with "
-                    f"{attn_mask_type} attention mask type. "
-                    f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
-                )
-        elif hasattr(self.submodules.mtp_model_layer.submodules, 'self_attention'):
-            self_attention_spec = self.submodules.mtp_model_layer.submodules.self_attention
-            if self_attention_spec is not None:
-                attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
-                assert attn_mask_type in SUPPORTED_ATTN_MASK, (
-                    f"Multi-Token Prediction (MTP) is not yet supported with "
-                    f"{attn_mask_type} attention mask type. "
-                    f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
-                )
+        self.mtp_layer_pattern = mtp_layer_pattern
+
+        # Validate attention mask type if using transformer-based inner layers
+        if (
+            self.submodules.mtp_model_layer is not None
+            and hasattr(self.submodules.mtp_model_layer, 'submodules')
+        ):
+            if hasattr(self.submodules.mtp_model_layer.submodules, 'attention_layer'):
+                self_attention_spec = self.submodules.mtp_model_layer.submodules.attention_layer
+                if self_attention_spec.submodules.self_attention is not None:
+                    self_attention_spec = self_attention_spec.submodules.self_attention
+                    attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
+                    assert attn_mask_type in SUPPORTED_ATTN_MASK, (
+                        f"Multi-Token Prediction (MTP) is not yet supported with "
+                        f"{attn_mask_type} attention mask type. "
+                        f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
+                    )
+            elif hasattr(self.submodules.mtp_model_layer.submodules, 'self_attention'):
+                self_attention_spec = self.submodules.mtp_model_layer.submodules.self_attention
+                if self_attention_spec is not None:
+                    attn_mask_type = self_attention_spec.params.get('attn_mask_type', '')
+                    assert attn_mask_type in SUPPORTED_ATTN_MASK, (
+                        f"Multi-Token Prediction (MTP) is not yet supported with "
+                        f"{attn_mask_type} attention mask type. "
+                        f"The supported attention mask types are {SUPPORTED_ATTN_MASK}."
+                    )
 
         self.enorm = build_module(
             self.submodules.enorm,
@@ -599,8 +635,27 @@ def __init__(
             is_expert=False,
             tp_comm_buffer_name="mtp_eh_proj"
         )
-        if self.config.mtp_num_layers is not None:
+
+        # Build inner layers: three possible paths
+        # 1. New Mamba path: build from pattern using shared layer_builder
+        # 2. Legacy Mamba path: build MambaStack with mtp_hybrid_override_pattern
+        # 3. GPT path: single TransformerLayer
+        if mtp_layer_pattern is not None and mamba_submodules is not None:
+            # New Mamba path: build inner layers from pattern using shared layer_builder
+            from megatron.core.ssm.layer_builder import build_layers_from_pattern
+
+            inner_layers = build_layers_from_pattern(
+                pattern=mtp_layer_pattern,
+                submodules=mamba_submodules,
+                config=config,
+                pg_collection=pg_collection,
+                layer_offset=0,
+                is_mtp_layer=True,
+            )
+            self.mtp_model_layer = MTPModelLayerContainer(inner_layers, config=config)
+        elif self.config.mtp_num_layers is not None:
             if self.mtp_hybrid_override_pattern is not None:
+                # Legacy Mamba path: build MambaStack with override pattern
                 pg_collection = ProcessGroupCollection.use_mpu_process_groups()
 
                 # We do not need pre and post process stage for MTP layer, given they are
@@ -620,9 +675,7 @@ def __init__(
                     is_mtp_layer=True
                 )
             else:
-                # Uses the transformer block spec for MTP layer. This option is only
-                # implemented for the GPT model. In hybrid model, we model transformer
-                # block spec for MTP layers with the hybrid override pattern.
+                # GPT path: Uses the transformer block spec for MTP layer
                 self.mtp_model_layer = build_module(
                     self.submodules.mtp_model_layer,
                     config=self.config,
@@ -747,7 +800,17 @@ def _proj_and_transformer_layer(
             # transformer layer is cudagraphed, the FP8GlobalStateManager.is_first_fp8_module() is
             # True so that the fp8 weight caching can be triggered correctly.
             with transformer_layer_fp8_context:
-                if self.mtp_hybrid_override_pattern is not None:
+                if self.mtp_layer_pattern is not None:
+                    # New Mamba path: run through pattern-based inner layers
+                    hidden_states = self._run_pattern_layers(
+                        hidden_states=hidden_states,
+                        attention_mask=attention_mask,
+                        rotary_pos_emb=rotary_pos_emb,
+                        inference_params=inference_params,
+                        packed_seq_params=packed_seq_params,
+                    )
+                elif self.mtp_hybrid_override_pattern is not None:
+                    # Legacy Mamba path: MambaStack with override pattern
                     # Since pre-process is set to False, we need to set the input tensor manually.
                     self.mtp_model_layer.set_input_tensor(hidden_states)
                     hidden_states = self.mtp_model_layer(
@@ -758,6 +821,7 @@ def _proj_and_transformer_layer(
                         inference_context=context,
                     )
                 else:
+                    # GPT path: single TransformerLayer
                     hidden_states, _ = self.mtp_model_layer(
                         hidden_states=hidden_states,
                         attention_mask=attention_mask,
@@ -790,6 +854,52 @@ def _postprocess(self, hidden_states: torch.Tensor):
 
         return hidden_states
 
+    def _run_pattern_layers(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        rotary_pos_emb: Optional[torch.Tensor] = None,
+        inference_params: Optional[InferenceParams] = None,
+        packed_seq_params: Optional[PackedSeqParams] = None,
+    ) -> torch.Tensor:
+        """Run through pattern-based inner layers (Mamba path).
+
+        This method iterates through the layers in MTPModelLayerContainer,
+        calling each layer with the appropriate interface based on its type.
+
+        Args:
+            hidden_states: Input hidden states [s, b, h]
+            attention_mask: Attention mask tensor
+            rotary_pos_emb: Rotary position embeddings
+            inference_params: Inference parameters
+            packed_seq_params: Packed sequence parameters
+
+        Returns:
+            Output hidden states after processing through all inner layers
+        """
+        from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols
+        from megatron.core.transformer.transformer_layer import TransformerLayer
+
+        for layer_type, layer in zip(self.mtp_layer_pattern, self.mtp_model_layer.layers):
+            if layer_type == Symbols.ATTENTION or isinstance(layer, TransformerLayer):
+                # TransformerLayer (attention or MoE) returns tuple (hidden_states, context)
+                hidden_states, _ = layer(
+                    hidden_states=hidden_states,
+                    attention_mask=attention_mask,
+                    rotary_pos_emb=rotary_pos_emb,
+                    inference_params=inference_params,
+                    packed_seq_params=packed_seq_params,
+                )
+            else:
+                # MambaLayer, MLPLayer - return hidden_states directly
+                hidden_states = layer(
+                    hidden_states=hidden_states,
+                    attention_mask=attention_mask,
+                    inference_context=inference_params,
+                )
+
+        return hidden_states
+
     def _checkpointed_forward(self, forward_func, *args, **kwargs):
         def checkpoint_handler():
             """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`"""
@@ -1005,6 +1115,11 @@ def __init__(
         vp_stage: Optional[int] = None,
         mtp_hybrid_override_pattern: str = None,
         pg_collection: ProcessGroupCollection = None,
+        # New: For Mamba path with unified pattern syntax
+        mtp_layer_pattern: Optional[str] = None,
+        mtp_num_depths: int = 0,
+        mamba_submodules: Optional["MambaStackSubmodules"] = None,
+        mtp_shared_weights: bool = False,
     ):
         super().__init__(config=config)
         self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern
@@ -1012,6 +1127,12 @@ def __init__(
         self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor
         self.vp_stage = vp_stage
 
+        # Store new parameters for pattern-based layer building
+        self.mtp_layer_pattern = mtp_layer_pattern
+        self.mtp_num_depths = mtp_num_depths
+        self.mamba_submodules = mamba_submodules
+        self.mtp_shared_weights = mtp_shared_weights
+
         self.mtp_use_repeated_layer = self.config.mtp_use_repeated_layer
 
         vp_size = config.virtual_pipeline_model_parallel_size
@@ -1038,7 +1159,14 @@ def __init__(
         self.cp_group = pg_collection.cp
 
     def _build_layers(self, pg_collection):
-        def build_layer(layer_spec, layer_number):
+        # Determine number of depths to build
+        if self.mtp_num_depths > 0:
+            num_depths = self.mtp_num_depths
+        else:
+            num_depths = self.config.mtp_num_layers or len(self.submodules.layer_specs)
+
+        def build_layer_legacy(layer_spec, layer_number):
+            """Build layer using legacy spec-based approach."""
             fp8_init_context = get_fp8_context(self.config, is_init=True)
             with fp8_init_context:
                 module = build_module(
@@ -1051,18 +1179,57 @@ def build_layer(layer_spec, layer_number):
                 )
             return module
 
-        if self.mtp_use_repeated_layer:
+        def build_layer_with_pattern(layer_spec, layer_number, mtp_layer_pattern, mamba_submodules):
+            """Build layer using pattern-based approach (new Mamba path)."""
+            fp8_init_context = get_fp8_context(self.config, is_init=True)
+            with fp8_init_context:
+                module = build_module(
+                    layer_spec,
+                    config=self.config,
+                    layer_number=layer_number,
+                    vp_stage=self.vp_stage,
+                    pg_collection=pg_collection,
+                    mtp_layer_pattern=mtp_layer_pattern,
+                    mamba_submodules=mamba_submodules,
+                )
+            return module
+
+        # New Mamba path: use mtp_layer_pattern and mamba_submodules
+        if self.mtp_layer_pattern is not None and self.mamba_submodules is not None:
+            if self.mtp_shared_weights:
+                # Shared weights: build one layer, use it for all depths
+                layer_spec = self.submodules.layer_specs[0]
+                shared_layer = build_layer_with_pattern(
+                    layer_spec, layer_number=1,
+                    mtp_layer_pattern=self.mtp_layer_pattern,
+                    mamba_submodules=self.mamba_submodules,
+                )
+                self.layers = torch.nn.ModuleList([shared_layer])
+            else:
+                # Non-shared: each depth gets its own layers
+                self.layers = torch.nn.ModuleList([
+                    build_layer_with_pattern(
+                        self.submodules.layer_specs[min(i, len(self.submodules.layer_specs) - 1)],
+                        layer_number=i + 1,
+                        mtp_layer_pattern=self.mtp_layer_pattern,
+                        mamba_submodules=self.mamba_submodules,
+                    )
+                    for i in range(num_depths)
+                ])
+        elif self.mtp_use_repeated_layer:
+            # Legacy repeated layer mode
             assert len(self.submodules.layer_specs) == 1, (
                 f"Repeated MTP mode requires exactly 1 layer spec, got {len(self.submodules.layer_specs)}. "
                 f"The layer will be applied {self.config.mtp_num_layers} times."
             )
             self.layers = torch.nn.ModuleList([
-                build_layer(self.submodules.layer_specs[0], layer_number=1)
+                build_layer_legacy(self.submodules.layer_specs[0], layer_number=1)
             ])
         else:
+            # Legacy mode: build from layer_specs
             self.layers = torch.nn.ModuleList(
                 [
-                    build_layer(layer_spec, i + 1)
+                    build_layer_legacy(layer_spec, i + 1)
                     for i, layer_spec in enumerate(self.submodules.layer_specs)
                 ]
             )

From c4917051d809952c89a38d4c226fb8fefe49cf97 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Mon, 12 Jan 2026 17:30:26 -0800
Subject: [PATCH 27/48] MambaModel Integration + Finalization

---
 gpt_builders.py                               | 39 +++++--------
 mamba_builders.py                             | 21 +++----
 megatron/core/models/mamba/mamba_model.py     | 40 ++++++++++---
 megatron/core/pipeline_parallel/schedules.py  |  4 +-
 .../transformer/multi_token_prediction.py     | 16 ++++--
 megatron/training/arguments.py                | 56 ++++++++++---------
 .../test_multi_token_prediction.py            | 26 ++++-----
 7 files changed, 111 insertions(+), 91 deletions(-)

diff --git a/gpt_builders.py b/gpt_builders.py
index 9ba7ac8f446..4b86e30e597 100644
--- a/gpt_builders.py
+++ b/gpt_builders.py
@@ -8,7 +8,6 @@
     get_gpt_layer_with_inference_spec,
     get_gpt_mtp_block_spec,
 )
-from megatron.core.models.mamba.mamba_layer_specs import get_mamba_mtp_block_spec
 from megatron.core.models.gpt.heterogeneous.heterogeneous_layer_specs import (
     get_gpt_heterogeneous_layer_spec,
 )
@@ -61,29 +60,22 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_
                 transformer_layer_spec = _get_transformer_layer_spec(use_te, config)
         mtp_block_spec = None
         if args.mtp_num_layers is not None:
-            if args.mtp_hybrid_override_pattern is not None:
-                # Uses the hybrid model for the MTP module.
-                mtp_mamba_stack_spec = import_module(args.mtp_spec)
-                mtp_block_spec = get_mamba_mtp_block_spec(
-                    config, mtp_mamba_stack_spec, use_transformer_engine=use_te, vp_stage=vp_stage
-                )
+            assert not (config.transformer_impl == "inference_optimized")
+            if (
+                hasattr(transformer_layer_spec, 'layer_specs')
+                and len(transformer_layer_spec.layer_specs) == 0
+            ):
+                # Get the decoder layer spec explicitly if no decoder layer in the last stage,
+                # Only happens with block spec (TransformerBlockSubmodules) when using MoE.
+                transformer_layer_spec_for_mtp = _get_transformer_layer_spec(use_te, config)
             else:
-                assert not (config.transformer_impl == "inference_optimized")
-                if (
-                    hasattr(transformer_layer_spec, 'layer_specs')
-                    and len(transformer_layer_spec.layer_specs) == 0
-                ):
-                    # Get the decoder layer spec explicitly if no decoder layer in the last stage,
-                    # Only happens with block spec (TransformerBlockSubmodules) when using MoE.
-                    transformer_layer_spec_for_mtp = _get_transformer_layer_spec(use_te, config)
-                else:
-                    transformer_layer_spec_for_mtp = transformer_layer_spec
-                mtp_block_spec = get_gpt_mtp_block_spec(
-                    config,
-                    transformer_layer_spec_for_mtp,
-                    use_transformer_engine=use_te,
-                    vp_stage=vp_stage,
-                )
+                transformer_layer_spec_for_mtp = transformer_layer_spec
+            mtp_block_spec = get_gpt_mtp_block_spec(
+                config,
+                transformer_layer_spec_for_mtp,
+                use_transformer_engine=use_te,
+                vp_stage=vp_stage,
+            )
 
         model = GPTModel(
             config=config,
@@ -102,7 +94,6 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_
             mtp_block_spec=mtp_block_spec,
             vp_stage=vp_stage,
             pg_collection=pg_collection,
-            mtp_hybrid_override_pattern=args.mtp_hybrid_override_pattern,
         )
 
     return model
diff --git a/mamba_builders.py b/mamba_builders.py
index ba769d43ca2..29fce24baac 100644
--- a/mamba_builders.py
+++ b/mamba_builders.py
@@ -1,7 +1,6 @@
 # Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 
 from model_provider import count_parameters_in_layer
-from megatron.core.models.mamba.mamba_layer_specs import get_mamba_mtp_block_spec
 from megatron.core.models.mamba import MambaModel
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.spec_utils import import_module
@@ -9,6 +8,7 @@
 from megatron.training.arguments import core_transformer_config_from_args
 from megatron.core.models.mamba.mamba_layer_specs import mamba_inference_stack_spec
 
+
 def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None):
     print_rank_0('building MAMBA model ...')
     if config is None:
@@ -16,22 +16,17 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
     assert args.use_legacy_models is False, "Mamba only supported in Mcore!"
 
     if config.transformer_impl == "inference_optimized":
-        mamba_stack_spec = mamba_inference_stack_spec 
-        assert not config.inference_fuse_tp_communication, "inference_fuse_tp_communication is not supported for Mamba"
+        mamba_stack_spec = mamba_inference_stack_spec
+        assert not config.inference_fuse_tp_communication, (
+            "inference_fuse_tp_communication is not supported for Mamba"
+        )
     elif args.spec is not None:
         mamba_stack_spec = import_module(args.spec)
     else:
         raise ValueError("You must provide a valid Mamba layer spec via --spec")
 
-    mtp_block_spec = None
-    if args.mtp_num_layers is not None:
-        use_te = args.transformer_impl == "transformer_engine"
-        assert args.mtp_spec is not None
-        assert args.mtp_hybrid_override_pattern is not None, "We need to set the override hybrid pattern for MTP layers!"
-        mtp_mamba_stack_spec = import_module(args.mtp_spec)
-        mtp_block_spec = get_mamba_mtp_block_spec(
-            config, mtp_mamba_stack_spec, use_transformer_engine=use_te, vp_stage=vp_stage
-        )
+    # MTP is now configured via unified pattern in hybrid_override_pattern
+    # e.g., "M*M*/MM/MM" means main="M*M*", mtp="MM", 2 depths
     model = MambaModel(
         config=config,
         mamba_stack_spec=mamba_stack_spec,
@@ -41,7 +36,6 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
         hybrid_attention_ratio=args.hybrid_attention_ratio,
         hybrid_mlp_ratio=args.hybrid_mlp_ratio,
         hybrid_override_pattern=args.hybrid_override_pattern,
-        mtp_hybrid_override_pattern=args.mtp_hybrid_override_pattern,
         post_process=post_process,
         fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
         parallel_output=True,
@@ -51,7 +45,6 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
         rotary_base=args.rotary_base,
         pg_collection=pg_collection,
         vp_stage=vp_stage,
-        mtp_block_spec=mtp_block_spec,
     )
 
     for l in range(model.decoder.num_layers_per_pipeline_rank):
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 023d337e3f2..a15d7b96ab5 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -44,7 +44,11 @@ class MambaModel(LanguageModule):
         hybrid_attention_ratio (float, optional): The target ratio of attention
             layers to total layers
         hybrid_mlp_ratio (float, optional): The target ratio of mlp layers to total layers
-        hybrid_override_pattern (str, optional): The hybrid layer pattern to override with
+        hybrid_override_pattern (str, optional): Unified hybrid layer pattern with optional MTP.
+            Format: "<main_pattern>/<mtp_pattern>/<mtp_pattern>/..."
+            Examples:
+                - "M*M*" -> main decoder only, no MTP
+                - "M*M*/MM/MM" -> main="M*M*", mtp="MM", 2 depths
         post_process (bool, optional): Include an output layer (used with pipeline parallelism).
             Defaults to True.
         fp16_lm_cross_entropy (bool, optional): Defaults to False.
@@ -74,7 +78,6 @@ def __init__(
         hybrid_attention_ratio: float = 0.0,
         hybrid_mlp_ratio: float = 0.0,
         hybrid_override_pattern: str = None,
-        mtp_hybrid_override_pattern: str = None,
         post_process: bool = True,
         fp16_lm_cross_entropy: bool = False,
         parallel_output: bool = True,
@@ -87,7 +90,6 @@ def __init__(
         seq_len_interpolation_factor: Optional[float] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
         vp_stage: Optional[int] = None,
-        mtp_block_spec: Optional[ModuleSpec] = None,
     ) -> None:
         super().__init__(config=config, pg_collection=pg_collection)
 
@@ -101,14 +103,24 @@ def __init__(
         self.hybrid_attention_ratio = hybrid_attention_ratio
         self.hybrid_mlp_ratio = hybrid_mlp_ratio
         self.hybrid_override_pattern = hybrid_override_pattern
-        self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern
         self.post_process = post_process
         self.fp16_lm_cross_entropy = fp16_lm_cross_entropy
         self.parallel_output = parallel_output
         self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
         self.position_embedding_type = position_embedding_type
-        self.mtp_block_spec = mtp_block_spec
-        self.mtp_process = mtp_block_spec is not None
+        self.vp_stage = vp_stage
+
+        # Parse unified pattern to extract main and MTP components
+        from megatron.core.ssm.mamba_hybrid_layer_allocation import parse_hybrid_pattern
+
+        parsed = parse_hybrid_pattern(hybrid_override_pattern)
+        self.mtp_pattern = parsed.mtp_pattern
+        self.mtp_num_depths = parsed.mtp_num_depths
+
+        # Determine if MTP is needed (based on pattern parsing)
+        self.mtp_process = (
+            self.mtp_pattern is not None and self.mtp_num_depths > 0 and self.post_process
+        )
 
         # Cache for RoPE tensors which do not change between iterations.
         self.rotary_pos_emb_cache = {}
@@ -150,9 +162,23 @@ def __init__(
             pg_collection=self.pg_collection,
         )
 
+        # MTP block (MambaModel creates it, MTP builds its own inner layers)
         if self.mtp_process:
+            from megatron.core.models.mamba.mamba_layer_specs import get_mamba_mtp_block_spec
+
+            mtp_block_spec = get_mamba_mtp_block_spec()
+
+            # Extract mamba_submodules from mamba_stack_spec for MTP to use
+            mamba_submodules = mamba_stack_spec.submodules
+
             self.mtp = MultiTokenPredictionBlock(
-                config=self.config, spec=self.mtp_block_spec, vp_stage=vp_stage, mtp_hybrid_override_pattern=self.mtp_hybrid_override_pattern
+                config=self.config,
+                spec=mtp_block_spec,
+                pg_collection=self.pg_collection,
+                vp_stage=self.vp_stage,
+                mtp_layer_pattern=self.mtp_pattern,
+                mtp_num_depths=self.mtp_num_depths,
+                mamba_submodules=mamba_submodules,
             )
 
         # Output
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index c8c520b2e51..cd07ba247c1 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -283,11 +283,11 @@ def forward_step_calc_loss(
             if config.grad_scale_func is not None
             else torch.ones(1, device=output_tensor.device)
         )
-        # Set the loss scale
+        # Set the loss scale (multiply by cp_group_size to match MoE scaling)
         if config.calculate_per_token_loss:
             MTPLossAutoScaler.set_loss_scale(loss_scale)
         else:
-            MTPLossAutoScaler.set_loss_scale(loss_scale / num_microbatches)
+            MTPLossAutoScaler.set_loss_scale(loss_scale * cp_group_size / num_microbatches)
 
     return output_tensor, num_tokens
 
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 92583da3c83..eaad51448a1 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -986,6 +986,15 @@ def forward(
             packed_seq_params=packed_seq_params,
         )
 
+        # Roll RoPE to match rolled positions (position_ids were rolled in _get_embeddings)
+        # After rolling, index i should use RoPE for position i+1
+        if rotary_pos_emb is not None:
+            rotary_pos_emb = torch.roll(rotary_pos_emb, shifts=-1, dims=0)
+        if rotary_pos_cos is not None:
+            rotary_pos_cos = torch.roll(rotary_pos_cos, shifts=-1, dims=0)
+        if rotary_pos_sin is not None:
+            rotary_pos_sin = torch.roll(rotary_pos_sin, shifts=-1, dims=0)
+
         if self.config.recompute_granularity == 'full' and self.training:
             hidden_states = self._checkpointed_forward(
                 self._proj_and_transformer_layer,
@@ -1119,7 +1128,6 @@ def __init__(
         mtp_layer_pattern: Optional[str] = None,
         mtp_num_depths: int = 0,
         mamba_submodules: Optional["MambaStackSubmodules"] = None,
-        mtp_shared_weights: bool = False,
     ):
         super().__init__(config=config)
         self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern
@@ -1131,8 +1139,8 @@ def __init__(
         self.mtp_layer_pattern = mtp_layer_pattern
         self.mtp_num_depths = mtp_num_depths
         self.mamba_submodules = mamba_submodules
-        self.mtp_shared_weights = mtp_shared_weights
 
+        # Use config.mtp_use_repeated_layer for both GPT and Mamba paths
         self.mtp_use_repeated_layer = self.config.mtp_use_repeated_layer
 
         vp_size = config.virtual_pipeline_model_parallel_size
@@ -1196,8 +1204,8 @@ def build_layer_with_pattern(layer_spec, layer_number, mtp_layer_pattern, mamba_
 
         # New Mamba path: use mtp_layer_pattern and mamba_submodules
         if self.mtp_layer_pattern is not None and self.mamba_submodules is not None:
-            if self.mtp_shared_weights:
-                # Shared weights: build one layer, use it for all depths
+            if self.mtp_use_repeated_layer:
+                # Shared/repeated layer: build one layer, use it for all depths
                 layer_spec = self.submodules.layer_specs[0]
                 shared_layer = build_layer_with_pattern(
                     layer_spec, layer_number=1,
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index d4a3337d423..12ab8647618 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -529,9 +529,27 @@ def validate_args(args, defaults={}):
                 args.global_batch_size), flush=True)
     assert args.global_batch_size > 0
 
+    # Deprecation warnings for legacy MTP arguments
     if args.mtp_hybrid_override_pattern is not None:
-        assert args.mtp_spec is not None, "--mtp-spec must be set for mtp hybrid MTP models"
-    
+        import warnings
+        warnings.warn(
+            "--mtp-hybrid-override-pattern is deprecated. "
+            "For new Mamba+MTP models, use unified --hybrid-override-pattern instead. "
+            "Example: 'M*M*/MM/MM' means main='M*M*', MTP pattern='MM' with 2 depths. "
+            "This argument is kept only for loading old checkpoints.",
+            DeprecationWarning,
+            stacklevel=2
+        )
+    if args.mtp_spec is not None:
+        import warnings
+        warnings.warn(
+            "--mtp-spec is deprecated. "
+            "For new Mamba+MTP models, use unified --hybrid-override-pattern instead. "
+            "This argument is kept only for loading old checkpoints.",
+            DeprecationWarning,
+            stacklevel=2
+        )
+
     # Uneven virtual pipeline parallelism
     assert (
         int(args.num_layers_per_virtual_pipeline_stage is not None)
@@ -602,18 +620,7 @@ def validate_args(args, defaults={}):
 
                 assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \
                     'Number of layers should be divisible by the pipeline-model-parallel size'
-    # Validate MTP hybrid pattern segment count
-    if args.mtp_hybrid_override_pattern is not None:
-        mtp_segments = len(args.mtp_hybrid_override_pattern.split('|'))
-        expected_segments = args.pipeline_model_parallel_size * (
-            args.virtual_pipeline_model_parallel_size or 1
-        )
-        assert mtp_segments == expected_segments, (
-            f'MTP hybrid pattern has {mtp_segments} segments but expected {expected_segments} '
-            f'(PP={args.pipeline_model_parallel_size} * VP={args.virtual_pipeline_model_parallel_size or 1}). '
-            f'Use: {"|||" * (expected_segments - 1)}*E'
-        )
-    
+
     if args.virtual_pipeline_model_parallel_size is not None:
         if args.overlap_p2p_comm:
             assert args.pipeline_model_parallel_size > 1, \
@@ -3391,19 +3398,18 @@ def _add_experimental_args(parser):
                             'hybrid ratio arguments, then the number of each type'
                             'of layer in the override pattern must match number in'
                             'the overidden pattern')
+    # DEPRECATED: Legacy MTP arguments kept for checkpoint compatibility.
+    # For new Mamba+MTP models, use unified pattern in --hybrid-override-pattern instead.
+    # Example: "M*M*/MM/MM" means main="M*M*", MTP pattern="MM" with 2 depths.
     group.add_argument("--mtp-hybrid-override-pattern", type=str, default=None,
-                       help='Force a specific hybrid layer pattern for MTP layers. '
-                       'If a value greater than 0.0 is supplied to any of the hybrid ratio '
-                       'arguments, then the number of each type of layer in the '
-                       'override pattern must match number in the overidden '
-                       'pattern')
+                       help='DEPRECATED: Use unified --hybrid-override-pattern instead. '
+                       'Legacy argument for loading old checkpoints. '
+                       'Force a specific hybrid layer pattern for MTP layers.')
     group.add_argument('--mtp-spec', type=str, default=None, nargs='*',
-                       help='Specify the <module_location function_name> pair '
-                       'that returns a spec for mtp layer to customize a model,' 
-                       'transformer block, or transformer layer, depending on '
-                       'the use case. To use local spec specify local as the argument.'
-                       'For more details, see the model class, '
-                       'TransformerBlock class, or TransformerLayer class.')
+                       help='DEPRECATED: Use unified --hybrid-override-pattern instead. '
+                       'Legacy argument for loading old checkpoints. '
+                       'Specify the <module_location function_name> pair '
+                       'that returns a spec for mtp layer.')
     group.add_argument('--mamba-state-dim', type=int, default=128,
                        help='State dimension for Mamba layers.')
     group.add_argument('--mamba-head-dim', type=int, default=64,
diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index a41ad3fd8e0..9c21e2a51cc 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -13,10 +13,7 @@
     get_gpt_mtp_block_spec,
 )
 from megatron.core.models.gpt.gpt_model import GPTModel
-from megatron.core.models.mamba.mamba_layer_specs import (
-    get_mamba_mtp_block_spec,
-    mamba_stack_spec,
-)
+from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec
 from megatron.core.models.mamba.mamba_model import MambaModel
 from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator
 from megatron.core.packed_seq_params import PackedSeqParams
@@ -711,15 +708,18 @@ def model_provider(
         post_process=True,
         **config_kwargs,
     ):
-        """Model provider for Mamba hybrid models with MTP."""
+        """Model provider for Mamba hybrid models with MTP.
+
+        Uses the unified pattern syntax where MTP is configured via hybrid_override_pattern:
+        Format: "<main_pattern>/<mtp_pattern>/<mtp_pattern>/..."
+        Example: "M*M*/M*/M*" = main decoder "M*M*", MTP pattern "M*" with 2 depths
+        """
         model_parallel_cuda_manual_seed(_SEED)
         args = get_args()
         config = core_transformer_config_from_args(args)
 
-        mtp_block_spec = get_mamba_mtp_block_spec(
-            config=config, spec=mamba_stack_spec, use_transformer_engine=True
-        )
-
+        # MTP is configured via unified pattern in hybrid_override_pattern
+        # MambaModel creates the MTP block internally based on the parsed pattern
         model = MambaModel(
             config=config,
             mamba_stack_spec=mamba_stack_spec,
@@ -730,13 +730,11 @@ def model_provider(
             hybrid_attention_ratio=args.hybrid_attention_ratio,
             hybrid_mlp_ratio=args.hybrid_mlp_ratio,
             hybrid_override_pattern=args.hybrid_override_pattern,
-            mtp_hybrid_override_pattern=args.mtp_hybrid_override_pattern,
             fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
             parallel_output=True,
             share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
             position_embedding_type=args.position_embedding_type,
             rotary_percent=args.rotary_percent,
-            mtp_block_spec=mtp_block_spec,
         )
         return model
 
@@ -776,11 +774,9 @@ def create_test_args(
         args.bf16 = True
         args.hybrid_attention_ratio = 0.5
         args.hybrid_mlp_ratio = 0.0
-        args.hybrid_override_pattern = "M*M*"
-        args.mtp_hybrid_override_pattern = "M*"
-        args.mtp_num_layers_per_layer = 2
+        # Unified pattern: "main/mtp/mtp" - main decoder "M*M*", MTP pattern "M*" with 2 depths
+        args.hybrid_override_pattern = "M*M*/M*/M*"
         args.spec = "megatron.core.models.mamba.mamba_layer_specs.mamba_stack_spec"
-        args.mtp_spec = "megatron.core.models.mamba.mamba_layer_specs.mamba_stack_spec"
 
         if fp8 is not None:
             args.fp8 = 'e4m3'

From 77ba45d5946ed6d8a5cdd0b9bd5ecc1a05a92569 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Tue, 13 Jan 2026 14:40:39 -0800
Subject: [PATCH 28/48] Consolidate MTP weight tying

---
 .../common/language_module/language_module.py | 16 +++++++--
 megatron/core/models/gpt/gpt_model.py         | 35 -------------------
 megatron/core/models/mamba/mamba_model.py     | 26 --------------
 3 files changed, 14 insertions(+), 63 deletions(-)

diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
index 6f145f7f32c..1e4ea3851b6 100644
--- a/megatron/core/models/common/language_module/language_module.py
+++ b/megatron/core/models/common/language_module/language_module.py
@@ -310,7 +310,13 @@ def sharded_state_dict(
         # processing stage.
         if self.mtp_process and not self.pre_process:
             emb_weight = self.embedding.word_embeddings.weight
-            tie_word_embeddings_state_dict(sharded_state_dict, emb_weight, first_stage_word_emb_key)
+            tie_word_embeddings_state_dict(
+                sharded_state_dict,
+                emb_weight,
+                first_stage_word_emb_key,
+                tp_group=self.tp_group,
+                dp_cp_group=metadata['dp_cp_group'],
+            )
 
         if self.mtp_process and not self.post_process:
             # We only need to tie the output layer weight if share_embeddings_and_output_weights
@@ -318,7 +324,13 @@ def sharded_state_dict(
             # will be stored in embedding layer, and output layer will not have any weight.
             if not self.share_embeddings_and_output_weights:
                 output_layer_weight = self.output_layer.weight
-                tie_output_layer_state_dict(sharded_state_dict, output_layer_weight, output_layer_weight_key)
+                tie_output_layer_state_dict(
+                    sharded_state_dict,
+                    output_layer_weight,
+                    output_layer_weight_key,
+                    tp_group=self.tp_group,
+                    dp_cp_group=metadata['dp_cp_group'],
+                )
         if self.share_embeddings_and_output_weights:
             self.tie_embeddings_and_output_weights_state_dict(
                 sharded_state_dict, output_layer_weight_key, first_stage_word_emb_key, metadata
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index 465103aafbf..16dad3e94ce 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -27,8 +27,6 @@
     MTPLossLoggingHelper,
     MultiTokenPredictionBlock,
     roll_tensor,
-    tie_output_layer_state_dict,
-    tie_word_embeddings_state_dict,
 )
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_block import TransformerBlock
@@ -710,37 +708,4 @@ def sharded_state_dict(
             output_extra_state and output_extra_state.data
         ), f'Expected output layer extra state to be empty, got: {output_extra_state}'
 
-        # Multi-Token Prediction (MTP) need both embedding layer and output layer in
-        # mtp process stage.
-        # If MTP is not placed in the pre processing stage, we need to maintain a copy of
-        # embedding layer in the mtp process stage and tie it to the embedding in the pre
-        # processing stage.
-        # Also, if MTP is not placed in the post processing stage, we need to maintain a copy
-        # of output layer in the mtp process stage and tie it to the output layer in the post
-        # processing stage.
-        if self.mtp_process and not self.pre_process:
-            emb_weight_key = f'{prefix}embedding.word_embeddings.weight'
-            emb_weight = self.embedding.word_embeddings.weight
-            tie_word_embeddings_state_dict(
-                sharded_state_dict,
-                emb_weight,
-                emb_weight_key,
-                tp_group=self.tp_group,
-                dp_cp_group=metadata['dp_cp_group'],
-            )
-        if self.mtp_process and not self.post_process:
-            # We only need to tie the output layer weight if share_embeddings_and_output_weights
-            # is False. Because if share_embeddings_and_output_weights is True, the shared weight
-            # will be stored in embedding layer, and output layer will not have any weight.
-            if not self.share_embeddings_and_output_weights:
-                output_layer_weight_key = f'{prefix}output_layer.weight'
-                output_layer_weight = self.output_layer.weight
-                tie_output_layer_state_dict(
-                    sharded_state_dict,
-                    output_layer_weight,
-                    output_layer_weight_key,
-                    tp_group=self.tp_group,
-                    dp_cp_group=metadata['dp_cp_group'],
-                )
-
         return sharded_state_dict
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index a15d7b96ab5..222c3e9fe95 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -24,8 +24,6 @@
     MTPLossLoggingHelper,
     MultiTokenPredictionBlock,
     roll_tensor,
-    tie_output_layer_state_dict,
-    tie_word_embeddings_state_dict,
 )
 from megatron.core.packed_seq_params import PackedSeqParams
 
@@ -507,28 +505,4 @@ def sharded_state_dict(
             ShardedStateDict: sharded state dict for the MambaModel
         """
         sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
-        
-        # Multi-Token Prediction (MTP) need both embedding layer and output layer in
-        # mtp process stage.
-        # If MTP is not placed in the pre processing stage, we need to maintain a copy of
-        # embedding layer in the mtp process stage and tie it to the embedding in the pre
-        # processing stage.
-        # Also, if MTP is not placed in the post processing stage, we need to maintain a copy
-        # of output layer in the mtp process stage and tie it to the output layer in the post
-        # processing stage.
-        if self.mtp_process and not self.pre_process:
-            emb_weight_key = f'{prefix}embedding.word_embeddings.weight'
-            emb_weight = self.embedding.word_embeddings.weight
-            tie_word_embeddings_state_dict(sharded_state_dict, emb_weight, emb_weight_key)
-        if self.mtp_process and not self.post_process:
-            # We only need to tie the output layer weight if share_embeddings_and_output_weights
-            # is False. Because if share_embeddings_and_output_weights is True, the shared weight
-            # will be stored in embedding layer, and output layer will not have any weight.
-            if not self.share_embeddings_and_output_weights:
-                output_layer_weight_key = f'{prefix}output_layer.weight'
-                output_layer_weight = self.output_layer.weight
-                tie_output_layer_state_dict(
-                    sharded_state_dict, output_layer_weight, output_layer_weight_key
-                )
-
         return sharded_state_dict

From 31f512b30d6ba3856d70adb1ff86fc860efcfb37 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Tue, 13 Jan 2026 15:27:15 -0800
Subject: [PATCH 29/48] consolidate shared_embedding_or_output_weight

---
 .../common/language_module/language_module.py |  6 ++---
 megatron/core/models/gpt/gpt_model.py         | 23 +------------------
 megatron/core/models/mamba/mamba_model.py     | 21 -----------------
 3 files changed, 4 insertions(+), 46 deletions(-)

diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
index 1e4ea3851b6..26537b2b5b9 100644
--- a/megatron/core/models/common/language_module/language_module.py
+++ b/megatron/core/models/common/language_module/language_module.py
@@ -260,11 +260,11 @@ def shared_embedding_or_output_weight(self) -> Tensor:
         Returns:
             Tensor: During pre processing or MTP process it returns the input embeddings weight while during post processing it returns the final output layers weight
         """
-        if self.pre_process or self.mtp_process:
+        if self.pre_process or getattr(self, 'mtp_process', False):
             # Multi-Token Prediction (MTP) need both embedding layer and output layer.
             # So there will be both embedding layer and output layer in the mtp process stage.
-            # In this case, if share_embeddings_and_output_weights is True, the shared weights
-            # will be stored in embedding layer, and output layer will not have any weight.
+            # When share_embeddings_and_output_weights is True, the embedding weight is the
+            # canonical shared weight and is passed to the output layer during forward.
             assert hasattr(
                 self, 'embedding'
             ), f"embedding is needed in this pipeline stage, but it is not initialized."
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index 16dad3e94ce..cebebd120b2 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -136,7 +136,7 @@ def __init__(
             self.rotary_base = rotary_base
         self.rotary_scaling = rope_scaling
         self.mtp_block_spec = mtp_block_spec
-        self.mtp_process = mtp_block_spec is not None
+        self.mtp_process = mtp_block_spec is not None and self.post_process
 
         if self.pre_process or self.mtp_process:
             self.embedding = LanguageModelEmbedding(
@@ -603,27 +603,6 @@ def _postprocess(
 
         return loss
 
-    def shared_embedding_or_output_weight(self) -> Tensor:
-        """Gets the embedding weight or output logit weights when share input embedding and
-        output weights set to True or when use Multi-Token Prediction (MTP) feature.
-
-        Returns:
-            Tensor: During pre processing or MTP process it returns the input embeddings weight.
-            Otherwise, during post processing it returns the final output layers weight.
-        """
-        if self.pre_process or self.mtp_process:
-            # Multi-Token Prediction (MTP) need both embedding layer and output layer.
-            # So there will be both embedding layer and output layer in the mtp process stage.
-            # In this case, if share_embeddings_and_output_weights is True, the shared weights
-            # will be stored in embedding layer, and output layer will not have any weight.
-            assert hasattr(
-                self, 'embedding'
-            ), "embedding is needed in this pipeline stage, but it is not initialized."
-            return self.embedding.word_embeddings.weight
-        elif self.post_process:
-            return self.output_layer.weight
-        return None
-
     def build_schedule_plan(
         self,
         input_ids: Tensor,
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 222c3e9fe95..54cde291320 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -467,27 +467,6 @@ def forward(
 
         return loss
 
-    def shared_embedding_or_output_weight(self) -> Optional[Tensor]:
-        """Gets the embedding weight or output logit weights when share input embedding and
-        output weights set to True or when use Multi-Token Prediction (MTP) feature.
-
-        Returns:
-            Optional[Tensor]: During pre processing or MTP process it returns the input embeddings weight.
-            Otherwise, during post processing it returns the final output layers weight.
-        """
-        if self.pre_process or self.mtp_process:
-            # Multi-Token Prediction (MTP) need both embedding layer and output layer.
-            # So there will be both embedding layer and output layer in the mtp process stage.
-            # In this case, if share_embeddings_and_output_weights is True, the shared weights
-            # will be stored in embedding layer, and output layer will not have any weight.
-            assert hasattr(
-                self, 'embedding'
-            ), f"embedding is needed in this pipeline stage, but it is not initialized."
-            return self.embedding.word_embeddings.weight
-        elif self.post_process:
-            return self.output_layer.weight
-        return None
-
     def sharded_state_dict(
         self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[Dict] = None
     ) -> ShardedStateDict:

From cf681fad944f64e1012395e0d0cb9520486b2b9d Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Tue, 13 Jan 2026 15:38:48 -0800
Subject: [PATCH 30/48] rename mtp_hybrid_override_pattern to mtp_layer_pattern

---
 megatron/core/models/gpt/gpt_model.py         |  4 ++--
 .../transformer/multi_token_prediction.py     | 19 +++++++++----------
 megatron/training/arguments.py                |  2 +-
 megatron/training/checkpointing.py            |  3 ++-
 megatron/training/training.py                 |  4 ++--
 5 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index cebebd120b2..1d19082d6eb 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -100,7 +100,7 @@ def __init__(
         mtp_block_spec: Optional[ModuleSpec] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
         vp_stage: Optional[int] = None,
-        mtp_hybrid_override_pattern: str = None,
+        mtp_layer_pattern: str = None,
     ) -> None:
         super().__init__(config=config, pg_collection=pg_collection)
 
@@ -209,7 +209,7 @@ def __init__(
 
         if self.mtp_process:
             self.mtp = MultiTokenPredictionBlock(
-                config=self.config, spec=self.mtp_block_spec, vp_stage=vp_stage, mtp_hybrid_override_pattern=mtp_hybrid_override_pattern
+                config=self.config, spec=self.mtp_block_spec, vp_stage=vp_stage, mtp_layer_pattern=mtp_layer_pattern
             )
 
         # Output
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index eaad51448a1..a1cf115012b 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -562,7 +562,7 @@ def __init__(
         self,
         config: TransformerConfig,
         submodules: MultiTokenPredictionLayerSubmodules,
-        mtp_hybrid_override_pattern: str = None,
+        mtp_layer_pattern: str = None,
         layer_number: int = 1,
         vp_stage: Optional[int] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
@@ -576,7 +576,7 @@ def __init__(
         self.layer_number = layer_number
         self.vp_stage = vp_stage
         self.cp_group = pg_collection.cp
-        self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern
+        self.mtp_layer_pattern = mtp_layer_pattern
         self.mtp_layer_pattern = mtp_layer_pattern
 
         # Validate attention mask type if using transformer-based inner layers
@@ -638,7 +638,7 @@ def __init__(
 
         # Build inner layers: three possible paths
         # 1. New Mamba path: build from pattern using shared layer_builder
-        # 2. Legacy Mamba path: build MambaStack with mtp_hybrid_override_pattern
+        # 2. Legacy Mamba path: build MambaStack with mtp_layer_pattern
         # 3. GPT path: single TransformerLayer
         if mtp_layer_pattern is not None and mamba_submodules is not None:
             # New Mamba path: build inner layers from pattern using shared layer_builder
@@ -654,20 +654,20 @@ def __init__(
             )
             self.mtp_model_layer = MTPModelLayerContainer(inner_layers, config=config)
         elif self.config.mtp_num_layers is not None:
-            if self.mtp_hybrid_override_pattern is not None:
+            if self.mtp_layer_pattern is not None:
                 # Legacy Mamba path: build MambaStack with override pattern
                 pg_collection = ProcessGroupCollection.use_mpu_process_groups()
 
                 # We do not need pre and post process stage for MTP layer, given they are
                 # handled in the MultiTokenPredictionLayer itself.
                 assert self.config.mtp_num_layers_per_layer is not None, \
-                    "mtp_num_layers_per_layer must be set when using mtp_hybrid_override_pattern"
+                    "mtp_num_layers_per_layer must be set when using mtp_layer_pattern"
                 self.mtp_model_layer = build_module(
                     self.submodules.mtp_model_layer,
                     self.config,
                     pre_process=False,
                     post_process=False,
-                    hybrid_override_pattern=self.mtp_hybrid_override_pattern,
+                    hybrid_override_pattern=self.mtp_layer_pattern,
                     dtype=self.config.params_dtype,
                     pg_collection=pg_collection,
                     vp_stage=self.vp_stage,
@@ -809,7 +809,7 @@ def _proj_and_transformer_layer(
                         inference_params=inference_params,
                         packed_seq_params=packed_seq_params,
                     )
-                elif self.mtp_hybrid_override_pattern is not None:
+                elif self.mtp_layer_pattern is not None:
                     # Legacy Mamba path: MambaStack with override pattern
                     # Since pre-process is set to False, we need to set the input tensor manually.
                     self.mtp_model_layer.set_input_tensor(hidden_states)
@@ -1122,7 +1122,6 @@ def __init__(
         config: TransformerConfig,
         spec: Union[TransformerBlockSubmodules, ModuleSpec],
         vp_stage: Optional[int] = None,
-        mtp_hybrid_override_pattern: str = None,
         pg_collection: ProcessGroupCollection = None,
         # New: For Mamba path with unified pattern syntax
         mtp_layer_pattern: Optional[str] = None,
@@ -1130,7 +1129,7 @@ def __init__(
         mamba_submodules: Optional["MambaStackSubmodules"] = None,
     ):
         super().__init__(config=config)
-        self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern
+        self.mtp_layer_pattern = mtp_layer_pattern
         self.submodules = _get_mtp_block_submodules(config, spec)
         self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor
         self.vp_stage = vp_stage
@@ -1183,7 +1182,7 @@ def build_layer_legacy(layer_spec, layer_number):
                     layer_number=layer_number,
                     vp_stage=self.vp_stage,
                     pg_collection=pg_collection,
-                    mtp_hybrid_override_pattern=self.mtp_hybrid_override_pattern
+                    mtp_layer_pattern=self.mtp_layer_pattern
                 )
             return module
 
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 12ab8647618..fea66b64834 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -530,7 +530,7 @@ def validate_args(args, defaults={}):
     assert args.global_batch_size > 0
 
     # Deprecation warnings for legacy MTP arguments
-    if args.mtp_hybrid_override_pattern is not None:
+    if args.mtp_layer_pattern is not None:
         import warnings
         warnings.warn(
             "--mtp-hybrid-override-pattern is deprecated. "
diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index 779f27fd581..d85f6e5767c 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -1347,7 +1347,8 @@ def _set_arg(arg_name, old_arg_name=None, force=False):
     _set_arg('hidden_dropout', force=True)
 
     _set_arg('hybrid_override_pattern', force=True)
-    _set_arg('mtp_hybrid_override_pattern', force=True)
+    _set_arg('mtp_layer_pattern', old_arg_name='mtp_hybrid_override_pattern', force=True)  # old name
+    _set_arg('mtp_layer_pattern', force=True)  # new mtp name, takes precedence
     _set_arg('spec', force=True)
     _set_arg('hybrid_attention_ratio', force=True)
     _set_arg('hybrid_mlp_ratio', force=True)
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 28e3d2b6e94..2484825b3a2 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -178,9 +178,9 @@ def calculate_layer_counts():
             for layer_type in args.hybrid_override_pattern:
                 if layer_type in counts:
                     counts[layer_type] += 1
-            if args.mtp_hybrid_override_pattern:
+            if args.mtp_layer_pattern:
                 assert args.mtp_num_layers is not None
-                for layer_type in args.mtp_hybrid_override_pattern:
+                for layer_type in args.mtp_layer_pattern:
                     if layer_type in counts:
                         counts[layer_type] += args.mtp_num_layers
             return counts['*'], counts['M'], counts['-'], counts['E']

From 82fab8771bcef6b255158d44da576d271a80a644 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Tue, 13 Jan 2026 16:31:17 -0800
Subject: [PATCH 31/48] Remove mamba-model's re-preprocessing

---
 megatron/core/models/mamba/mamba_model.py | 136 +---------------------
 1 file changed, 3 insertions(+), 133 deletions(-)

diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 54cde291320..ed86d42a36b 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -3,10 +3,8 @@
 from typing import Literal, Optional
 
 from torch import Tensor
-import torch
-from typing import Dict
-from megatron.core.dist_checkpointing.mapping import ShardedStateDict
-from megatron.core import parallel_state, tensor_parallel
+
+from megatron.core import tensor_parallel
 from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
 from megatron.core.inference.contexts import BaseInferenceContext
 from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
@@ -17,15 +15,9 @@
 from megatron.core.tensor_parallel import gather_from_sequence_parallel_region
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.enums import ModelType
+from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionBlock
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.utils import WrappedTensor, deprecate_inference_params
-from megatron.core.transformer.multi_token_prediction import (
-    MTPLossAutoScaler,
-    MTPLossLoggingHelper,
-    MultiTokenPredictionBlock,
-    roll_tensor,
-)
-from megatron.core.packed_seq_params import PackedSeqParams
 
 
 class MambaModel(LanguageModule):
@@ -218,94 +210,6 @@ def set_input_tensor(self, input_tensor: Tensor) -> None:
         assert len(input_tensor) == 1, 'input_tensor should only be length 1 for gpt/bert'
         self.decoder.set_input_tensor(input_tensor[0])
 
-    def _preprocess(
-        self,
-        input_ids: Tensor,
-        position_ids: Tensor,
-        decoder_input: Tensor = None,
-        inference_context: BaseInferenceContext = None,
-        packed_seq_params: PackedSeqParams = None,
-    ):
-        """Preprocesses inputs for the transformer decoder.
-
-        Applies embeddings to input tokens, or uses `decoder_input` from a previous
-        pipeline stage. Also sets up rotary positional embeddings.
-        """
-
-        # If decoder_input is provided (not None), then input_ids and position_ids are ignored.
-        # Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input.
-
-        # Decoder embedding.
-        if decoder_input is not None:
-            pass
-        elif self.pre_process:
-            decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
-        else:
-            # intermediate stage of pipeline
-            # decoder will get hidden_states from encoder.input_tensor
-            decoder_input = None
-
-        # Rotary positional embeddings (embedding is None for PP intermediate devices)
-        rotary_pos_emb = None
-        rotary_pos_cos = None
-        rotary_pos_sin = None
-        if self.position_embedding_type == 'rope' and not self.config.multi_latent_attention:
-            if not self.training and self.config.flash_decode and inference_context:
-                assert (
-                    inference_context.is_static_batching()
-                ), "GPTModel currently only supports static inference batching."
-                # Flash decoding uses precomputed cos and sin for RoPE
-                rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb_cache.setdefault(
-                    inference_context.max_sequence_length,
-                    self.rotary_pos_emb.get_cos_sin(inference_context.max_sequence_length),
-                )
-            else:
-                rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
-                    inference_context, self.decoder, decoder_input, self.config, packed_seq_params
-                )
-                rotary_pos_emb = self.rotary_pos_emb(
-                    rotary_seq_len,
-                    packed_seq=packed_seq_params is not None
-                    and packed_seq_params.qkv_format == 'thd',
-                )
-        elif self.position_embedding_type == 'mrope' and not self.config.multi_latent_attention:
-            if self.training or not self.config.flash_decode:
-                rotary_pos_emb = self.rotary_pos_emb(position_ids, self.mrope_section)
-            else:
-                # Flash decoding uses precomputed cos and sin for RoPE
-                raise NotImplementedError(
-                    "Flash decoding uses precomputed cos and sin for RoPE, not implmented in "
-                    "MultimodalRotaryEmbedding yet."
-                )
-
-        if (
-            (self.config.enable_cuda_graph or self.config.flash_decode)
-            and rotary_pos_cos is not None
-            and inference_context
-            and inference_context.is_static_batching()
-            and not self.training
-        ):
-            current_batch_size = input_ids.shape[0]
-            sequence_len_offset = torch.tensor(
-                [inference_context.sequence_len_offset] * current_batch_size,
-                dtype=torch.int32,
-                device=rotary_pos_cos.device,  # Co-locate this with the rotary tensors
-            )
-        else:
-            sequence_len_offset = None
-
-        # Wrap decoder_input to allow the decoder (TransformerBlock) to delete the
-        # reference held by this caller function, enabling early garbage collection for
-        # inference. Skip wrapping if decoder_input is logged after decoder completion.
-        if (
-            inference_context is not None
-            and not self.training
-            and not has_config_logger_enabled(self.config)
-        ):
-            decoder_input = WrappedTensor(decoder_input)
-
-        return decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset
-
     def forward(
         self,
         input_ids: Tensor,
@@ -382,16 +286,6 @@ def forward(
             output_weight = self.shared_embedding_or_output_weight()            
 
         if self.mtp_process:
-            extra_block_kwargs = None
-            packed_seq_params = None
-            decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset = self._preprocess(
-                input_ids=input_ids,
-                position_ids=position_ids,
-                decoder_input=decoder_input,
-                inference_context=inference_context,
-                packed_seq_params=packed_seq_params,
-            )
-
 
             hidden_states = self.mtp(
                 input_ids=input_ids,
@@ -400,12 +294,7 @@ def forward(
                 attention_mask=attention_mask,
                 inference_params=inference_params,
                 rotary_pos_emb=rotary_pos_emb,
-                rotary_pos_cos=rotary_pos_cos,
-                rotary_pos_sin=rotary_pos_sin,
-                packed_seq_params=packed_seq_params,
-                sequence_len_offset=sequence_len_offset,
                 embedding=self.embedding,
-                **(extra_block_kwargs or {}),
             )
 
         if not self.post_process:
@@ -466,22 +355,3 @@ def forward(
         loss = self.compute_language_model_loss(labels, logits)
 
         return loss
-
-    def sharded_state_dict(
-        self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[Dict] = None
-    ) -> ShardedStateDict:
-        """Sharded state dict implementation for MambaModel backward-compatibility.
-
-        Removing extra state.
-        Tie word embeddings and output layer in mtp process stage.
-
-        Args:
-            prefix (str): Module name prefix.
-            sharded_offsets (tuple): PP related offsets, expected to be empty at this module level.
-            metadata (Optional[Dict]): metadata controlling sharded state dict creation.
-
-        Returns:
-            ShardedStateDict: sharded state dict for the MambaModel
-        """
-        sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
-        return sharded_state_dict

From 7035cd8a93ab4dd7b3a7d6fcfe9db10e4a345c18 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Tue, 13 Jan 2026 17:21:05 -0800
Subject: [PATCH 32/48] MLPLayer inherits from TransformerLayer

---
 megatron/core/ssm/mamba_block.py | 16 ++--------------
 1 file changed, 2 insertions(+), 14 deletions(-)

diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index f606fdfefd7..09e08803e00 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -88,8 +88,6 @@ def __init__(
         dtype=None,
         pg_collection: ProcessGroupCollection = None,
         vp_stage: Optional[int] = None,
-        num_layers: int = None,
-        is_mtp_layer: bool = False,
     ) -> None:
         super().__init__(config=config)
         self.residual_in_fp32 = residual_in_fp32
@@ -97,7 +95,6 @@ def __init__(
         self.post_layer_norm = post_layer_norm
         self.post_process = post_process
         self.vp_stage = vp_stage
-        self.is_mtp_layer = is_mtp_layer
 
         assert pg_collection is not None, "pg_collection must be provided for MambaStack"
 
@@ -117,7 +114,7 @@ def __init__(
         main_pattern = parsed.main_pattern  # MTP pattern ignored here
 
         self.layer_type_list = allocate_layers(
-            self.config.num_layers if num_layers is None else num_layers,
+            self.config.num_layers if self.config.num_layers is not None else len(main_pattern),
             self.hybrid_attention_ratio,
             self.hybrid_mlp_ratio,
             main_pattern,
@@ -138,7 +135,7 @@ def __init__(
             pg_collection=pg_collection,
             layer_offset=pp_layer_offset,
             residual_in_fp32=residual_in_fp32,
-            is_mtp_layer=is_mtp_layer,
+            is_mtp_layer=False,
         )
 
         # Required for activation recomputation
@@ -275,15 +272,6 @@ def forward(
                             rotary_pos_emb=rotary_pos_emb,
                             sequence_len_offset=sequence_len_offset,
                         )
-                    elif isinstance(layer, MLPLayer):
-                        # MLPLayer (standalone MLP without attention)
-                        hidden_states = layer(
-                            hidden_states=hidden_states,
-                            attention_mask=attention_mask,
-                            inference_context=inference_context,
-                            rotary_pos_emb=rotary_pos_emb,
-                            sequence_len_offset=sequence_len_offset,
-                        )
                     else:  # MambaLayer
                         hidden_states = layer(
                             hidden_states=hidden_states,

From af53f49fd42c9669f9c3ca07ae9d8982d0f747f4 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Tue, 13 Jan 2026 17:38:36 -0800
Subject: [PATCH 33/48] gpt model to not support layer patterns

---
 megatron/core/models/gpt/gpt_model.py         |  3 +-
 .../transformer/multi_token_prediction.py     | 39 ++++---------------
 .../core/transformer/transformer_config.py    |  3 --
 3 files changed, 9 insertions(+), 36 deletions(-)

diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index 1d19082d6eb..d5a06462730 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -100,7 +100,6 @@ def __init__(
         mtp_block_spec: Optional[ModuleSpec] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
         vp_stage: Optional[int] = None,
-        mtp_layer_pattern: str = None,
     ) -> None:
         super().__init__(config=config, pg_collection=pg_collection)
 
@@ -209,7 +208,7 @@ def __init__(
 
         if self.mtp_process:
             self.mtp = MultiTokenPredictionBlock(
-                config=self.config, spec=self.mtp_block_spec, vp_stage=vp_stage, mtp_layer_pattern=mtp_layer_pattern
+                config=self.config, spec=self.mtp_block_spec, vp_stage=vp_stage
             )
 
         # Output
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index a1cf115012b..9ff3d0d585e 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -562,11 +562,10 @@ def __init__(
         self,
         config: TransformerConfig,
         submodules: MultiTokenPredictionLayerSubmodules,
-        mtp_layer_pattern: str = None,
         layer_number: int = 1,
         vp_stage: Optional[int] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
-        # New: For Mamba path - pattern and submodules to build inner layers directly
+        # For Mamba path - pattern and submodules to build inner layers directly
         mtp_layer_pattern: Optional[str] = None,
         mamba_submodules: Optional["MambaStackSubmodules"] = None,
     ):
@@ -577,7 +576,6 @@ def __init__(
         self.vp_stage = vp_stage
         self.cp_group = pg_collection.cp
         self.mtp_layer_pattern = mtp_layer_pattern
-        self.mtp_layer_pattern = mtp_layer_pattern
 
         # Validate attention mask type if using transformer-based inner layers
         if (
@@ -654,34 +652,13 @@ def __init__(
             )
             self.mtp_model_layer = MTPModelLayerContainer(inner_layers, config=config)
         elif self.config.mtp_num_layers is not None:
-            if self.mtp_layer_pattern is not None:
-                # Legacy Mamba path: build MambaStack with override pattern
-                pg_collection = ProcessGroupCollection.use_mpu_process_groups()
-
-                # We do not need pre and post process stage for MTP layer, given they are
-                # handled in the MultiTokenPredictionLayer itself.
-                assert self.config.mtp_num_layers_per_layer is not None, \
-                    "mtp_num_layers_per_layer must be set when using mtp_layer_pattern"
-                self.mtp_model_layer = build_module(
-                    self.submodules.mtp_model_layer,
-                    self.config,
-                    pre_process=False,
-                    post_process=False,
-                    hybrid_override_pattern=self.mtp_layer_pattern,
-                    dtype=self.config.params_dtype,
-                    pg_collection=pg_collection,
-                    vp_stage=self.vp_stage,
-                    num_layers=self.config.mtp_num_layers_per_layer,
-                    is_mtp_layer=True
-                )
-            else:
-                # GPT path: Uses the transformer block spec for MTP layer
-                self.mtp_model_layer = build_module(
-                    self.submodules.mtp_model_layer,
-                    config=self.config,
-                    vp_stage=self.vp_stage,
-                    is_mtp_layer=True
-                )
+            # GPT path: Uses the transformer block spec for MTP layer
+            self.mtp_model_layer = build_module(
+                self.submodules.mtp_model_layer,
+                config=self.config,
+                vp_stage=self.vp_stage,
+                is_mtp_layer=True
+            )
 
         self.final_layernorm = build_module(
             self.submodules.layer_norm,
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index 7c0450bddfe..71885a8a382 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -54,9 +54,6 @@ class TransformerConfig(ModelParallelConfig):
     mtp_use_repeated_layer: bool = False
     """Use a single MTP layer repeatedly instead of multiple separate layers."""
 
-    mtp_num_layers_per_layer: Optional[int] = None
-    """Number of layers inside each MTP layer (e.g., for hybrid models)."""
-
     num_layers_in_first_pipeline_stage: Optional[int] = None
     """Number of transformer layers on first pipeline stage.
     None implies equal layer division across PP ranks."""

From 3611f2a8ed71254c43d7f7435d0568a87e6f78e6 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Tue, 13 Jan 2026 21:50:40 -0800
Subject: [PATCH 34/48] More thorough backward compatibility warnings for the
 args

---
 mamba_builders.py                  | 20 ++++++++--
 megatron/training/arguments.py     | 59 ++++++++++++++++++++++++------
 megatron/training/checkpointing.py |  3 +-
 megatron/training/training.py      | 25 +++++++------
 4 files changed, 80 insertions(+), 27 deletions(-)

diff --git a/mamba_builders.py b/mamba_builders.py
index 29fce24baac..73b8e26265b 100644
--- a/mamba_builders.py
+++ b/mamba_builders.py
@@ -25,8 +25,22 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
     else:
         raise ValueError("You must provide a valid Mamba layer spec via --spec")
 
-    # MTP is now configured via unified pattern in hybrid_override_pattern
-    # e.g., "M*M*/MM/MM" means main="M*M*", mtp="MM", 2 depths
+    # Backward compatibility: if old checkpoint has separate mtp_hybrid_override_pattern,
+    # construct the unified pattern
+    hybrid_override_pattern = args.hybrid_override_pattern
+    if (
+        getattr(args, 'mtp_hybrid_override_pattern', None) is not None
+        and args.mtp_num_layers is not None
+        and args.mtp_num_layers > 0
+        and (hybrid_override_pattern is None or '/' not in hybrid_override_pattern)
+    ):
+        # Old checkpoint format: combine main pattern with MTP pattern
+        main_pattern = hybrid_override_pattern or ''
+        mtp_pattern = args.mtp_hybrid_override_pattern
+        # Build unified pattern: main/mtp/mtp/... (mtp_num_layers times)
+        hybrid_override_pattern = main_pattern + '/' + '/'.join([mtp_pattern] * args.mtp_num_layers)
+        print_rank_0(f"Converted legacy MTP pattern to unified: {hybrid_override_pattern}")
+
     model = MambaModel(
         config=config,
         mamba_stack_spec=mamba_stack_spec,
@@ -35,7 +49,7 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
         pre_process=pre_process,
         hybrid_attention_ratio=args.hybrid_attention_ratio,
         hybrid_mlp_ratio=args.hybrid_mlp_ratio,
-        hybrid_override_pattern=args.hybrid_override_pattern,
+        hybrid_override_pattern=hybrid_override_pattern,
         post_process=post_process,
         fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
         parallel_output=True,
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index fea66b64834..6165b1eeb5a 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -530,24 +530,20 @@ def validate_args(args, defaults={}):
     assert args.global_batch_size > 0
 
     # Deprecation warnings for legacy MTP arguments
-    if args.mtp_layer_pattern is not None:
-        import warnings
-        warnings.warn(
+    if args.mtp_hybrid_override_pattern is not None:
+        warn_rank_0(
             "--mtp-hybrid-override-pattern is deprecated. "
             "For new Mamba+MTP models, use unified --hybrid-override-pattern instead. "
             "Example: 'M*M*/MM/MM' means main='M*M*', MTP pattern='MM' with 2 depths. "
             "This argument is kept only for loading old checkpoints.",
-            DeprecationWarning,
-            stacklevel=2
+            args.rank,
         )
     if args.mtp_spec is not None:
-        import warnings
-        warnings.warn(
+        warn_rank_0(
             "--mtp-spec is deprecated. "
             "For new Mamba+MTP models, use unified --hybrid-override-pattern instead. "
             "This argument is kept only for loading old checkpoints.",
-            DeprecationWarning,
-            stacklevel=2
+            args.rank,
         )
 
     # Uneven virtual pipeline parallelism
@@ -1259,6 +1255,20 @@ def validate_args(args, defaults={}):
                 "when enabling delay_wgrad_compute"
             )
 
+    # For Mamba/hybrid models with unified pattern, infer mtp_num_layers from pattern
+    if args.hybrid_override_pattern and '/' in args.hybrid_override_pattern:
+        from megatron.core.ssm.mamba_hybrid_layer_allocation import parse_hybrid_pattern
+        parsed = parse_hybrid_pattern(args.hybrid_override_pattern)
+        if parsed.mtp_pattern and parsed.mtp_num_depths > 0:
+            inferred_mtp_num_layers = parsed.mtp_num_depths
+            if args.mtp_num_layers is None:
+                args.mtp_num_layers = inferred_mtp_num_layers
+            elif args.mtp_num_layers != inferred_mtp_num_layers:
+                raise ValueError(
+                    f"--mtp-num-layers ({args.mtp_num_layers}) conflicts with "
+                    f"MTP depth count ({inferred_mtp_num_layers}) in pattern '{args.hybrid_override_pattern}'"
+                )
+
     if args.mtp_num_layers:
         assert not args.use_legacy_models, "The legacy Megatron models does not support Multi-Token Prediction (MTP)."
         assert args.position_embedding_type == "rope" or args.position_embedding_type == "none", (
@@ -1266,6 +1276,35 @@ def validate_args(args, defaults={}):
             + f"The supported position embedding types are rope and none."
         )
 
+    # Validate MTP args for hybrid vs non-hybrid models
+    if args.is_hybrid_model:
+        # Mamba/hybrid model MTP validation
+        if args.mtp_num_layers and not (args.hybrid_override_pattern and '/' in args.hybrid_override_pattern):
+            # Hybrid model wants MTP but no unified pattern - check for legacy args
+            if args.mtp_hybrid_override_pattern is None:
+                warn_rank_0(
+                    "Hybrid model with --mtp-num-layers but no MTP pattern. "
+                    "Use unified --hybrid-override-pattern with '/' separator (e.g., 'M*M*/MM/MM') "
+                    "or legacy --mtp-hybrid-override-pattern for old checkpoints.",
+                    args.rank
+                )
+    else:
+        # Non-hybrid (GPT) model MTP validation
+        if args.mtp_hybrid_override_pattern is not None:
+            warn_rank_0(
+                "--mtp-hybrid-override-pattern is for Mamba/hybrid models only. "
+                "For GPT models, MTP replicates the main transformer layer structure. "
+                "This argument will be ignored.",
+                args.rank
+            )
+        if args.mtp_spec is not None:
+            warn_rank_0(
+                "--mtp-spec is for Mamba/hybrid models only. "
+                "For GPT models, use the standard layer spec. "
+                "This argument will be ignored.",
+                args.rank
+            )
+
     if args.cpu_offloading_num_layers > 0:
         args.cpu_offloading = True
 
@@ -1807,8 +1846,6 @@ def _add_network_size_args(parser):
                        help='Use a single MTP layer repeatedly instead of multiple separate layers. '
                        'This is more parameter-efficient. When enabled, only 1 MTP layer is created '
                        'and applied --mtp-num-layers times.')
-    group.add_argument('--mtp-num-layers-per-layer', type=int, default=None,
-                       help='Number of layers inside each MTP layer (e.g., for hybrid models).')
     group.add_argument('--moe-latent-size', type=int, default=None,
                        help='Latent projection dimension for MoE. If None, MoE latent projections are not used.')
     return parser
diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index d85f6e5767c..a46604ed6fd 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -1347,8 +1347,7 @@ def _set_arg(arg_name, old_arg_name=None, force=False):
     _set_arg('hidden_dropout', force=True)
 
     _set_arg('hybrid_override_pattern', force=True)
-    _set_arg('mtp_layer_pattern', old_arg_name='mtp_hybrid_override_pattern', force=True)  # old name
-    _set_arg('mtp_layer_pattern', force=True)  # new mtp name, takes precedence
+    _set_arg('mtp_hybrid_override_pattern', force=True)  # Legacy MTP pattern for old checkpoints
     _set_arg('spec', force=True)
     _set_arg('hybrid_attention_ratio', force=True)
     _set_arg('hybrid_mlp_ratio', force=True)
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 2484825b3a2..1716c5a733c 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -174,15 +174,20 @@ def num_floating_point_operations(args, batch_size):
     def calculate_layer_counts():
         """Calculate the number of attention, Mamba, and MLP layers."""
         if args.hybrid_override_pattern:
-            counts = {'M': 0, '*': 0, '-': 0, 'E':0}
-            for layer_type in args.hybrid_override_pattern:
-                if layer_type in counts:
-                    counts[layer_type] += 1
-            if args.mtp_layer_pattern:
-                assert args.mtp_num_layers is not None
-                for layer_type in args.mtp_layer_pattern:
+            from megatron.core.ssm.mamba_hybrid_layer_allocation import parse_hybrid_pattern
+            # Parse unified pattern to separate main and MTP components
+            parsed = parse_hybrid_pattern(args.hybrid_override_pattern)
+            counts = {'M': 0, '*': 0, '-': 0, 'E': 0}
+            # Count main decoder layers
+            if parsed.main_pattern:
+                for layer_type in parsed.main_pattern:
                     if layer_type in counts:
-                        counts[layer_type] += args.mtp_num_layers
+                        counts[layer_type] += 1
+            # Count MTP layers (pattern repeated mtp_num_depths times)
+            if parsed.mtp_pattern and parsed.mtp_num_depths > 0:
+                for layer_type in parsed.mtp_pattern:
+                    if layer_type in counts:
+                        counts[layer_type] += parsed.mtp_num_depths
             return counts['*'], counts['M'], counts['-'], counts['E']
         else:
             num_attn_layers = round(args.num_layers * args.hybrid_attention_ratio)
@@ -259,10 +264,8 @@ def hybrid_flops(batch_size, seq_len, hidden_size,
                      mlp_expansion=4.0, swiglu=False,
                      moe_latent_size=None,
                      moe_ffn_hidden_size=2048, shared_expert_ffn_hidden_size=2048, num_experts_routed_to=1,
-                     vocab_size=256000, mtp_num_layers=None):
+                     vocab_size=256000, mtp_num_layers=0):
         """Calculate total FLOPs for the hybrid model."""
-        if mtp_num_layers is None:
-            mtp_num_layers = 0
         flops_fwd = (
                 num_attn_layers * attn_layer_flops(batch_size, seq_len, hidden_size,
                                                    num_attn_heads, gqa, gqa_groups, kv_channels) +

From d81b58a9c90930aa0c5e4a3a24c0425cd5140bfe Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Wed, 14 Jan 2026 14:54:14 -0800
Subject: [PATCH 35/48] Remove addition of cp_group_size to mtp loss

---
 megatron/core/models/mamba/mamba_model.py     |  8 ++------
 megatron/core/pipeline_parallel/schedules.py  |  4 ++--
 .../core/ssm/mamba_hybrid_layer_allocation.py |  1 -
 .../transformer/multi_token_prediction.py     | 19 ++++++++++++++-----
 .../test_multi_token_prediction.py            |  8 ++++----
 5 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 90050a9756e..f0a68eba0a4 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -162,8 +162,6 @@ def __init__(
             from megatron.core.models.mamba.mamba_layer_specs import get_mamba_mtp_block_spec
 
             mtp_block_spec = get_mamba_mtp_block_spec()
-
-            # Extract mamba_submodules from mamba_stack_spec for MTP to use
             mamba_submodules = mamba_stack_spec.submodules
 
             self.mtp = MultiTokenPredictionBlock(
@@ -298,13 +296,11 @@ def forward(
             rotary_pos_emb=rotary_pos_emb,
         )
 
-        # Get output weight for shared embeddings
         output_weight = None
         if self.share_embeddings_and_output_weights:
-            output_weight = self.shared_embedding_or_output_weight()            
+            output_weight = self.shared_embedding_or_output_weight()
 
         if self.mtp_process:
-
             hidden_states = self.mtp(
                 input_ids=input_ids,
                 position_ids=position_ids,
@@ -312,6 +308,7 @@ def forward(
                 attention_mask=attention_mask,
                 inference_params=inference_params,
                 rotary_pos_emb=rotary_pos_emb,
+                packed_seq_params=packed_seq_params,
                 embedding=self.embedding,
             )
 
@@ -330,7 +327,6 @@ def forward(
                 compute_language_model_loss=self.compute_language_model_loss,
             )
 
-        
         sequence_parallel_override = False
         if in_inference_mode and inference_context.materialize_only_last_token_logits:
             if inference_context.is_static_batching():
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index 1222e760eb4..4133e72ab17 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -292,11 +292,11 @@ def forward_step_calc_loss(
             if config.grad_scale_func is not None
             else torch.ones(1, device=output_tensor.device)
         )
-        # Set the loss scale (multiply by cp_group_size to match MoE scaling)
+        # Set the loss scale
         if config.calculate_per_token_loss:
             MTPLossAutoScaler.set_loss_scale(loss_scale)
         else:
-            MTPLossAutoScaler.set_loss_scale(loss_scale * cp_group_size / num_microbatches)
+            MTPLossAutoScaler.set_loss_scale(loss_scale / num_microbatches)
 
     return output_tensor, num_tokens
 
diff --git a/megatron/core/ssm/mamba_hybrid_layer_allocation.py b/megatron/core/ssm/mamba_hybrid_layer_allocation.py
index 10aea459d36..146bc41390d 100644
--- a/megatron/core/ssm/mamba_hybrid_layer_allocation.py
+++ b/megatron/core/ssm/mamba_hybrid_layer_allocation.py
@@ -94,7 +94,6 @@ def parse_hybrid_pattern(pattern: Optional[str]) -> ParsedHybridPattern:
     if pattern is None:
         return ParsedHybridPattern(main_pattern=None, mtp_pattern=None, mtp_num_depths=0)
 
-    # Split by MTP separator
     parts = pattern.split(Symbols.MTP_SEPARATOR)
 
     if len(parts) == 1:
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index fc40307be46..9d8801640bb 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -538,6 +538,20 @@ class MTPModelLayerContainer(MegatronModule):
     def __init__(self, layers: torch.nn.ModuleList, config: TransformerConfig):
         super().__init__(config=config)
         self.layers = layers
+        self._current_microbatch = 0
+
+    @property
+    def current_microbatch(self):
+        """Get current microbatch index for CUDA graph management."""
+        return self._current_microbatch
+
+    @current_microbatch.setter
+    def current_microbatch(self, value):
+        """Set current microbatch and propagate to inner layers for CUDA graphs."""
+        self._current_microbatch = value
+        for layer in self.layers:
+            if hasattr(layer, 'current_microbatch'):
+                layer.current_microbatch = value
 
     def set_input_tensor(self, input_tensor: Tensor):
         """Set input tensor for pipeline parallelism compatibility."""
@@ -1174,17 +1188,12 @@ def __init__(
         mamba_submodules: Optional["MambaStackSubmodules"] = None,
     ):
         super().__init__(config=config)
-        self.mtp_layer_pattern = mtp_layer_pattern
         self.submodules = _get_mtp_block_submodules(config, spec)
         self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor
         self.vp_stage = vp_stage
-
-        # Store new parameters for pattern-based layer building
         self.mtp_layer_pattern = mtp_layer_pattern
         self.mtp_num_depths = mtp_num_depths
         self.mamba_submodules = mamba_submodules
-
-        # Use config.mtp_use_repeated_layer for both GPT and Mamba paths
         self.mtp_use_repeated_layer = self.config.mtp_use_repeated_layer
 
         vp_size = config.virtual_pipeline_model_parallel_size
diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index 9c21e2a51cc..9e9a79521b2 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -164,7 +164,7 @@ def model_provider(
             config=config,
             transformer_layer_spec=transformer_layer_spec,
             mtp_block_spec=mtp_block_spec,
-            vocab_size=args.vocal_size,
+            vocab_size=args.vocab_size,
             max_sequence_length=args.max_position_embeddings,
             pre_process=pre_process,
             post_process=post_process,
@@ -188,7 +188,7 @@ def create_test_args(
         args.num_layers = 2
         args.mtp_num_layers = 2
         args.mtp_loss_scaling_factor = 0.1
-        args.vocal_size = 128800
+        args.vocab_size = 128800
         args.hidden_size = 128
         args.num_attention_heads = 8
         args.max_position_embeddings = 256
@@ -723,7 +723,7 @@ def model_provider(
         model = MambaModel(
             config=config,
             mamba_stack_spec=mamba_stack_spec,
-            vocab_size=args.vocal_size,
+            vocab_size=args.vocab_size,
             max_sequence_length=args.max_position_embeddings,
             pre_process=pre_process,
             post_process=post_process,
@@ -749,7 +749,7 @@ def create_test_args(
         args.num_layers = 4
         args.mtp_num_layers = 2
         args.mtp_loss_scaling_factor = 0.1
-        args.vocal_size = 128800
+        args.vocab_size = 128800
         args.hidden_size = 128
         args.num_attention_heads = 8
         args.num_query_groups = 8

From a8a8eae216593db5738fda2ab81eabcfbb10577b Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Thu, 15 Jan 2026 21:22:40 -0800
Subject: [PATCH 36/48] Refactor arg parsing to arguments.py

---
 mamba_builders.py                             |  18 +--
 .../transformer/multi_token_prediction.py     |   1 +
 megatron/training/arguments.py                | 118 ++++++++++--------
 megatron/training/checkpointing.py            |   7 +-
 4 files changed, 76 insertions(+), 68 deletions(-)

diff --git a/mamba_builders.py b/mamba_builders.py
index 73b8e26265b..00e2ac953b5 100644
--- a/mamba_builders.py
+++ b/mamba_builders.py
@@ -25,22 +25,6 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
     else:
         raise ValueError("You must provide a valid Mamba layer spec via --spec")
 
-    # Backward compatibility: if old checkpoint has separate mtp_hybrid_override_pattern,
-    # construct the unified pattern
-    hybrid_override_pattern = args.hybrid_override_pattern
-    if (
-        getattr(args, 'mtp_hybrid_override_pattern', None) is not None
-        and args.mtp_num_layers is not None
-        and args.mtp_num_layers > 0
-        and (hybrid_override_pattern is None or '/' not in hybrid_override_pattern)
-    ):
-        # Old checkpoint format: combine main pattern with MTP pattern
-        main_pattern = hybrid_override_pattern or ''
-        mtp_pattern = args.mtp_hybrid_override_pattern
-        # Build unified pattern: main/mtp/mtp/... (mtp_num_layers times)
-        hybrid_override_pattern = main_pattern + '/' + '/'.join([mtp_pattern] * args.mtp_num_layers)
-        print_rank_0(f"Converted legacy MTP pattern to unified: {hybrid_override_pattern}")
-
     model = MambaModel(
         config=config,
         mamba_stack_spec=mamba_stack_spec,
@@ -49,7 +33,7 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
         pre_process=pre_process,
         hybrid_attention_ratio=args.hybrid_attention_ratio,
         hybrid_mlp_ratio=args.hybrid_mlp_ratio,
-        hybrid_override_pattern=hybrid_override_pattern,
+        hybrid_override_pattern=args.hybrid_override_pattern,
         post_process=post_process,
         fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
         parallel_output=True,
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 9d8801640bb..d84f7566e3b 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -14,6 +14,7 @@
 from megatron.core.fp8_utils import get_fp8_context
 from megatron.core.models.backends import BackendSpecProvider, LocalSpecProvider
 from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.pipeline_parallel.utils import is_vp_last_stage
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel import (
     gather_from_tensor_model_parallel_region,
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index d19fbf71715..8bcc1597290 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -554,6 +554,74 @@ def validate_args(args, defaults={}):
             args.rank,
         )
 
+    # === MTP validation ===
+    # Backward compatibility: convert legacy mtp_hybrid_override_pattern to unified format
+    from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols, parse_hybrid_pattern
+    sep = Symbols.MTP_SEPARATOR
+    if (
+        getattr(args, 'mtp_hybrid_override_pattern', None) is not None
+        and args.mtp_num_layers is not None
+        and args.mtp_num_layers > 0
+        and (args.hybrid_override_pattern is None or sep not in args.hybrid_override_pattern)
+    ):
+        main_pattern = args.hybrid_override_pattern or ''
+        mtp_pattern = args.mtp_hybrid_override_pattern
+        args.hybrid_override_pattern = main_pattern + sep + sep.join([mtp_pattern] * args.mtp_num_layers)
+        args.mtp_num_layers = None
+        args.mtp_hybrid_override_pattern = None
+        print_rank_0(f"Converted legacy MTP pattern to unified: {args.hybrid_override_pattern}")
+
+    # Infer mtp_num_layers from unified pattern
+    if args.hybrid_override_pattern and sep in args.hybrid_override_pattern:
+        parsed = parse_hybrid_pattern(args.hybrid_override_pattern)
+        if parsed.mtp_pattern and parsed.mtp_num_depths > 0:
+            inferred_mtp_num_layers = parsed.mtp_num_depths
+            if args.mtp_num_layers is None:
+                args.mtp_num_layers = inferred_mtp_num_layers
+            elif args.mtp_num_layers != inferred_mtp_num_layers:
+                raise ValueError(
+                    f"--mtp-num-layers ({args.mtp_num_layers}) conflicts with "
+                    f"MTP depth count ({inferred_mtp_num_layers}) in pattern '{args.hybrid_override_pattern}'"
+                )
+
+    # MTP validation
+    if args.mtp_num_layers:
+        assert not args.use_legacy_models, "The legacy Megatron models does not support Multi-Token Prediction (MTP)."
+        assert args.position_embedding_type == "rope" or args.position_embedding_type == "none", (
+            f"Multi-Token Prediction (MTP) is not supported with {args.position_embedding_type} position embedding type."
+            + f"The supported position embedding types are rope and none."
+        )
+
+    # Validate MTP args for hybrid vs non-hybrid models
+    if args.is_hybrid_model:
+        # Mamba/hybrid model MTP validation
+        if args.mtp_num_layers and not (args.hybrid_override_pattern and sep in args.hybrid_override_pattern):
+            # Hybrid model wants MTP but no unified pattern - check for legacy args
+            if args.mtp_hybrid_override_pattern is None:
+                warn_rank_0(
+                    "Hybrid model with --mtp-num-layers but no MTP pattern. "
+                    "Use unified --hybrid-override-pattern with '/' separator (e.g., 'M*M*/MM/MM') "
+                    "or legacy --mtp-hybrid-override-pattern for old checkpoints.",
+                    args.rank
+                )
+    else:
+        # Non-hybrid (GPT) model MTP validation
+        if args.mtp_hybrid_override_pattern is not None:
+            warn_rank_0(
+                "--mtp-hybrid-override-pattern is for Mamba/hybrid models only. "
+                "For GPT models, MTP replicates the main transformer layer structure. "
+                "This argument will be ignored.",
+                args.rank
+            )
+        if args.mtp_spec is not None:
+            warn_rank_0(
+                "--mtp-spec is for Mamba/hybrid models only. "
+                "For GPT models, use the standard layer spec. "
+                "This argument will be ignored.",
+                args.rank
+            )
+    # === End of MTP validation ===
+    
     # Uneven virtual pipeline parallelism
     assert (
         int(args.num_layers_per_virtual_pipeline_stage is not None)
@@ -1279,56 +1347,6 @@ def validate_args(args, defaults={}):
                 "when enabling delay_wgrad_compute"
             )
 
-    # For Mamba/hybrid models with unified pattern, infer mtp_num_layers from pattern
-    if args.hybrid_override_pattern and '/' in args.hybrid_override_pattern:
-        from megatron.core.ssm.mamba_hybrid_layer_allocation import parse_hybrid_pattern
-        parsed = parse_hybrid_pattern(args.hybrid_override_pattern)
-        if parsed.mtp_pattern and parsed.mtp_num_depths > 0:
-            inferred_mtp_num_layers = parsed.mtp_num_depths
-            if args.mtp_num_layers is None:
-                args.mtp_num_layers = inferred_mtp_num_layers
-            elif args.mtp_num_layers != inferred_mtp_num_layers:
-                raise ValueError(
-                    f"--mtp-num-layers ({args.mtp_num_layers}) conflicts with "
-                    f"MTP depth count ({inferred_mtp_num_layers}) in pattern '{args.hybrid_override_pattern}'"
-                )
-
-    if args.mtp_num_layers:
-        assert not args.use_legacy_models, "The legacy Megatron models does not support Multi-Token Prediction (MTP)."
-        assert args.position_embedding_type == "rope" or args.position_embedding_type == "none", (
-            f"Multi-Token Prediction (MTP) is not supported with {args.position_embedding_type} position embedding type."
-            + f"The supported position embedding types are rope and none."
-        )
-
-    # Validate MTP args for hybrid vs non-hybrid models
-    if args.is_hybrid_model:
-        # Mamba/hybrid model MTP validation
-        if args.mtp_num_layers and not (args.hybrid_override_pattern and '/' in args.hybrid_override_pattern):
-            # Hybrid model wants MTP but no unified pattern - check for legacy args
-            if args.mtp_hybrid_override_pattern is None:
-                warn_rank_0(
-                    "Hybrid model with --mtp-num-layers but no MTP pattern. "
-                    "Use unified --hybrid-override-pattern with '/' separator (e.g., 'M*M*/MM/MM') "
-                    "or legacy --mtp-hybrid-override-pattern for old checkpoints.",
-                    args.rank
-                )
-    else:
-        # Non-hybrid (GPT) model MTP validation
-        if args.mtp_hybrid_override_pattern is not None:
-            warn_rank_0(
-                "--mtp-hybrid-override-pattern is for Mamba/hybrid models only. "
-                "For GPT models, MTP replicates the main transformer layer structure. "
-                "This argument will be ignored.",
-                args.rank
-            )
-        if args.mtp_spec is not None:
-            warn_rank_0(
-                "--mtp-spec is for Mamba/hybrid models only. "
-                "For GPT models, use the standard layer spec. "
-                "This argument will be ignored.",
-                args.rank
-            )
-
     if args.cpu_offloading_num_layers > 0:
         args.cpu_offloading = True
 
diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index b6d2f742494..3be427fb4e2 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -1356,7 +1356,12 @@ def _set_arg(arg_name, old_arg_name=None, force=False):
     _set_arg('hidden_dropout', force=True)
 
     _set_arg('hybrid_override_pattern', force=True)
-    _set_arg('mtp_hybrid_override_pattern', force=True)  # Legacy MTP pattern for old checkpoints
+
+    # Legacy MTP pattern for old checkpoints
+    _set_arg('mtp_hybrid_override_pattern', force=True)
+    _set_arg('mtp_num_layers', force=True)
+    _set_arg('mtp_use_repeated_layer', force=True)
+
     _set_arg('spec', force=True)
     _set_arg('hybrid_attention_ratio', force=True)
     _set_arg('hybrid_mlp_ratio', force=True)

From 590315f0e0aaa225b6e7046de7d0b8b3cf5f6f05 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Tue, 20 Jan 2026 15:31:34 -0800
Subject: [PATCH 37/48] Handle mtp_standalone case, TELayerNormColumnParallel
 bug ..

---
 mamba_builders.py                             |  6 +-
 megatron/core/models/gpt/gpt_model.py         |  5 +-
 .../core/models/mamba/mamba_layer_specs.py    | 74 +++++++------------
 megatron/core/models/mamba/mamba_model.py     | 21 ++++--
 megatron/core/ssm/mamba_block.py              |  9 +--
 megatron/training/arguments.py                |  5 +-
 6 files changed, 51 insertions(+), 69 deletions(-)

diff --git a/mamba_builders.py b/mamba_builders.py
index 00e2ac953b5..5d31af60475 100644
--- a/mamba_builders.py
+++ b/mamba_builders.py
@@ -17,9 +17,9 @@ def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, p
 
     if config.transformer_impl == "inference_optimized":
         mamba_stack_spec = mamba_inference_stack_spec
-        assert not config.inference_fuse_tp_communication, (
-            "inference_fuse_tp_communication is not supported for Mamba"
-        )
+        assert (
+            not config.inference_fuse_tp_communication
+        ), "inference_fuse_tp_communication is not supported for Mamba"
     elif args.spec is not None:
         mamba_stack_spec = import_module(args.spec)
     else:
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index ae405a9aae6..9661bba9c8c 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -26,6 +26,7 @@
     MTPLossAutoScaler,
     MTPLossLoggingHelper,
     MultiTokenPredictionBlock,
+    mtp_on_this_rank,
     roll_tensor,
 )
 from megatron.core.transformer.spec_utils import ModuleSpec
@@ -139,7 +140,9 @@ def __init__(
             self.rotary_base = rotary_base
         self.rotary_scaling = rope_scaling
         self.mtp_block_spec = mtp_block_spec
-        self.mtp_process = mtp_block_spec is not None and self.post_process
+        self.mtp_process = mtp_block_spec is not None and mtp_on_this_rank(
+            self.config, vp_stage=vp_stage
+        )
 
         if self.pre_process or self.mtp_process:
             self.embedding = LanguageModelEmbedding(
diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index ef217cb963f..0fbaf9b2856 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -28,6 +28,7 @@
     MultiTokenPredictionLayerSubmodules,
 )
 
+# This should be private and should not be used outside of this file.
 moe = get_moe_module_spec(
     use_te=True,
     num_experts=8,  # Can be any positive integer (must not be None).
@@ -35,6 +36,29 @@
     moe_use_legacy_grouped_gemm=False,
 )
 
+
+# MTP block spec for Mamba - provides norms and projection only.
+# Inner layers are built by MultiTokenPredictionLayer using the shared
+# layer_builder with mtp_layer_pattern and mamba_submodules.
+_mamba_mtp_block_spec = ModuleSpec(
+    module=MultiTokenPredictionBlock,
+    submodules=MultiTokenPredictionBlockSubmodules(
+        layer_specs=[
+            ModuleSpec(
+                module=MultiTokenPredictionLayer,
+                submodules=MultiTokenPredictionLayerSubmodules(
+                    enorm=TENorm,
+                    hnorm=TENorm,
+                    eh_proj=TEColumnParallelLinear,
+                    mtp_model_layer=None,  # Built via pattern + mamba_submodules
+                    layer_norm=TENorm,
+                ),
+            )
+        ]
+    ),
+)
+
+
 mamba_stack_spec = ModuleSpec(
     module=MambaStack,
     submodules=MambaStackSubmodules(
@@ -90,6 +114,7 @@
                 pre_mlp_layernorm=TENorm, mlp=moe, mlp_bda=get_bias_dropout_add
             ),
         ),
+        mtp_block_spec=_mamba_mtp_block_spec,
     ),
 )
 
@@ -151,53 +176,6 @@
                 pre_mlp_layernorm=TENorm, mlp=moe, mlp_bda=get_bias_dropout_add
             ),
         ),
+        mtp_block_spec=_mamba_mtp_block_spec,
     ),
 )
-
-
-def get_mamba_mtp_block_spec(use_te: bool = True) -> ModuleSpec:
-    """MTP block spec for Mamba using unified pattern syntax.
-
-    This spec provides norms and projection only - inner layers are built
-    by MultiTokenPredictionLayer using the shared layer_builder with
-    mtp_layer_pattern and mamba_submodules passed from MambaModel.
-
-    The number of MTP depths is determined by the parsed unified pattern
-    (e.g., "M*M*/MM/MM" -> main="M*M*", mtp="MM", 2 depths).
-
-    Args:
-        use_te: Whether to use TransformerEngine modules (default: True)
-
-    Returns:
-        ModuleSpec for MultiTokenPredictionBlock
-
-    Example:
-        >>> mtp_spec = get_mamba_mtp_block_spec()
-        >>> mtp_block = MultiTokenPredictionBlock(
-        ...     config=config,
-        ...     spec=mtp_spec,
-        ...     mtp_layer_pattern="MM",
-        ...     mtp_num_depths=2,
-        ...     mamba_submodules=mamba_stack_spec.submodules,
-        ... )
-    """
-    norm = TENorm if use_te else TENorm  # Fallback to TENorm for now
-    linear = TELayerNormColumnParallelLinear if use_te else TELayerNormColumnParallelLinear
-
-    return ModuleSpec(
-        module=MultiTokenPredictionBlock,
-        submodules=MultiTokenPredictionBlockSubmodules(
-            layer_specs=[
-                ModuleSpec(
-                    module=MultiTokenPredictionLayer,
-                    submodules=MultiTokenPredictionLayerSubmodules(
-                        enorm=norm,
-                        hnorm=norm,
-                        eh_proj=linear,
-                        mtp_model_layer=None,  # Built via pattern + mamba_submodules
-                        layer_norm=norm,
-                    ),
-                )
-            ]
-        ),
-    )
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index f0a68eba0a4..5e483a1494e 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -16,7 +16,10 @@
 from megatron.core.tensor_parallel import gather_from_sequence_parallel_region
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.enums import ModelType
-from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionBlock
+from megatron.core.transformer.multi_token_prediction import (
+    mtp_on_this_rank,
+    MultiTokenPredictionBlock,
+)
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.utils import (
     WrappedTensor,
@@ -114,7 +117,7 @@ def __init__(
 
         # Determine if MTP is needed (based on pattern parsing)
         self.mtp_process = (
-            self.mtp_pattern is not None and self.mtp_num_depths > 0 and self.post_process
+            self.mtp_pattern is not None and self.mtp_num_depths > 0 and mtp_on_this_rank(self.config, vp_stage=self.vp_stage)
         )
 
         # Cache for RoPE tensors which do not change between iterations.
@@ -151,18 +154,20 @@ def __init__(
             pre_process=self.pre_process,
             hybrid_attention_ratio=self.hybrid_attention_ratio,
             hybrid_mlp_ratio=self.hybrid_mlp_ratio,
-            hybrid_override_pattern=self.hybrid_override_pattern,
+            hybrid_override_pattern=parsed.main_pattern,
             post_process=self.post_process,
             dtype=config.params_dtype,
             pg_collection=self.pg_collection,
         )
 
-        # MTP block (MambaModel creates it, MTP builds its own inner layers)
+        # MTP block - uses mtp_block_spec from mamba_stack_spec.submodules
         if self.mtp_process:
-            from megatron.core.models.mamba.mamba_layer_specs import get_mamba_mtp_block_spec
-
-            mtp_block_spec = get_mamba_mtp_block_spec()
             mamba_submodules = mamba_stack_spec.submodules
+            mtp_block_spec = mamba_submodules.mtp_block_spec
+            assert mtp_block_spec is not None, (
+                "MTP pattern specified but mtp_block_spec is None in mamba_stack_spec.submodules. "
+                "Ensure mamba_stack_spec includes mtp_block_spec for MTP support."
+            )
 
             self.mtp = MultiTokenPredictionBlock(
                 config=self.config,
@@ -189,7 +194,7 @@ def __init__(
                 tp_group=self.pg_collection.tp,
             )
 
-        if self.pre_process or self.post_process:
+        if self.pre_process or self.post_process or self.mtp_process:
             self.setup_embeddings_and_output_layer()
 
         for name, module in self.named_modules():
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index ba5ecf287d6..f015f0dab4b 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -45,6 +45,7 @@ class MambaStackSubmodules:
     attention_layer: Union[ModuleSpec, type] = IdentityOp
     mlp_layer: Union[ModuleSpec, type] = IdentityOp
     moe_layer: Union[ModuleSpec, type] = IdentityOp
+    mtp_block_spec: Optional[ModuleSpec] = None
 
 
 class MambaStack(MegatronModule):
@@ -110,15 +111,11 @@ def __init__(
         self.hybrid_override_pattern = hybrid_override_pattern
         self.pg_collection = pg_collection
 
-        # Parse pattern - extract only main pattern, MTP handled by MambaModel
-        parsed = parse_hybrid_pattern(hybrid_override_pattern)
-        main_pattern = parsed.main_pattern  # MTP pattern ignored here
-
         self.layer_type_list = allocate_layers(
-            self.config.num_layers if self.config.num_layers is not None else len(main_pattern),
+            self.config.num_layers if self.config.num_layers is not None else len(self.hybrid_override_pattern),
             self.hybrid_attention_ratio,
             self.hybrid_mlp_ratio,
-            main_pattern,
+            self.hybrid_override_pattern,
         )
 
         pp_layer_offset = 0
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 8bcc1597290..174c25255c5 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -537,6 +537,7 @@ def validate_args(args, defaults={}):
                 args.global_batch_size), flush=True)
     assert args.global_batch_size > 0
 
+    # === MTP validation ===
     # Deprecation warnings for legacy MTP arguments
     if args.mtp_hybrid_override_pattern is not None:
         warn_rank_0(
@@ -554,7 +555,6 @@ def validate_args(args, defaults={}):
             args.rank,
         )
 
-    # === MTP validation ===
     # Backward compatibility: convert legacy mtp_hybrid_override_pattern to unified format
     from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols, parse_hybrid_pattern
     sep = Symbols.MTP_SEPARATOR
@@ -567,7 +567,6 @@ def validate_args(args, defaults={}):
         main_pattern = args.hybrid_override_pattern or ''
         mtp_pattern = args.mtp_hybrid_override_pattern
         args.hybrid_override_pattern = main_pattern + sep + sep.join([mtp_pattern] * args.mtp_num_layers)
-        args.mtp_num_layers = None
         args.mtp_hybrid_override_pattern = None
         print_rank_0(f"Converted legacy MTP pattern to unified: {args.hybrid_override_pattern}")
 
@@ -1896,7 +1895,7 @@ def _add_network_size_args(parser):
                        help='Untie embeddings and output weights.')
     group.add_argument('--multi-latent-attention', action='store_true',
                        help='Use multi-latent attention for model.')
-    group.add_argument('--mtp-num-layers', type=int, default=None,
+    group.add_argument('--mtp-num-layers', type=int, default=0,
                        help='Number of Multi-Token Prediction (MTP) Layers.'
                        'MTP extends the prediction scope to multiple future tokens at each position.'
                        'This MTP implementation sequentially predict additional tokens '

From 0520329ba57894cede45d6ff8f7944b44a539756 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Wed, 21 Jan 2026 09:57:00 -0800
Subject: [PATCH 38/48] Address review comments

---
 megatron/core/models/mamba/mamba_layer_specs.py  |  1 +
 .../core/transformer/multi_token_prediction.py   | 16 ++++++----------
 megatron/core/transformer/transformer_layer.py   | 13 ++++++++++---
 megatron/training/arguments.py                   |  6 +++---
 4 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index 0fbaf9b2856..c7f9f0bc163 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -5,6 +5,7 @@
     TELayerNormColumnParallelLinear,
     TENorm,
     TERowParallelLinear,
+    TEColumnParallelLinear,
 )
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index d84f7566e3b..90d8c09ea6b 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -31,7 +31,6 @@
     from megatron.core.ssm.mamba_block import MambaStackSubmodules
 from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
 from megatron.core.transformer.transformer_config import TransformerConfig
-from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
 from megatron.core.utils import (
     get_pg_rank,
     is_torch_min_version,
@@ -532,8 +531,6 @@ class MTPModelLayerContainer(MegatronModule):
     from a pattern string. It provides proper state dict key structure:
     `mtp_model_layer.layers.{i}.*` for checkpoint compatibility.
 
-    This is used when MTP builds its own layers using the shared layer_builder,
-    instead of receiving a pre-built MambaStack.
     """
 
     def __init__(self, layers: torch.nn.ModuleList, config: TransformerConfig):
@@ -724,10 +721,9 @@ def __init__(
             tp_comm_buffer_name="mtp_eh_proj"
         )
 
-        # Build inner layers: three possible paths
+        # Build inner layers: two possible paths
         # 1. New Mamba path: build from pattern using shared layer_builder
-        # 2. Legacy Mamba path: build MambaStack with mtp_layer_pattern
-        # 3. GPT path: single TransformerLayer
+        # 2. GPT path: single TransformerLayer
         if mtp_layer_pattern is not None and mamba_submodules is not None:
             # New Mamba path: build inner layers from pattern using shared layer_builder
             from megatron.core.ssm.layer_builder import build_layers_from_pattern
@@ -743,14 +739,14 @@ def __init__(
             self.mtp_model_layer = MTPModelLayerContainer(inner_layers, config=config)
         elif self.config.mtp_num_layers is not None:
             # GPT path: Uses the transformer block spec for MTP layer
-            diff_transformer_layer_offset = self.config.num_layers - get_transformer_layer_offset(
-                self.config, vp_stage
-            )
+            # MTP inner layers use their own layer numbering (self.layer_number = 1, 2, etc.)
+            # rather than continuing from decoder layer numbers. This is consistent with the
+            # Mamba path and ensures proper aux loss tracking in router.py.
             self.mtp_model_layer = build_module(
                 self.submodules.mtp_model_layer,
                 config=self.config,
                 vp_stage=self.vp_stage,
-                layer_number=self.layer_number + diff_transformer_layer_offset,
+                layer_number=self.layer_number,
                 is_mtp_layer=True
             )
 
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index ed466a14f1a..eebd8101cfe 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -277,9 +277,16 @@ def __init__(
         self.tp_group = pg_collection.tp
 
         self.submodules_config = submodules
-        self.layer_number = layer_number + get_transformer_layer_offset(
-            self.config, vp_stage, get_pg_rank(pg_collection.pp)
-        )
+        # MTP inner layers use their own layer numbering (starting from 1 within each MTP depth),
+        # so they should NOT add the decoder layer offset. The router.py handles MTP layer
+        # numbering separately by adding config.num_layers to distinguish MTP layers from decoder
+        # layers in the aux loss tracker.
+        if is_mtp_layer:
+            self.layer_number = layer_number
+        else:
+            self.layer_number = layer_number + get_transformer_layer_offset(
+                self.config, vp_stage, get_pg_rank(pg_collection.pp)
+            )
         self.hidden_dropout = config.hidden_dropout if hidden_dropout is None else hidden_dropout
         self.is_mtp_layer = is_mtp_layer
 
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 174c25255c5..4e14ca9d253 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -542,7 +542,7 @@ def validate_args(args, defaults={}):
     if args.mtp_hybrid_override_pattern is not None:
         warn_rank_0(
             "--mtp-hybrid-override-pattern is deprecated. "
-            "For new Mamba+MTP models, use unified --hybrid-override-pattern instead. "
+            "For new hybrid models with MTP models, use unified --hybrid-override-pattern instead. "
             "Example: 'M*M*/MM/MM' means main='M*M*', MTP pattern='MM' with 2 depths. "
             "This argument is kept only for loading old checkpoints.",
             args.rank,
@@ -550,7 +550,7 @@ def validate_args(args, defaults={}):
     if args.mtp_spec is not None:
         warn_rank_0(
             "--mtp-spec is deprecated. "
-            "For new Mamba+MTP models, use unified --hybrid-override-pattern instead. "
+            "For new hybrid models with MTP models, use unified --hybrid-override-pattern instead. "
             "This argument is kept only for loading old checkpoints.",
             args.rank,
         )
@@ -3505,7 +3505,7 @@ def _add_experimental_args(parser):
                             'of layer in the override pattern must match number in'
                             'the overidden pattern')
     # DEPRECATED: Legacy MTP arguments kept for checkpoint compatibility.
-    # For new Mamba+MTP models, use unified pattern in --hybrid-override-pattern instead.
+    # For new hybrid models with MTP models, use unified pattern in --hybrid-override-pattern instead.
     # Example: "M*M*/MM/MM" means main="M*M*", MTP pattern="MM" with 2 depths.
     group.add_argument("--mtp-hybrid-override-pattern", type=str, default=None,
                        help='DEPRECATED: Use unified --hybrid-override-pattern instead. '

From 2bfa108a37096c3ff9c4c1a6609a7a1c8773e9c1 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Wed, 21 Jan 2026 12:11:23 -0800
Subject: [PATCH 39/48] standalone broken because loss not on non-mtp stage

---
 megatron/core/models/gpt/gpt_model.py         |  10 +-
 megatron/core/models/mamba/mamba_model.py     |   8 +-
 .../transformer/multi_token_prediction.py     | 158 +++++++++---------
 megatron/training/arguments.py                |   7 +-
 4 files changed, 96 insertions(+), 87 deletions(-)

diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index 9661bba9c8c..7048c17e26a 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -27,6 +27,7 @@
     MTPLossLoggingHelper,
     MultiTokenPredictionBlock,
     mtp_on_this_rank,
+    process_mtp_loss,
     roll_tensor,
 )
 from megatron.core.transformer.spec_utils import ModuleSpec
@@ -141,7 +142,7 @@ def __init__(
         self.rotary_scaling = rope_scaling
         self.mtp_block_spec = mtp_block_spec
         self.mtp_process = mtp_block_spec is not None and mtp_on_this_rank(
-            self.config, vp_stage=vp_stage
+            self.config, ignore_virtual=False, vp_stage=vp_stage
         )
 
         if self.pre_process or self.mtp_process:
@@ -550,8 +551,8 @@ def _postprocess(
         if not self.post_process:
             return hidden_states
 
-        if self.mtp_process:
-            hidden_states = self.mtp.process_loss(
+        if self.config.mtp_num_layers is not None:
+            hidden_states = process_mtp_loss(
                 hidden_states=hidden_states,
                 labels=labels,
                 loss_mask=loss_mask,
@@ -560,6 +561,9 @@ def _postprocess(
                 runtime_gather_output=runtime_gather_output,
                 is_training=self.training,
                 compute_language_model_loss=self.compute_language_model_loss,
+                config=self.config,
+                cp_group=self.pg_collection.cp,
+                packed_seq_params=packed_seq_params,
             )
         sequence_parallel_override = False
 
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 5e483a1494e..8411c7fbbd4 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -19,6 +19,7 @@
 from megatron.core.transformer.multi_token_prediction import (
     mtp_on_this_rank,
     MultiTokenPredictionBlock,
+    process_mtp_loss,
 )
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.utils import (
@@ -320,8 +321,8 @@ def forward(
         if not self.post_process:
             return hidden_states
 
-        if self.mtp_process:
-            hidden_states = self.mtp.process_loss(
+        if self.config.mtp_num_layers is not None:
+            hidden_states = process_mtp_loss(
                 hidden_states=hidden_states,
                 labels=labels,
                 loss_mask=loss_mask,
@@ -330,6 +331,9 @@ def forward(
                 runtime_gather_output=runtime_gather_output,
                 is_training=self.training,
                 compute_language_model_loss=self.compute_language_model_loss,
+                config=self.config,
+                cp_group=self.pg_collection.cp,
+                packed_seq_params=packed_seq_params,
             )
 
         sequence_parallel_override = False
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 90d8c09ea6b..2403f44f500 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -625,6 +625,79 @@ def set_loss_scale(scale: torch.Tensor):
         MTPLossAutoScaler.main_loss_backward_scale = scale
 
 
+def process_mtp_loss(
+    hidden_states: Tensor,
+    labels: Tensor,
+    loss_mask: Optional[Tensor],
+    output_layer: Callable,
+    output_weight: Optional[Tensor],
+    runtime_gather_output: Optional[bool],
+    is_training: bool,
+    compute_language_model_loss: Callable,
+    config: TransformerConfig,
+    cp_group: Optional[torch.distributed.ProcessGroup] = None,
+    packed_seq_params: Optional[PackedSeqParams] = None,
+) -> Tensor:
+    """Process Multi-Token Prediction (MTP) loss computation.
+
+    This is a standalone function that handles MTP loss computation. It's used on the
+    post_process rank to split concatenated hidden states and compute MTP losses.
+
+    Args:
+        hidden_states (Tensor): Hidden states tensor (concatenated with MTP outputs).
+        labels (Tensor): Ground truth labels.
+        loss_mask (Optional[Tensor]): Mask for loss computation. If None, uses all ones.
+        output_layer (Callable): Output layer method to compute logits.
+        output_weight (Optional[Tensor]): Optional output weight for shared embeddings.
+        runtime_gather_output (Optional[bool]): Whether to gather output at runtime.
+        is_training (bool): Whether the model is in training mode.
+        compute_language_model_loss (Callable): Method to compute language model loss.
+        config (TransformerConfig): Model configuration containing mtp_num_layers etc.
+        cp_group (Optional[ProcessGroup]): Context parallelism process group.
+        packed_seq_params (Optional[PackedSeqParams]): Packed sequence parameters.
+
+    Returns:
+        Tensor: Updated hidden states after MTP loss processing (first chunk only).
+    """
+    mtp_labels = labels.clone()
+    hidden_states_list = torch.chunk(hidden_states, 1 + config.mtp_num_layers, dim=0)
+    hidden_states = hidden_states_list[0]
+
+    if loss_mask is None:
+        loss_mask = torch.ones_like(mtp_labels)
+
+    for mtp_layer_number in range(config.mtp_num_layers):
+        mtp_logits, _ = output_layer(
+            hidden_states_list[mtp_layer_number + 1],
+            weight=output_weight,
+            runtime_gather_output=runtime_gather_output,
+        )
+        mtp_labels, _ = roll_tensor(
+            mtp_labels, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params
+        )
+        loss_mask, num_tokens = roll_tensor(
+            loss_mask, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params
+        )
+        mtp_loss = compute_language_model_loss(mtp_labels, mtp_logits)
+        mtp_loss = loss_mask * mtp_loss
+        if is_training:
+            MTPLossLoggingHelper.save_loss_to_tracker(
+                torch.sum(mtp_loss) / num_tokens,
+                mtp_layer_number,
+                config.mtp_num_layers,
+                avg_group=parallel_state.get_data_parallel_group(with_context_parallel=True),
+            )
+        mtp_loss_scale = config.mtp_loss_scaling_factor / config.mtp_num_layers
+        if config.calculate_per_token_loss:
+            hidden_states = MTPLossAutoScaler.apply(hidden_states, mtp_loss_scale * mtp_loss)
+        else:
+            hidden_states = MTPLossAutoScaler.apply(
+                hidden_states, mtp_loss_scale * mtp_loss / num_tokens
+            )
+
+    return hidden_states
+
+
 class MultiTokenPredictionLayer(MegatronModule):
     """The implementation for Multi-Token Prediction (MTP) which extends
     the prediction scope to multiple future tokens at each position.
@@ -1276,10 +1349,11 @@ def build_layer_with_pattern(layer_spec, layer_number, mtp_layer_pattern, mamba_
                 ])
         elif self.mtp_use_repeated_layer:
             # Legacy repeated layer mode
-            assert len(self.submodules.layer_specs) == 1, (
-                f"Repeated MTP mode requires exactly 1 layer spec, got {len(self.submodules.layer_specs)}. "
-                f"The layer will be applied {self.config.mtp_num_layers} times."
-            )
+            if len(self.submodules.layer_specs) != 1:
+                warnings.warn(
+                    f"Repeated MTP mode expects exactly 1 layer spec, got {len(self.submodules.layer_specs)}. "
+                    f"The first layer will be applied {self.config.mtp_num_layers} times."
+                )
             self.layers = torch.nn.ModuleList([
                 build_layer_legacy(self.submodules.layer_specs[0], layer_number=1)
             ])
@@ -1351,82 +1425,6 @@ def forward(
         hidden_states = torch.cat(hidden_states_list, dim=0)
         return hidden_states
 
-    def process_loss(
-        self,
-        hidden_states: Tensor,
-        labels: Tensor,
-        loss_mask: Optional[Tensor],
-        output_layer: Callable,
-        output_weight: Optional[Tensor],
-        runtime_gather_output: Optional[bool],
-        is_training: bool,
-        compute_language_model_loss: Callable,
-        packed_seq_params: Optional[PackedSeqParams] = None,
-    ) -> Tensor:
-        """Process Multi-Token Prediction (MTP) loss computation.
-
-        This method handles the MTP loss computation for multiple prediction layers.
-        It chunks the hidden states, computes logits and losses for each MTP layer,
-        and applies loss scaling.
-
-        Args:
-            hidden_states (Tensor): Hidden states tensor from the model.
-            labels (Tensor): Ground truth labels.
-            loss_mask (Optional[Tensor]): Mask for loss computation. If None, uses all ones.
-            output_layer (Callable): Output layer method to compute logits.
-            output_weight (Optional[Tensor]): Optional output weight for shared embeddings.
-            runtime_gather_output (Optional[bool]): Whether to gather output at runtime.
-            is_training (bool): Whether the model is in training mode.
-            compute_language_model_loss (Callable): Method to compute language model loss.
-
-        Returns:
-            Tensor: Updated hidden states after MTP loss processing.
-        """
-        mtp_labels = labels.clone()
-        hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0)
-        hidden_states = hidden_states_list[0]
-
-        if loss_mask is None:
-            # if loss_mask is not provided, use all ones as loss_mask
-            loss_mask = torch.ones_like(mtp_labels)
-
-        for mtp_layer_number in range(self.config.mtp_num_layers):
-            # output
-            mtp_logits, _ = output_layer(
-                hidden_states_list[mtp_layer_number + 1],
-                weight=output_weight,
-                runtime_gather_output=runtime_gather_output,
-            )
-            # Calc loss for the current Multi-Token Prediction (MTP) layers.
-            mtp_labels, _ = roll_tensor(mtp_labels, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params)
-            loss_mask, num_tokens = roll_tensor(
-                loss_mask, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params
-            )
-            mtp_loss = compute_language_model_loss(mtp_labels, mtp_logits)
-            mtp_loss = loss_mask * mtp_loss
-            if is_training:
-                # TODO(shifangx): remove the use of parallel_state here
-                # after moving loss logging to loss_func in pretrain_gpt.py
-                MTPLossLoggingHelper.save_loss_to_tracker(
-                    torch.sum(mtp_loss) / num_tokens,
-                    mtp_layer_number,
-                    self.config.mtp_num_layers,
-                    avg_group=parallel_state.get_data_parallel_group(
-                        with_context_parallel=True
-                    ),
-                )
-            mtp_loss_scale = self.config.mtp_loss_scaling_factor / self.config.mtp_num_layers
-            if self.config.calculate_per_token_loss:
-                hidden_states = MTPLossAutoScaler.apply(
-                    hidden_states, mtp_loss_scale * mtp_loss
-                )
-            else:
-                hidden_states = MTPLossAutoScaler.apply(
-                    hidden_states, mtp_loss_scale * mtp_loss / num_tokens
-                )
-
-        return hidden_states
-
     def sharded_state_dict(
         self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None
     ) -> ShardedStateDict:
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 4e14ca9d253..b589839ff82 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -578,10 +578,13 @@ def validate_args(args, defaults={}):
             if args.mtp_num_layers is None:
                 args.mtp_num_layers = inferred_mtp_num_layers
             elif args.mtp_num_layers != inferred_mtp_num_layers:
-                raise ValueError(
+                warn_rank_0(
                     f"--mtp-num-layers ({args.mtp_num_layers}) conflicts with "
-                    f"MTP depth count ({inferred_mtp_num_layers}) in pattern '{args.hybrid_override_pattern}'"
+                    f"MTP depth count ({inferred_mtp_num_layers}) in pattern '{args.hybrid_override_pattern}'. "
+                    f"Using the inferred value ({inferred_mtp_num_layers}).",
+                    args.rank
                 )
+                args.mtp_num_layers = inferred_mtp_num_layers
 
     # MTP validation
     if args.mtp_num_layers:

From ae7b12a7f9d5e602b6e92ee1d5725c86a4592e7d Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Thu, 22 Jan 2026 15:21:17 -0800
Subject: [PATCH 40/48] Switch to using mamba-stack for MTP inner layers

---
 megatron/core/ssm/mamba_block.py              |  17 ++-
 .../transformer/multi_token_prediction.py     | 103 ++----------------
 2 files changed, 26 insertions(+), 94 deletions(-)

diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index e635789e321..90b4f71c453 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -91,6 +91,7 @@ def __init__(
         dtype=None,
         pg_collection: ProcessGroupCollection = None,
         vp_stage: Optional[int] = None,
+        is_mtp_layer: bool = False,
     ) -> None:
         super().__init__(config=config)
         self.residual_in_fp32 = residual_in_fp32
@@ -98,6 +99,7 @@ def __init__(
         self.post_layer_norm = post_layer_norm
         self.post_process = post_process
         self.vp_stage = vp_stage
+        self.is_mtp_layer = is_mtp_layer
 
         assert pg_collection is not None, "pg_collection must be provided for MambaStack"
 
@@ -112,8 +114,17 @@ def __init__(
         self.hybrid_override_pattern = hybrid_override_pattern
         self.pg_collection = pg_collection
 
+        # For MTP layers, always use pattern length (not config.num_layers which is for main decoder)
+        if self.is_mtp_layer:
+            num_layers_for_allocation = len(self.hybrid_override_pattern)
+        else:
+            num_layers_for_allocation = (
+                self.config.num_layers if self.config.num_layers is not None
+                else len(self.hybrid_override_pattern)
+            )
+
         self.layer_type_list = allocate_layers(
-            self.config.num_layers if self.config.num_layers is not None else len(self.hybrid_override_pattern),
+            num_layers_for_allocation,
             self.hybrid_attention_ratio,
             self.hybrid_mlp_ratio,
             self.hybrid_override_pattern,
@@ -134,7 +145,7 @@ def __init__(
             pg_collection=pg_collection,
             layer_offset=pp_layer_offset,
             residual_in_fp32=residual_in_fp32,
-            is_mtp_layer=False,
+            is_mtp_layer=self.is_mtp_layer,
         )
 
         # Required for activation recomputation
@@ -149,7 +160,7 @@ def __init__(
             )
 
     def _select_layers_for_pipeline_parallel(self, layer_type_list):
-        num_layers_per_pipeline_rank = self.config.num_layers // self.pp_group.size()
+        num_layers_per_pipeline_rank = len(layer_type_list) // self.pp_group.size()
 
         assert self.config.virtual_pipeline_model_parallel_size is None, (
             "The Mamba hybrid model does not currently support "
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 2403f44f500..5d3a82ee8a9 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -524,39 +524,6 @@ def get_mtp_layer_offset(config: TransformerConfig, vp_stage: Optional[int] = No
     return offset
 
 
-class MTPModelLayerContainer(MegatronModule):
-    """Container for MTP inner layers built from a pattern.
-
-    This container wraps the inner layers (e.g., Mamba, Attention layers) built
-    from a pattern string. It provides proper state dict key structure:
-    `mtp_model_layer.layers.{i}.*` for checkpoint compatibility.
-
-    """
-
-    def __init__(self, layers: torch.nn.ModuleList, config: TransformerConfig):
-        super().__init__(config=config)
-        self.layers = layers
-        self._current_microbatch = 0
-
-    @property
-    def current_microbatch(self):
-        """Get current microbatch index for CUDA graph management."""
-        return self._current_microbatch
-
-    @current_microbatch.setter
-    def current_microbatch(self, value):
-        """Set current microbatch and propagate to inner layers for CUDA graphs."""
-        self._current_microbatch = value
-        for layer in self.layers:
-            if hasattr(layer, 'current_microbatch'):
-                layer.current_microbatch = value
-
-    def set_input_tensor(self, input_tensor: Tensor):
-        """Set input tensor for pipeline parallelism compatibility."""
-        # Inner layers handle their own input tensors
-        pass
-
-
 def get_mtp_num_layers_to_build(
     config: TransformerConfig, vp_stage: Optional[int] = None, pp_rank: Optional[int] = None
 ) -> int:
@@ -795,21 +762,22 @@ def __init__(
         )
 
         # Build inner layers: two possible paths
-        # 1. New Mamba path: build from pattern using shared layer_builder
+        # 1. Mamba path: use MambaStack for hybrid pattern support (and future PP)
         # 2. GPT path: single TransformerLayer
         if mtp_layer_pattern is not None and mamba_submodules is not None:
-            # New Mamba path: build inner layers from pattern using shared layer_builder
-            from megatron.core.ssm.layer_builder import build_layers_from_pattern
+            from megatron.core.ssm.mamba_block import MambaStack
 
-            inner_layers = build_layers_from_pattern(
-                pattern=mtp_layer_pattern,
+            self.mtp_model_layer = MambaStack(
+                config=self.config,
                 submodules=mamba_submodules,
-                config=config,
+                hybrid_override_pattern=mtp_layer_pattern,
+                pre_process=True,       # Always receives input from eh_proj
+                post_layer_norm=False,  # MTP has its own final_layernorm
+                post_process=True,      # MTP layer is self-contained
                 pg_collection=pg_collection,
-                layer_offset=0,
-                is_mtp_layer=True,
+                vp_stage=vp_stage,
+                is_mtp_layer=True,      # Mark as MTP for router aux loss scaling
             )
-            self.mtp_model_layer = MTPModelLayerContainer(inner_layers, config=config)
         elif self.config.mtp_num_layers is not None:
             # GPT path: Uses the transformer block spec for MTP layer
             # MTP inner layers use their own layer numbering (self.layer_number = 1, 2, etc.)
@@ -941,12 +909,11 @@ def _proj_and_transformer_layer(
             # True so that the fp8 weight caching can be triggered correctly.
             with transformer_layer_fp8_context:
                 if self.mtp_layer_pattern is not None:
-                    # New Mamba path: run through pattern-based inner layers
-                    hidden_states = self._run_pattern_layers(
+                    hidden_states = self.mtp_model_layer(
                         hidden_states=hidden_states,
                         attention_mask=attention_mask,
                         rotary_pos_emb=rotary_pos_emb,
-                        inference_params=inference_params,
+                        inference_context=inference_params,
                         packed_seq_params=packed_seq_params,
                     )
                 else:
@@ -983,52 +950,6 @@ def _postprocess(self, hidden_states: torch.Tensor):
 
         return hidden_states
 
-    def _run_pattern_layers(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        rotary_pos_emb: Optional[torch.Tensor] = None,
-        inference_params: Optional[InferenceParams] = None,
-        packed_seq_params: Optional[PackedSeqParams] = None,
-    ) -> torch.Tensor:
-        """Run through pattern-based inner layers (Mamba path).
-
-        This method iterates through the layers in MTPModelLayerContainer,
-        calling each layer with the appropriate interface based on its type.
-
-        Args:
-            hidden_states: Input hidden states [s, b, h]
-            attention_mask: Attention mask tensor
-            rotary_pos_emb: Rotary position embeddings
-            inference_params: Inference parameters
-            packed_seq_params: Packed sequence parameters
-
-        Returns:
-            Output hidden states after processing through all inner layers
-        """
-        from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols
-        from megatron.core.transformer.transformer_layer import TransformerLayer
-
-        for layer_type, layer in zip(self.mtp_layer_pattern, self.mtp_model_layer.layers):
-            if layer_type == Symbols.ATTENTION or isinstance(layer, TransformerLayer):
-                # TransformerLayer (attention or MoE) returns tuple (hidden_states, context)
-                hidden_states, _ = layer(
-                    hidden_states=hidden_states,
-                    attention_mask=attention_mask,
-                    rotary_pos_emb=rotary_pos_emb,
-                    inference_params=inference_params,
-                    packed_seq_params=packed_seq_params,
-                )
-            else:
-                # MambaLayer, MLPLayer - return hidden_states directly
-                hidden_states = layer(
-                    hidden_states=hidden_states,
-                    attention_mask=attention_mask,
-                    inference_context=inference_params,
-                )
-
-        return hidden_states
-
     def _checkpointed_forward(self, forward_func, *args, **kwargs):
         def checkpoint_handler():
             """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`"""

From 6683247e72eac508e35d36587eec38b3937716b0 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Thu, 22 Jan 2026 22:50:11 -0800
Subject: [PATCH 41/48] Address review comments

---
 .../core/models/mamba/mamba_layer_specs.py    |   4 +-
 megatron/core/ssm/layer_builder.py            | 139 ------------------
 megatron/core/ssm/mamba_block.py              |   7 +-
 .../core/ssm/mamba_hybrid_layer_allocation.py |  24 +--
 .../transformer/multi_token_prediction.py     |   4 +-
 5 files changed, 21 insertions(+), 157 deletions(-)
 delete mode 100644 megatron/core/ssm/layer_builder.py

diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index 41f3dd3412d..9ea00872056 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -20,6 +20,7 @@
 from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.multi_token_prediction import (
     MultiTokenPredictionBlock,
@@ -38,8 +39,7 @@
 
 
 # MTP block spec for Mamba - provides norms and projection only.
-# Inner layers are built by MultiTokenPredictionLayer using the shared
-# layer_builder with mtp_layer_pattern and mamba_submodules.
+# Inner layers are built by MultiTokenPredictionLayer using nested MambaStack
 _mamba_mtp_block_spec = ModuleSpec(
     module=MultiTokenPredictionBlock,
     submodules=MultiTokenPredictionBlockSubmodules(
diff --git a/megatron/core/ssm/layer_builder.py b/megatron/core/ssm/layer_builder.py
deleted file mode 100644
index 52342d0ec9d..00000000000
--- a/megatron/core/ssm/layer_builder.py
+++ /dev/null
@@ -1,139 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-
-"""
-Shared layer-building utilities for Mamba hybrid architectures.
-
-Used by:
-- MambaStack: for main decoder layers
-- MultiTokenPredictionLayer: for MTP inner layers
-
-This module provides a single source of truth for converting layer pattern strings
-(e.g., "M*M*" or "MM") into actual layer modules.
-"""
-
-from typing import TYPE_CHECKING, List
-
-from torch import nn
-
-from megatron.core.fp8_utils import get_fp8_context
-from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols
-from megatron.core.transformer.spec_utils import build_module
-
-if TYPE_CHECKING:
-    from megatron.core.process_groups_config import ProcessGroupCollection
-    from megatron.core.ssm.mamba_block import MambaStackSubmodules
-    from megatron.core.transformer import TransformerConfig
-
-
-def build_layers_from_pattern(
-    pattern: str,
-    submodules: "MambaStackSubmodules",
-    config: "TransformerConfig",
-    pg_collection: "ProcessGroupCollection",
-    layer_offset: int = 0,
-    residual_in_fp32: bool = False,
-    is_mtp_layer: bool = False,
-) -> nn.ModuleList:
-    """
-    Build layers from a pattern string. Single source of truth for pattern → layers.
-
-    This function is shared between:
-    - MambaStack (for main decoder layers)
-    - MultiTokenPredictionLayer (for MTP inner layers)
-
-    Args:
-        pattern: Layer pattern string, e.g., "M*M*" or "MM"
-            - M = Mamba layer
-            - * = Attention (TransformerLayer)
-            - - = MLP layer
-            - E = MoE layer
-        submodules: MambaStackSubmodules containing specs for each layer type
-        config: TransformerConfig for the model
-        pg_collection: Process group collection for distributed training
-        layer_offset: Starting layer number offset (for layer_number assignment)
-        residual_in_fp32: Whether to use FP32 for residual connections (Mamba layers)
-        is_mtp_layer: Whether these layers are for MTP (affects attention layer building)
-
-    Returns:
-        nn.ModuleList of layers matching the pattern
-
-    Example:
-        >>> submodules = mamba_stack_spec.submodules
-        >>> layers = build_layers_from_pattern("M*M", submodules, config, pg_collection)
-        >>> # Returns [MambaLayer, TransformerLayer, MambaLayer]
-    """
-    layers = nn.ModuleList()
-
-    for i, layer_type in enumerate(pattern):
-        layer_number = i + 1 + layer_offset
-        fp8_init_context = get_fp8_context(config, i + layer_offset, is_init=True)
-
-        with fp8_init_context:
-            if layer_type == Symbols.MAMBA:
-                layer = build_module(
-                    submodules.mamba_layer,
-                    config=config,
-                    residual_in_fp32=residual_in_fp32,
-                    layer_number=layer_number,
-                    pp_layer_offset=layer_offset,
-                    pg_collection=pg_collection,
-                )
-            elif layer_type == Symbols.ATTENTION:
-                # Transformer layers apply their own pp_layer_offset
-                layer = build_module(
-                    submodules.attention_layer,
-                    config=config,
-                    layer_number=i + 1,
-                    pg_collection=pg_collection,
-                    is_mtp_layer=is_mtp_layer,
-                )
-            elif layer_type == Symbols.MLP:
-                # MLP layers apply their own pp_layer_offset
-                layer = build_module(
-                    submodules.mlp_layer,
-                    config=config,
-                    layer_number=i + 1,
-                    pg_collection=pg_collection,
-                )
-            elif layer_type == Symbols.MOE:
-                # MoE layers apply their own pp_layer_offset
-                layer = build_module(
-                    submodules.moe_layer,
-                    config=config,
-                    layer_number=i + 1,
-                    pg_collection=pg_collection,
-                    is_mtp_layer=is_mtp_layer,
-                )
-            else:
-                raise ValueError(
-                    f"Unknown layer type: '{layer_type}'. "
-                    f"Valid types are: {Symbols.VALID}"
-                )
-
-        layers.append(layer)
-
-    return layers
-
-
-def get_layer_types_from_pattern(pattern: str) -> List[str]:
-    """
-    Convert a pattern string to a list of layer type symbols.
-
-    This is a simple utility that validates and returns the pattern as a list.
-
-    Args:
-        pattern: Layer pattern string, e.g., "M*M*"
-
-    Returns:
-        List of layer type symbols, e.g., ["M", "*", "M", "*"]
-
-    Raises:
-        ValueError: If pattern contains invalid symbols
-    """
-    layer_types = list(pattern)
-    for char in layer_types:
-        if char not in Symbols.VALID:
-            raise ValueError(
-                f"Invalid layer type: '{char}'. Valid types are: {Symbols.VALID}"
-            )
-    return layer_types
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index 70866b877b6..59d88fa3243 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -21,7 +21,7 @@
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols as LayerSymbols
-from megatron.core.ssm.mamba_hybrid_layer_allocation import allocate_layers, parse_hybrid_pattern
+from megatron.core.ssm.mamba_hybrid_layer_allocation import allocate_layers
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.identity_op import IdentityOp
@@ -127,10 +127,11 @@ def __init__(
             self.hybrid_attention_ratio,
             self.hybrid_mlp_ratio,
             self.hybrid_override_pattern,
+            silent=self.is_mtp_layer,
         )
 
         pp_layer_offset = 0
-        if self.pp_group.size() > 1:
+        if self.pp_group.size() > 1 and not self.is_mtp_layer:
             pp_layer_offset, self.layer_type_list = self._select_layers_for_pipeline_parallel(
                 self.layer_type_list
             )
@@ -189,7 +190,7 @@ def __init__(
             )
 
     def _select_layers_for_pipeline_parallel(self, layer_type_list):
-        num_layers_per_pipeline_rank = len(layer_type_list) // self.pp_group.size()
+        num_layers_per_pipeline_rank = self.config.num_layers // self.pp_group.size()
 
         assert self.config.virtual_pipeline_model_parallel_size is None, (
             "The Mamba hybrid model does not currently support "
diff --git a/megatron/core/ssm/mamba_hybrid_layer_allocation.py b/megatron/core/ssm/mamba_hybrid_layer_allocation.py
index 146bc41390d..07d0fd381ed 100644
--- a/megatron/core/ssm/mamba_hybrid_layer_allocation.py
+++ b/megatron/core/ssm/mamba_hybrid_layer_allocation.py
@@ -218,19 +218,21 @@ def allocate_layers(
     target_attention_ratio: float,
     target_mlp_ratio: float,
     override_pattern: str = None,
+    silent: bool = False,
 ) -> list:
     """Allocates layers according to the requested distribution of layer types."""
     assert total_layers_count > 0
     assert target_attention_ratio >= 0.0 and target_attention_ratio <= 1.0
     assert target_mlp_ratio >= 0.0 and target_mlp_ratio <= 1.0
     assert target_attention_ratio + target_mlp_ratio <= 1.0
+    maybe_log_single_rank = (lambda *args, **kwargs: None) if silent else log_single_rank
     # Note: target_mamba_ratio = 1.0 - target_attention_ratio - target_mlp_ratio
 
     layer_type_list = _allocate_auto(total_layers_count, target_attention_ratio, target_mlp_ratio)
 
     if override_pattern is not None:
         layer_type_list_override = _allocate_override(total_layers_count, override_pattern)
-        log_single_rank(logger, logging.INFO, "Using hybrid override pattern")
+        maybe_log_single_rank(logger, logging.INFO, "Using hybrid override pattern")
         if (target_attention_ratio > 0.0 or target_mlp_ratio > 0.0) and not _layer_counts_match(
             layer_type_list_override, layer_type_list
         ):
@@ -240,13 +242,13 @@ def allocate_layers(
                 "pattern."
             )
         if layer_type_list_override == layer_type_list:
-            log_single_rank(
+            maybe_log_single_rank(
                 logger, logging.INFO, "The override pattern matches the overridden pattern"
             )
         else:
-            log_single_rank(logger, logging.INFO, "Warning: overriding pattern A with pattern B")
-            log_single_rank(logger, logging.INFO, f"A: {''.join(layer_type_list)}")
-            log_single_rank(logger, logging.INFO, f"B: {''.join(layer_type_list_override)}")
+            maybe_log_single_rank(logger, logging.INFO, "Warning: overriding pattern A with pattern B")
+            maybe_log_single_rank(logger, logging.INFO, f"A: {''.join(layer_type_list)}")
+            maybe_log_single_rank(logger, logging.INFO, f"B: {''.join(layer_type_list_override)}")
         layer_type_list = layer_type_list_override
 
     if target_attention_ratio > 0.0 or target_mlp_ratio > 0.0 or override_pattern is not None:
@@ -255,32 +257,32 @@ def allocate_layers(
         actual_mlp_layers_count = layer_type_list.count(Symbols.MLP)
         actual_mlp_ratio = actual_mlp_layers_count / total_layers_count
         allocation_string = "".join(layer_type_list)
-        log_single_rank(
+        maybe_log_single_rank(
             logger,
             logging.INFO,
             f"Hybrid allocation ({Symbols.MAMBA} is mamba, "
             f"{Symbols.ATTENTION} is attention, "
             f"{Symbols.MLP} is mlp):",
         )
-        log_single_rank(logger, logging.INFO, allocation_string)
-        log_single_rank(
+        maybe_log_single_rank(logger, logging.INFO, allocation_string)
+        maybe_log_single_rank(
             logger,
             logging.INFO,
             f"{actual_attention_layers_count} attention layers in "
             f"{total_layers_count} total layers.",
         )
-        log_single_rank(
+        maybe_log_single_rank(
             logger,
             logging.INFO,
             f"Target attention ratio: {target_attention_ratio:.2f}. "
             f"Actual attention ratio: {actual_attention_ratio:.2f}.",
         )
-        log_single_rank(
+        maybe_log_single_rank(
             logger,
             logging.INFO,
             f"{actual_mlp_layers_count} mlp layers in " f"{total_layers_count} total layers.",
         )
-        log_single_rank(
+        maybe_log_single_rank(
             logger,
             logging.INFO,
             f"Target mlp ratio: {target_mlp_ratio:.2f}. "
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 5d3a82ee8a9..99c1534db3f 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -762,7 +762,7 @@ def __init__(
         )
 
         # Build inner layers: two possible paths
-        # 1. Mamba path: use MambaStack for hybrid pattern support (and future PP)
+        # 1. Mamba path: use MambaStack for hybrid pattern support
         # 2. GPT path: single TransformerLayer
         if mtp_layer_pattern is not None and mamba_submodules is not None:
             from megatron.core.ssm.mamba_block import MambaStack
@@ -776,7 +776,7 @@ def __init__(
                 post_process=True,      # MTP layer is self-contained
                 pg_collection=pg_collection,
                 vp_stage=vp_stage,
-                is_mtp_layer=True,      # Mark as MTP for router aux loss scaling
+                is_mtp_layer=True,
             )
         elif self.config.mtp_num_layers is not None:
             # GPT path: Uses the transformer block spec for MTP layer

From 7bfa95f8600407bf324044920963d6a9df0362f5 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Mon, 26 Jan 2026 11:58:58 -0800
Subject: [PATCH 42/48] Address review comments

---
 megatron/core/models/gpt/gpt_model.py         |  3 ---
 megatron/core/models/mamba/mamba_model.py     |  4 ----
 megatron/core/ssm/mamba_block.py              |  5 -----
 .../transformer/multi_token_prediction.py     |  7 -------
 megatron/training/arguments.py                | 19 -------------------
 5 files changed, 38 deletions(-)

diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index 1183c8fba93..cfbe6aa16d4 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -26,12 +26,9 @@
 from megatron.core.tensor_parallel import gather_from_sequence_parallel_region
 from megatron.core.transformer.enums import CudaGraphScope, ModelType
 from megatron.core.transformer.multi_token_prediction import (
-    MTPLossAutoScaler,
-    MTPLossLoggingHelper,
     MultiTokenPredictionBlock,
     mtp_on_this_rank,
     process_mtp_loss,
-    roll_tensor,
 )
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_block import TransformerBlock
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 1f81893b1d3..fc51f19c3d6 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -121,10 +121,6 @@ def __init__(
             self.mtp_pattern is not None and self.mtp_num_depths > 0 and mtp_on_this_rank(self.config, vp_stage=self.vp_stage)
         )
 
-        # Cache for RoPE tensors which do not change between iterations.
-        self.rotary_pos_emb_cache = {}
-
-
         # megatron core pipelining currently depends on model type
         # TODO: remove this dependency ?
         self.model_type = ModelType.encoder_or_decoder
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index 59d88fa3243..93ee8ad95f8 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -29,9 +29,6 @@
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.transformer.transformer_layer import TransformerLayer
 from megatron.core.transformer.utils import sharded_state_dict_default
-
-from megatron.core.ssm.mlp_layer import MLPLayer
-
 from megatron.core.utils import WrappedTensor, deprecate_inference_params, make_viewless_tensor
 
 
@@ -89,7 +86,6 @@ def __init__(
         device=None,
         dtype=None,
         pg_collection: ProcessGroupCollection = None,
-        vp_stage: Optional[int] = None,
         is_mtp_layer: bool = False,
     ) -> None:
         super().__init__(config=config)
@@ -97,7 +93,6 @@ def __init__(
         self.pre_process = pre_process
         self.post_layer_norm = post_layer_norm
         self.post_process = post_process
-        self.vp_stage = vp_stage
         self.is_mtp_layer = is_mtp_layer
 
         assert pg_collection is not None, "pg_collection must be provided for MambaStack"
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 99c1534db3f..8556e50a5c1 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -23,12 +23,6 @@
 from megatron.core.transformer.enums import AttnMaskType, LayerType
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
-
-# Type hint imports for Mamba support
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from megatron.core.ssm.mamba_block import MambaStackSubmodules
 from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import (
@@ -775,7 +769,6 @@ def __init__(
                 post_layer_norm=False,  # MTP has its own final_layernorm
                 post_process=True,      # MTP layer is self-contained
                 pg_collection=pg_collection,
-                vp_stage=vp_stage,
                 is_mtp_layer=True,
             )
         elif self.config.mtp_num_layers is not None:
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 27926dd2c2e..bcc5b3f303a 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -592,13 +592,6 @@ def validate_args(args, defaults={}):
             "This argument is kept only for loading old checkpoints.",
             args.rank,
         )
-    if args.mtp_spec is not None:
-        warn_rank_0(
-            "--mtp-spec is deprecated. "
-            "For new hybrid models with MTP models, use unified --hybrid-override-pattern instead. "
-            "This argument is kept only for loading old checkpoints.",
-            args.rank,
-        )
 
     # Backward compatibility: convert legacy mtp_hybrid_override_pattern to unified format
     from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols, parse_hybrid_pattern
@@ -660,13 +653,6 @@ def validate_args(args, defaults={}):
                 "This argument will be ignored.",
                 args.rank
             )
-        if args.mtp_spec is not None:
-            warn_rank_0(
-                "--mtp-spec is for Mamba/hybrid models only. "
-                "For GPT models, use the standard layer spec. "
-                "This argument will be ignored.",
-                args.rank
-            )
     # === End of MTP validation ===
     
     # Uneven virtual pipeline parallelism
@@ -3678,11 +3664,6 @@ def _add_experimental_args(parser):
                        help='DEPRECATED: Use unified --hybrid-override-pattern instead. '
                        'Legacy argument for loading old checkpoints. '
                        'Force a specific hybrid layer pattern for MTP layers.')
-    group.add_argument('--mtp-spec', type=str, default=None, nargs='*',
-                       help='DEPRECATED: Use unified --hybrid-override-pattern instead. '
-                       'Legacy argument for loading old checkpoints. '
-                       'Specify the <module_location function_name> pair '
-                       'that returns a spec for mtp layer.')
     group.add_argument('--mamba-state-dim', type=int, default=128,
                        help='State dimension for Mamba layers.')
     group.add_argument('--mamba-head-dim', type=int, default=64,

From ea3c2bd8e2db7c88a090fff76900d3454fd980d3 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Sat, 31 Jan 2026 00:17:48 -0800
Subject: [PATCH 43/48] Fix linting errors

---
 .../common/language_module/language_module.py |  8 ++--
 .../core/models/mamba/mamba_layer_specs.py    | 14 +++---
 megatron/core/models/mamba/mamba_model.py     |  6 ++-
 megatron/core/pipeline_parallel/schedules.py  |  6 +--
 megatron/core/ssm/mamba_block.py              |  3 +-
 .../core/ssm/mamba_hybrid_layer_allocation.py |  4 +-
 megatron/core/transformer/attention.py        |  8 +++-
 megatron/core/transformer/moe/moe_layer.py    | 10 ++--
 megatron/core/transformer/moe/router.py       | 24 +++++++---
 .../transformer/multi_token_prediction.py     | 46 ++++++++++---------
 .../ssm/test_mamba_hybrid_layer_allocation.py |  5 +-
 .../test_multi_token_prediction.py            | 18 ++------
 12 files changed, 84 insertions(+), 68 deletions(-)

diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
index bbc7de8e724..f94e51df69d 100644
--- a/megatron/core/models/common/language_module/language_module.py
+++ b/megatron/core/models/common/language_module/language_module.py
@@ -23,6 +23,10 @@
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.enums import AttnBackend, CudaGraphScope
 from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.multi_token_prediction import (
+    tie_output_layer_state_dict,
+    tie_word_embeddings_state_dict,
+)
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.transformer.utils import ensure_metadata_has_dp_cp_group
 from megatron.core.utils import (
@@ -31,10 +35,6 @@
     make_tp_sharded_tensor_for_checkpoint,
 )
 
-from megatron.core.transformer.multi_token_prediction import (
-    tie_output_layer_state_dict,
-    tie_word_embeddings_state_dict,
-)
 
 class LanguageModule(MegatronModule):
     """Base language module that has common helper functions used across GPT, BERT etc.
diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index e7781816dfc..6ca628475be 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -1,11 +1,11 @@
 # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
 
 from megatron.core.extensions.transformer_engine import (
+    TEColumnParallelLinear,
     TEDotProductAttention,
     TELayerNormColumnParallelLinear,
     TENorm,
     TERowParallelLinear,
-    TEColumnParallelLinear,
 )
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec
@@ -20,18 +20,18 @@
 from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.core.transformer.mlp import MLP, MLPSubmodules
-from megatron.core.transformer.transformer_layer import (
-    MoETransformerLayer,
-    TransformerLayer,
-    TransformerLayerSubmodules,
-)
-from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.multi_token_prediction import (
     MultiTokenPredictionBlock,
     MultiTokenPredictionBlockSubmodules,
     MultiTokenPredictionLayer,
     MultiTokenPredictionLayerSubmodules,
 )
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_layer import (
+    MoETransformerLayer,
+    TransformerLayer,
+    TransformerLayerSubmodules,
+)
 
 # This should be private and should not be used outside of this file.
 moe = get_moe_module_spec(
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index bd6b8904142..0a783391437 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -17,8 +17,8 @@
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.enums import ModelType
 from megatron.core.transformer.multi_token_prediction import (
-    mtp_on_this_rank,
     MultiTokenPredictionBlock,
+    mtp_on_this_rank,
     process_mtp_loss,
 )
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
@@ -118,7 +118,9 @@ def __init__(
 
         # Determine if MTP is needed (based on pattern parsing)
         self.mtp_process = (
-            self.mtp_pattern is not None and self.mtp_num_depths > 0 and mtp_on_this_rank(self.config, vp_stage=self.vp_stage)
+            self.mtp_pattern is not None
+            and self.mtp_num_depths > 0
+            and mtp_on_this_rank(self.config, vp_stage=self.vp_stage)
         )
 
         # megatron core pipelining currently depends on model type
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index 36fcb175c21..fd39b1136fe 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -212,9 +212,9 @@ def set_current_microbatch(model, microbatch_id):
             layer.current_microbatch = microbatch_id
         if hasattr(model_with_decoder, 'mtp'):
             for layer in model_with_decoder.mtp.layers:
-                assert hasattr(layer, 'mtp_model_layer'), (
-                    f"MTP layer {layer} must have 'mtp_model_layer' attribute"
-                )
+                assert hasattr(
+                    layer, 'mtp_model_layer'
+                ), f"MTP layer {layer} must have 'mtp_model_layer' attribute"
                 layer.mtp_model_layer.current_microbatch = microbatch_id
 
 
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index 66234dc16e9..f070a049b06 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -113,7 +113,8 @@ def __init__(
             num_layers_for_allocation = len(self.hybrid_override_pattern)
         else:
             num_layers_for_allocation = (
-                self.config.num_layers if self.config.num_layers is not None
+                self.config.num_layers
+                if self.config.num_layers is not None
                 else len(self.hybrid_override_pattern)
             )
 
diff --git a/megatron/core/ssm/mamba_hybrid_layer_allocation.py b/megatron/core/ssm/mamba_hybrid_layer_allocation.py
index 07d0fd381ed..d7002b2915d 100644
--- a/megatron/core/ssm/mamba_hybrid_layer_allocation.py
+++ b/megatron/core/ssm/mamba_hybrid_layer_allocation.py
@@ -246,7 +246,9 @@ def allocate_layers(
                 logger, logging.INFO, "The override pattern matches the overridden pattern"
             )
         else:
-            maybe_log_single_rank(logger, logging.INFO, "Warning: overriding pattern A with pattern B")
+            maybe_log_single_rank(
+                logger, logging.INFO, "Warning: overriding pattern A with pattern B"
+            )
             maybe_log_single_rank(logger, logging.INFO, f"A: {''.join(layer_type_list)}")
             maybe_log_single_rank(logger, logging.INFO, f"B: {''.join(layer_type_list_override)}")
         layer_type_list = layer_type_list_override
diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py
index 2200b558225..d4ed4935774 100644
--- a/megatron/core/transformer/attention.py
+++ b/megatron/core/transformer/attention.py
@@ -59,7 +59,9 @@
     rearrange = None
 
 try:
-    from flash_attn_3.flash_attn_interface import _flash_attn_forward
+    from flash_attn_3.flash_attn_interface import (
+        _flash_attn_forward,
+    )
     from flash_attn_3.flash_attn_interface import (
         flash_attn_with_kvcache as flash_attn3_with_kvcache,
     )
@@ -70,7 +72,9 @@
 
 if not HAVE_FA3:
     try:
-        from flashattn_hopper.flash_attn_interface import _flash_attn_forward
+        from flashattn_hopper.flash_attn_interface import (
+            _flash_attn_forward,
+        )
         from flashattn_hopper.flash_attn_interface import (
             flash_attn_with_kvcache as flash_attn3_with_kvcache,
         )
diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index 6d3c840e8aa..b137605c614 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -150,8 +150,10 @@ def __init__(
         if pg_collection is None:
             pg_collection = get_default_pg_collection()
         super(MoELayer, self).__init__(
-            config=config, layer_number=layer_number, pg_collection=pg_collection,
-            is_mtp_layer=is_mtp_layer
+            config=config,
+            layer_number=layer_number,
+            pg_collection=pg_collection,
+            is_mtp_layer=is_mtp_layer,
         )
         # If using mcore cudagraphs, recompute is handled by transformer_layer.MoETransformerLayer
         self.moe_layer_recompute = (
@@ -167,7 +169,9 @@ def __init__(
         self.tp_group = pg_collection.tp
 
         # Initialize router.
-        self.router = submodules.router(config=self.config, pg_collection=pg_collection, is_mtp_layer=is_mtp_layer)
+        self.router = submodules.router(
+            config=self.config, pg_collection=pg_collection, is_mtp_layer=is_mtp_layer
+        )
         self.tp_group = pg_collection.tp
 
         # Initialize latent projections.
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index 8e874236ae2..a8b84f2c351 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -29,8 +29,10 @@ class Router(ABC, MegatronModule):
     """Base Router class"""
 
     def __init__(
-        self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None,
-        is_mtp_layer: bool = False
+        self,
+        config: TransformerConfig,
+        pg_collection: Optional[ProcessGroupCollection] = None,
+        is_mtp_layer: bool = False,
     ) -> None:
         """
         Initialize the Router module.
@@ -148,8 +150,10 @@ class TopKRouter(Router):
     """
 
     def __init__(
-        self, config: TransformerConfig, pg_collection: Optional[ProcessGroupCollection] = None,
-        is_mtp_layer: bool = False
+        self,
+        config: TransformerConfig,
+        pg_collection: Optional[ProcessGroupCollection] = None,
+        is_mtp_layer: bool = False,
     ) -> None:
         """Initialize the zero token dropping router.
 
@@ -445,7 +449,11 @@ def attach_and_log_load_balancing_loss(
         """
         # When using repeated MTP layers, the loss is counted "mtp_num_layers" times. To avoid accumulating the
         # load balancing loss multiple times, we scale it by 1/mtp_num_layers so the total loss is correct.
-        if self.is_mtp_layer and self.config.mtp_use_repeated_layer and self.config.mtp_num_layers is not None:
+        if (
+            self.is_mtp_layer
+            and self.config.mtp_use_repeated_layer
+            and self.config.mtp_num_layers is not None
+        ):
             aux_loss = aux_loss / self.config.mtp_num_layers
 
         # TODO (zijiey): fix the per_layer_logging for MTP, currently it will incorrectly
@@ -517,7 +525,11 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None):
 
             # When using repeated MTP layers, the same MTP layer is called mtp_num_layers times. To avoid
             # accumulating the z_loss multiple times, we scale it by 1/mtp_num_layers so the total loss is correct.
-            if self.is_mtp_layer and self.config.mtp_use_repeated_layer and self.config.mtp_num_layers is not None:
+            if (
+                self.is_mtp_layer
+                and self.config.mtp_use_repeated_layer
+                and self.config.mtp_num_layers is not None
+            ):
                 z_loss = z_loss / self.config.mtp_num_layers
 
             num_layers = self.config.num_layers
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index 8556e50a5c1..d6147d867ea 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -699,9 +699,8 @@ def __init__(
         self.mtp_layer_pattern = mtp_layer_pattern
 
         # Validate attention mask type if using transformer-based inner layers
-        if (
-            self.submodules.mtp_model_layer is not None
-            and hasattr(self.submodules.mtp_model_layer, 'submodules')
+        if self.submodules.mtp_model_layer is not None and hasattr(
+            self.submodules.mtp_model_layer, 'submodules'
         ):
             if hasattr(self.submodules.mtp_model_layer.submodules, 'attention_layer'):
                 self_attention_spec = self.submodules.mtp_model_layer.submodules.attention_layer
@@ -752,7 +751,7 @@ def __init__(
             bias=False,
             skip_bias_add=False,
             is_expert=False,
-            tp_comm_buffer_name="mtp_eh_proj"
+            tp_comm_buffer_name="mtp_eh_proj",
         )
 
         # Build inner layers: two possible paths
@@ -765,9 +764,9 @@ def __init__(
                 config=self.config,
                 submodules=mamba_submodules,
                 hybrid_override_pattern=mtp_layer_pattern,
-                pre_process=True,       # Always receives input from eh_proj
+                pre_process=True,  # Always receives input from eh_proj
                 post_layer_norm=False,  # MTP has its own final_layernorm
-                post_process=True,      # MTP layer is self-contained
+                post_process=True,  # MTP layer is self-contained
                 pg_collection=pg_collection,
                 is_mtp_layer=True,
             )
@@ -781,7 +780,7 @@ def __init__(
                 config=self.config,
                 vp_stage=self.vp_stage,
                 layer_number=self.layer_number,
-                is_mtp_layer=True
+                is_mtp_layer=True,
             )
 
         self.final_layernorm = build_module(
@@ -1220,7 +1219,7 @@ def build_layer_legacy(layer_spec, layer_number):
                     layer_number=layer_number,
                     vp_stage=self.vp_stage,
                     pg_collection=pg_collection,
-                    mtp_layer_pattern=self.mtp_layer_pattern
+                    mtp_layer_pattern=self.mtp_layer_pattern,
                 )
             return module
 
@@ -1245,22 +1244,27 @@ def build_layer_with_pattern(layer_spec, layer_number, mtp_layer_pattern, mamba_
                 # Shared/repeated layer: build one layer, use it for all depths
                 layer_spec = self.submodules.layer_specs[0]
                 shared_layer = build_layer_with_pattern(
-                    layer_spec, layer_number=1,
+                    layer_spec,
+                    layer_number=1,
                     mtp_layer_pattern=self.mtp_layer_pattern,
                     mamba_submodules=self.mamba_submodules,
                 )
                 self.layers = torch.nn.ModuleList([shared_layer])
             else:
                 # Non-shared: each depth gets its own layers
-                self.layers = torch.nn.ModuleList([
-                    build_layer_with_pattern(
-                        self.submodules.layer_specs[min(i, len(self.submodules.layer_specs) - 1)],
-                        layer_number=i + 1,
-                        mtp_layer_pattern=self.mtp_layer_pattern,
-                        mamba_submodules=self.mamba_submodules,
-                    )
-                    for i in range(num_depths)
-                ])
+                self.layers = torch.nn.ModuleList(
+                    [
+                        build_layer_with_pattern(
+                            self.submodules.layer_specs[
+                                min(i, len(self.submodules.layer_specs) - 1)
+                            ],
+                            layer_number=i + 1,
+                            mtp_layer_pattern=self.mtp_layer_pattern,
+                            mamba_submodules=self.mamba_submodules,
+                        )
+                        for i in range(num_depths)
+                    ]
+                )
         elif self.mtp_use_repeated_layer:
             # Legacy repeated layer mode
             if len(self.submodules.layer_specs) != 1:
@@ -1268,9 +1272,9 @@ def build_layer_with_pattern(layer_spec, layer_number, mtp_layer_pattern, mamba_
                     f"Repeated MTP mode expects exactly 1 layer spec, got {len(self.submodules.layer_specs)}. "
                     f"The first layer will be applied {self.config.mtp_num_layers} times."
                 )
-            self.layers = torch.nn.ModuleList([
-                build_layer_legacy(self.submodules.layer_specs[0], layer_number=1)
-            ])
+            self.layers = torch.nn.ModuleList(
+                [build_layer_legacy(self.submodules.layer_specs[0], layer_number=1)]
+            )
         else:
             # Legacy mode: build from layer_specs
             self.layers = torch.nn.ModuleList(
diff --git a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
index 7edb46a7e80..2119c0f6d23 100644
--- a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
+++ b/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
@@ -160,10 +160,7 @@ def test_invalid_symbols_in_main_pattern(self):
 
     def test_invalid_symbols_in_mtp_pattern(self):
         """Test that invalid symbols in MTP pattern raise ValueError."""
-        invalid_patterns = [
-            "M*M*/MX",  # X is not valid
-            "M*M*/MM/Ma",  # a is not valid
-        ]
+        invalid_patterns = ["M*M*/MX", "M*M*/MM/Ma"]  # X is not valid  # a is not valid
         for pattern in invalid_patterns:
             with pytest.raises(ValueError, match="not a valid layer symbol"):
                 parse_hybrid_pattern(pattern)
diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py
index 9e9a79521b2..ec72d713eb1 100644
--- a/tests/unit_tests/transformer/test_multi_token_prediction.py
+++ b/tests/unit_tests/transformer/test_multi_token_prediction.py
@@ -702,12 +702,7 @@ def teardown_method(self, method):
         destroy_num_microbatches_calculator()
         MTPLossLoggingHelper.tracker = {}
 
-    def model_provider(
-        self,
-        pre_process=True,
-        post_process=True,
-        **config_kwargs,
-    ):
+    def model_provider(self, pre_process=True, post_process=True, **config_kwargs):
         """Model provider for Mamba hybrid models with MTP.
 
         Uses the unified pattern syntax where MTP is configured via hybrid_override_pattern:
@@ -868,9 +863,7 @@ def set_ckpt_path(ckpt_path):
             args.save = ckpt_path
             args.load = ckpt_path
 
-        with TempNamedDir(
-            tmp_path_dist_ckpt / 'test_mtp_mamba_model_reconfiguration'
-        ) as ckpt_dir:
+        with TempNamedDir(tmp_path_dist_ckpt / 'test_mtp_mamba_model_reconfiguration') as ckpt_dir:
             set_ckpt_path(ckpt_dir)
             save_checkpoint(
                 iteration,
@@ -925,8 +918,7 @@ def set_ckpt_path(ckpt_path):
 
     @pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available")
     def test_attention_mask_validation_mamba(self):
-        """Test that attention mask type validation works for Mamba hybrid models.
-        """
+        """Test that attention mask type validation works for Mamba hybrid models."""
         tp = 1
         cp = 1
         args = self.create_test_args(tp, cp, self.seq_length, self.micro_batch_size)
@@ -940,8 +932,6 @@ def test_attention_mask_validation_mamba(self):
             assert mamba_model[0].mtp is not None
         except AssertionError as e:
             if "Multi-Token Prediction (MTP) is not yet supported" in str(e):
-                pytest.fail(
-                    f"Attention mask validation failed for Mamba hybrid model: {e}"
-                )
+                pytest.fail(f"Attention mask validation failed for Mamba hybrid model: {e}")
             else:
                 raise

From 19a56e5c9b3dc2e512af5fa7a17f29577e3d4616 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Sat, 31 Jan 2026 00:21:54 -0800
Subject: [PATCH 44/48] isort then black

---
 megatron/core/transformer/attention.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py
index d4ed4935774..2200b558225 100644
--- a/megatron/core/transformer/attention.py
+++ b/megatron/core/transformer/attention.py
@@ -59,9 +59,7 @@
     rearrange = None
 
 try:
-    from flash_attn_3.flash_attn_interface import (
-        _flash_attn_forward,
-    )
+    from flash_attn_3.flash_attn_interface import _flash_attn_forward
     from flash_attn_3.flash_attn_interface import (
         flash_attn_with_kvcache as flash_attn3_with_kvcache,
     )
@@ -72,9 +70,7 @@
 
 if not HAVE_FA3:
     try:
-        from flashattn_hopper.flash_attn_interface import (
-            _flash_attn_forward,
-        )
+        from flashattn_hopper.flash_attn_interface import _flash_attn_forward
         from flashattn_hopper.flash_attn_interface import (
             flash_attn_with_kvcache as flash_attn3_with_kvcache,
         )

From e50409bb290abf742a69bdd76d739448cb6b2eb9 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Sat, 31 Jan 2026 00:33:24 -0800
Subject: [PATCH 45/48] Fix pylint errors

---
 megatron/core/models/gpt/gpt_model.py               |  2 +-
 megatron/core/ssm/mamba_block.py                    |  2 +-
 megatron/core/transformer/moe/router.py             | 10 ++++++----
 megatron/core/transformer/multi_token_prediction.py |  3 ++-
 4 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index 7f90517c5b1..4b96465a31e 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -6,7 +6,7 @@
 import torch
 from torch import Tensor
 
-from megatron.core import parallel_state, tensor_parallel
+from megatron.core import tensor_parallel
 from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
 from megatron.core.dist_checkpointing.mapping import ShardedStateDict
 from megatron.core.inference.contexts import BaseInferenceContext
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index f070a049b06..ffb7b8f6fdb 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -108,7 +108,7 @@ def __init__(
         self.hybrid_override_pattern = hybrid_override_pattern
         self.pg_collection = pg_collection
 
-        # For MTP layers, always use pattern length (not config.num_layers which is for main decoder)
+        # For MTP layers, always use pattern length (config.num_layers is for main decoder)
         if self.is_mtp_layer:
             num_layers_for_allocation = len(self.hybrid_override_pattern)
         else:
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
index a8b84f2c351..e42fd1ca8aa 100644
--- a/megatron/core/transformer/moe/router.py
+++ b/megatron/core/transformer/moe/router.py
@@ -447,8 +447,9 @@ def attach_and_log_load_balancing_loss(
                 padding tokens. Can be a Python int or a torch.Tensor (typically 0-d tensor).
                 If None, uses activation.shape[0]. Defaults to None.
         """
-        # When using repeated MTP layers, the loss is counted "mtp_num_layers" times. To avoid accumulating the
-        # load balancing loss multiple times, we scale it by 1/mtp_num_layers so the total loss is correct.
+        # When using repeated MTP layers, the loss is counted "mtp_num_layers" times.
+        # To avoid accumulating the load balancing loss multiple times, we scale it by
+        # 1/mtp_num_layers so the total loss is correct.
         if (
             self.is_mtp_layer
             and self.config.mtp_use_repeated_layer
@@ -523,8 +524,9 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None):
             else:
                 logits = MoEAuxLossAutoScaler.apply(logits, z_loss)
 
-            # When using repeated MTP layers, the same MTP layer is called mtp_num_layers times. To avoid
-            # accumulating the z_loss multiple times, we scale it by 1/mtp_num_layers so the total loss is correct.
+            # When using repeated MTP layers, the same MTP layer is called mtp_num_layers times.
+            # To avoid accumulating the z_loss multiple times, we scale it by 1/mtp_num_layers
+            # so the total loss is correct.
             if (
                 self.is_mtp_layer
                 and self.config.mtp_use_repeated_layer
diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
index d6147d867ea..6432af36cde 100755
--- a/megatron/core/transformer/multi_token_prediction.py
+++ b/megatron/core/transformer/multi_token_prediction.py
@@ -1269,7 +1269,8 @@ def build_layer_with_pattern(layer_spec, layer_number, mtp_layer_pattern, mamba_
             # Legacy repeated layer mode
             if len(self.submodules.layer_specs) != 1:
                 warnings.warn(
-                    f"Repeated MTP mode expects exactly 1 layer spec, got {len(self.submodules.layer_specs)}. "
+                    "Repeated MTP mode expects exactly 1 layer spec, got "
+                    f"{len(self.submodules.layer_specs)} instead. "
                     f"The first layer will be applied {self.config.mtp_num_layers} times."
                 )
             self.layers = torch.nn.ModuleList(

From 906e51356bf32ab22c44c696766f085b6bc900cf Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Sat, 31 Jan 2026 10:31:56 -0800
Subject: [PATCH 46/48] bigfix: test failures

---
 .../core/models/common/language_module/language_module.py     | 4 ++--
 megatron/core/models/common/model_chunk_schedule_plan.py      | 2 +-
 megatron/core/models/gpt/fine_grained_callables.py            | 4 ++--
 megatron/core/pipeline_parallel/schedules.py                  | 2 +-
 tests/unit_tests/models/test_mamba_moe_model.py               | 2 ++
 5 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
index f94e51df69d..850639c0e75 100644
--- a/megatron/core/models/common/language_module/language_module.py
+++ b/megatron/core/models/common/language_module/language_module.py
@@ -313,7 +313,7 @@ def sharded_state_dict(
         # Also, if MTP is not placed in the post processing stage, we need to maintain a copy
         # of output layer in the mtp process stage and tie it to the output layer in the post
         # processing stage.
-        if self.mtp_process and not self.pre_process:
+        if getattr(self, 'mtp_process', False) and not self.pre_process:
             emb_weight = self.embedding.word_embeddings.weight
             tie_word_embeddings_state_dict(
                 sharded_state_dict,
@@ -323,7 +323,7 @@ def sharded_state_dict(
                 dp_cp_group=metadata['dp_cp_group'],
             )
 
-        if self.mtp_process and not self.post_process:
+        if getattr(self, 'mtp_process', False) and not self.post_process:
             # We only need to tie the output layer weight if share_embeddings_and_output_weights
             # is False. Because if share_embeddings_and_output_weights is True, the shared weight
             # will be stored in embedding layer, and output layer will not have any weight.
diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py
index 033e8e808f9..3b0e3a13b76 100644
--- a/megatron/core/models/common/model_chunk_schedule_plan.py
+++ b/megatron/core/models/common/model_chunk_schedule_plan.py
@@ -123,7 +123,7 @@ def _build_callable_nodes(self, event, comp_stream, comm_stream, extra_args):
         # get flags for latter use
         is_mtp = isinstance(self.layer, MultiTokenPredictionLayer)
         is_moe = (
-            isinstance(self.layer.transformer_layer.mlp, MoELayer)
+            isinstance(self.layer.mtp_model_layer.mlp, MoELayer)
             if is_mtp
             else isinstance(self.layer.mlp, MoELayer)
         )
diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py
index 7cee9d2973c..e17ed0a5d40 100644
--- a/megatron/core/models/gpt/fine_grained_callables.py
+++ b/megatron/core/models/gpt/fine_grained_callables.py
@@ -613,9 +613,9 @@ def build_mtp_layer_callables(layer):
     multi-token prediction layer nodes (attention, MLP, etc.)
     """
 
-    forward_funcs, backward_dw = build_transformer_layer_callables(layer.transformer_layer)
+    forward_funcs, backward_dw = build_transformer_layer_callables(layer.mtp_model_layer)
     attn_forward, dispatch_forward, mlp_forward, combine_forward, _ = forward_funcs
-    is_moe = isinstance(layer.transformer_layer.mlp, MoELayer)
+    is_moe = isinstance(layer.mtp_model_layer.mlp, MoELayer)
     assert is_moe, "MTP layer in a2a overlap only supports MoE layer for now."
 
     def submodule_mtp_attn_forward(node, hidden_states):
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index fd39b1136fe..15c5adfc7a2 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -406,7 +406,7 @@ def forward_step(
 
     if is_first_microbatch and hasattr(model, 'set_is_first_microbatch'):
         model.set_is_first_microbatch()
-    if config.external_cuda_graph and current_microbatch is not None:
+    if current_microbatch is not None:
         set_current_microbatch(model, current_microbatch)
 
     unwrap_output_tensor = False
diff --git a/tests/unit_tests/models/test_mamba_moe_model.py b/tests/unit_tests/models/test_mamba_moe_model.py
index 3c7ae93a17c..9c581ec6cb4 100644
--- a/tests/unit_tests/models/test_mamba_moe_model.py
+++ b/tests/unit_tests/models/test_mamba_moe_model.py
@@ -193,9 +193,11 @@
     "moe_z_loss_coeff": None,
     "moe_enable_routing_replay": False,
     "mrope_section": None,
+    "mtp_hybrid_override_pattern": None,
     "mtp_loss_scaling_factor": 0.1,
     "mtp_num_layers": None,
     "mtp_standalone": False,
+    "mtp_use_repeated_layer": False,
     "multi_latent_attention": False,
     "no_rope_freq": None,
     "no_sync_func": None,

From 4bfc5ec42668f6f409d63990e7988bc66a5210ac Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Sat, 31 Jan 2026 23:04:40 -0800
Subject: [PATCH 47/48] bugfix: fixing standalone mtp test failures

---
 .../common/language_module/language_module.py | 27 +++----------------
 .../ssm/test_mamba_hybrid_layer_allocation.py |  2 +-
 2 files changed, 5 insertions(+), 24 deletions(-)

diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
index 850639c0e75..57975b2958b 100644
--- a/megatron/core/models/common/language_module/language_module.py
+++ b/megatron/core/models/common/language_module/language_module.py
@@ -23,10 +23,7 @@
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.enums import AttnBackend, CudaGraphScope
 from megatron.core.transformer.module import MegatronModule
-from megatron.core.transformer.multi_token_prediction import (
-    tie_output_layer_state_dict,
-    tie_word_embeddings_state_dict,
-)
+from megatron.core.transformer.multi_token_prediction import tie_word_embeddings_state_dict
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.transformer.utils import ensure_metadata_has_dp_cp_group
 from megatron.core.utils import (
@@ -305,14 +302,12 @@ def sharded_state_dict(
         output_layer_weight_key = f'{prefix}output_layer.weight'
         output_layer_bias_key = f'{prefix}output_layer.bias'
 
-        # Multi-Token Prediction (MTP) need both embedding layer and output layer in
-        # mtp process stage.
+        # Multi-Token Prediction (MTP) needs embedding layer in mtp process stage.
         # If MTP is not placed in the pre processing stage, we need to maintain a copy of
         # embedding layer in the mtp process stage and tie it to the embedding in the pre
         # processing stage.
-        # Also, if MTP is not placed in the post processing stage, we need to maintain a copy
-        # of output layer in the mtp process stage and tie it to the output layer in the post
-        # processing stage.
+        # Note: MTP loss is computed at post_process stage, so the output_layer on mtp_process
+        # rank doesn't need special tying - it's not used for loss computation.
         if getattr(self, 'mtp_process', False) and not self.pre_process:
             emb_weight = self.embedding.word_embeddings.weight
             tie_word_embeddings_state_dict(
@@ -322,20 +317,6 @@ def sharded_state_dict(
                 tp_group=self.tp_group,
                 dp_cp_group=metadata['dp_cp_group'],
             )
-
-        if getattr(self, 'mtp_process', False) and not self.post_process:
-            # We only need to tie the output layer weight if share_embeddings_and_output_weights
-            # is False. Because if share_embeddings_and_output_weights is True, the shared weight
-            # will be stored in embedding layer, and output layer will not have any weight.
-            if not self.share_embeddings_and_output_weights:
-                output_layer_weight = self.output_layer.weight
-                tie_output_layer_state_dict(
-                    sharded_state_dict,
-                    output_layer_weight,
-                    output_layer_weight_key,
-                    tp_group=self.tp_group,
-                    dp_cp_group=metadata['dp_cp_group'],
-                )
         if self.share_embeddings_and_output_weights:
             self.tie_embeddings_and_output_weights_state_dict(
                 sharded_state_dict, output_layer_weight_key, first_stage_word_emb_key, metadata
diff --git a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
index 2119c0f6d23..afa5d38c19d 100644
--- a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
+++ b/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
@@ -127,7 +127,7 @@ def test_main_with_multiple_mtp_depths(self):
             ("M*M*/MM/MM", "M*M*", "MM", 2),
             ("M*M*/MM/MM/MM", "M*M*", "MM", 3),
             ("MMMM/*M/*M/*M", "MMMM", "*M", 3),
-            ("M*/*/*/*", "M*", "*", 4),
+            ("M*/*/*/*", "M*", "*", 3),
             ("M/M/M/M/M", "M", "M", 4),
         ]
         for pattern, expected_main, expected_mtp, expected_depths in test_cases:

From 9cc1668ca68dac49027e7ef40dd13eb675def782 Mon Sep 17 00:00:00 2001
From: Sanjeev Satheesh <sasatheesh@nvidia.com>
Date: Sun, 1 Feb 2026 09:58:48 -0800
Subject: [PATCH 48/48] bugfix: expected value error strings

---
 .../ssm/test_mamba_hybrid_layer_allocation.py    | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
index afa5d38c19d..77c106c3bee 100644
--- a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
+++ b/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
@@ -160,10 +160,18 @@ def test_invalid_symbols_in_main_pattern(self):
 
     def test_invalid_symbols_in_mtp_pattern(self):
         """Test that invalid symbols in MTP pattern raise ValueError."""
-        invalid_patterns = ["M*M*/MX", "M*M*/MM/Ma"]  # X is not valid  # a is not valid
-        for pattern in invalid_patterns:
-            with pytest.raises(ValueError, match="not a valid layer symbol"):
-                parse_hybrid_pattern(pattern)
+        # Single MTP depth with invalid symbol - should raise "not a valid layer symbol"
+        with pytest.raises(ValueError, match="not a valid layer symbol"):
+            parse_hybrid_pattern("M*M*/MX")  # X is not valid
+
+        # Multiple MTP depths with invalid symbol and matching patterns
+        with pytest.raises(ValueError, match="not a valid layer symbol"):
+            parse_hybrid_pattern("M*M*/Ma/Ma")  # a is not valid
+
+        # Multiple MTP depths with invalid symbol but mismatched patterns
+        # This raises "All MTP patterns must be identical" before checking symbols
+        with pytest.raises(ValueError, match="All MTP patterns must be identical"):
+            parse_hybrid_pattern("M*M*/MM/Ma")
 
     def test_empty_main_pattern_with_mtp(self):
         """Test pattern that starts with / (empty main pattern)."""