fix(mtp): resolve the MTP template from the full decoder layer list

src/mcore_bridge/model/gpts/deepseek_v4.py

@@ -30,6 +30,12 @@
     apply_module = None
     off_interface = None
 
+try:
+    from megatron.core.models.gpt.experimental_attention_variant_module_specs import \
+        get_transformer_layer_with_experimental_attention_variant_spec as _build_all_layer_specs
+except ImportError:  # megatron-core 0.18: block-level entry point only
+    _build_all_layer_specs = None
+
 
 @contextmanager
 def _patch_YarnRotaryEmbedding(config):
@@ -99,9 +105,11 @@ def _apply_mla_rope(t, freqs, *, config, cu_seqlens, cp_group, inverse=False):
 class DSv4HybridSelfAttention(McoreDSv4HybridSelfAttention):
 
     def __init__(self, config, *args, **kwargs):
-        assert McoreDSv4HybridSelfAttention is not object, (
-            'Please install the Megatron-Core dev branch: '
-            '`pip install git+https://github.com/NVIDIA/Megatron-LM@dev`')
+        if McoreDSv4HybridSelfAttention is object:
+            raise ImportError('DeepSeek-V4 requires megatron-core >= 0.20: '
+                              '`experimental_attention_variant.deepseek_v4_hybrid_attention` entered the release in '
+                              '0.20 (0.18/0.19 only ship the standard GPT skeleton). '
+                              'Install `megatron-core>=0.20`, or Megatron-LM@dev.')
         with _patch_YarnRotaryEmbedding(config):
             super().__init__(config, *args, **kwargs)
         # ``layer_types`` is an HF-space (length ``num_layers``) list. On the HybridStack the layer
@@ -530,19 +538,23 @@ def _set_inv_freq(self):
 class DeepseekV4Loader(ModelLoader):
     model_cls = DeepseekV4GPTModel
 
-    def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
-        from megatron.core.models.gpt.experimental_attention_variant_module_specs import \
-            get_transformer_block_with_experimental_attention_variant_spec
-        transformer_layer_spec = get_transformer_block_with_experimental_attention_variant_spec(self.config, vp_stage)
-        for layer_spec in transformer_layer_spec.layer_specs:
+    def _build_decoder_layer_specs(self):
+        if _build_all_layer_specs is None or McoreDSv4HybridSelfAttention is object:
+            raise NotImplementedError(
+                'DeepSeek-V4 requires megatron-core >= 0.20: it needs both the layer-level '
+                'experimental entry point and `experimental_attention_variant.deepseek_v4_hybrid_attention`.')
+        return self._patch_dsv4_layer_specs([copy.deepcopy(spec) for spec in _build_all_layer_specs(self.config)])
+
+    def _patch_dsv4_layer_specs(self, layer_specs):
+        for layer_spec in layer_specs:
             layer_spec.submodules.self_attention.module = DSv4HybridSelfAttention
             core_attention_submodules = layer_spec.submodules.self_attention.submodules.core_attention.submodules
             if getattr(core_attention_submodules, 'compressor', None) is not None:
                 core_attention_submodules.compressor.module = Compressor
             if getattr(core_attention_submodules, 'indexer', None) is not None:
                 core_attention_submodules.indexer.module = CSAIndexer
                 core_attention_submodules.indexer.submodules.compressor.module = Compressor
-        return transformer_layer_spec
+        return layer_specs
 
 
 class DeepseekV4Bridge(GPTBridge):

src/mcore_bridge/model/gpts/glm4.py

@@ -89,12 +89,12 @@ def _set_layer_attn(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: boo
 
 class Glm4Loader(ModelLoader):
 
-    def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
-        transformer_layer_spec = super().get_transformer_layer_spec(vp_stage)
-        for layer_spec in transformer_layer_spec.layer_specs:
+    def _build_decoder_layer_specs(self):
+        layer_specs = super()._build_decoder_layer_specs()
+        for layer_spec in layer_specs:
             layer_spec.submodules.self_attention.module = Glm4SelfAttention
             self._set_mlp_spec(layer_spec.submodules, Glm4MLP)
-        return transformer_layer_spec
+        return layer_specs
 
 
 register_model(ModelMeta(

src/mcore_bridge/model/gpts/glm_moe_dsa.py

@@ -3,7 +3,6 @@
                                                                           DSAIndexerLossLoggingHelper)
 from megatron.core.transformer.experimental_attention_variant.dsa import DSAttention as McoreDSAttention
 from megatron.core.transformer.experimental_attention_variant.dsa import FusedDSAIndexerLoss, unfused_dsa_fn
-from typing import Optional
 
 from ..constant import ModelType
 from ..register import ModelLoader, ModelMeta, register_model
@@ -20,19 +19,18 @@ def _get_index_share_carrier(self, packed_seq_params, attention_mask):
 
 class GlmMoeDsaLoader(ModelLoader):
 
-    def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
-        transformer_layer_spec = super().get_transformer_layer_spec(vp_stage)
-        if self.config.dsa_indexer_topk_freq > 1 and getattr(DSAttention, '_HOLDER_ATTR', None) is None:
-            raise ImportError(
-                'Please install the megatron-core main branch to support the "shared" indexer layer of `glm_moe_dsa`: '
-                '`pip install git+https://github.com/NVIDIA/Megatron-LM.git`')
+    def _build_decoder_layer_specs(self):
+        layer_specs = super()._build_decoder_layer_specs()
         if self.config.dsa_indexer_topk_freq > 1:
-            for layer_spec in transformer_layer_spec.layer_specs:
+            if getattr(DSAttention, '_HOLDER_ATTR', None) is None:
+                raise ImportError('glm_moe_dsa with dsa_indexer_topk_freq > 1 requires megatron-core >= 0.19: '
+                                  'the shared-indexer carrier (`DSAttention._HOLDER_ATTR`) entered the release in 0.19 '
+                                  'and is absent from 0.18.')
+            for layer_spec in layer_specs:
                 core_attn = layer_spec.submodules.self_attention.submodules.core_attention
                 if hasattr(core_attn, 'module') and issubclass(core_attn.module, McoreDSAttention):
                     core_attn.module = DSAttention
-
-        return transformer_layer_spec
+        return layer_specs
 
 
 register_model(ModelMeta(

src/mcore_bridge/model/gpts/minimax_m2.py

@@ -5,7 +5,6 @@
 from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.spec_utils import build_module
-from typing import Optional
 
 from mcore_bridge.bridge import GPTBridge
 from mcore_bridge.config import ModelConfig
@@ -103,11 +102,11 @@ def _set_moe_state(
 
 class MinimaxM2Loader(ModelLoader):
 
-    def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
-        transformer_layer_spec = super().get_transformer_layer_spec(vp_stage)
-        for layer_spec in transformer_layer_spec.layer_specs:
+    def _build_decoder_layer_specs(self):
+        layer_specs = super()._build_decoder_layer_specs()
+        for layer_spec in layer_specs:
             layer_spec.submodules.self_attention.module = MinimaxM2SelfAttention
-        return transformer_layer_spec
+        return layer_specs
 
 
 register_model(ModelMeta(

src/mcore_bridge/model/gpts/olmoe.py

@@ -5,7 +5,6 @@
 from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.spec_utils import build_module
-from typing import Optional
 
 from mcore_bridge.bridge import GPTBridge
 from mcore_bridge.config import ModelConfig
@@ -80,11 +79,11 @@ def _set_qk_layernorm(self, mg_attn, hf_state_dict, to_mcore, **kwargs):
 
 class OLMoELoader(ModelLoader):
 
-    def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
-        transformer_layer_spec = super().get_transformer_layer_spec(vp_stage)
-        for layer_spec in transformer_layer_spec.layer_specs:
+    def _build_decoder_layer_specs(self):
+        layer_specs = super()._build_decoder_layer_specs()
+        for layer_spec in layer_specs:
             layer_spec.submodules.self_attention.module = OLMoESelfAttention
-        return transformer_layer_spec
+        return layer_specs
 
 
 register_model(ModelMeta(

src/mcore_bridge/model/gpts/qwen3_next.py

@@ -13,14 +13,13 @@
 from megatron.core.tensor_parallel.random import get_cuda_rng_tracker
 from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
 from megatron.core.transformer.spec_utils import build_module
-from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
 from megatron.core.utils import deprecate_inference_params, is_fa_min_version, nvtx_range_pop, nvtx_range_push
 from transformers.utils import is_torch_npu_available
 from typing import Optional, Tuple, Union
 
 from mcore_bridge.bridge import GPTBridge
 from mcore_bridge.config import ModelConfig
-from mcore_bridge.utils import get_env_args, get_local_layer_specs, get_logger, get_num_samples
+from mcore_bridge.utils import get_env_args, get_logger, get_num_samples
 
 from ..constant import ModelType
 from ..register import ModelLoader, ModelMeta, register_model
@@ -525,17 +524,16 @@ def _convert_mtp_extra(self, mtp_layer, hf_state_dict, to_mcore, origin_hf_state
 
 class Qwen3NextLoader(ModelLoader):
     gated_delta_net = Qwen3NextGatedDeltaNet
+    final_layer_norm_cls = Qwen3NextRMSNorm
 
-    def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
+    def _build_decoder_layer_specs(self):
         config = self.config
         config.hetereogenous_dist_checkpoint = True
         # compat Qwen3NextGatedDeltaNet
         config.hidden_act = 'silu'
         config.rms_norm_eps = config.layernorm_epsilon
         config.dtype = config.params_dtype
 
-        # Use Zero-Centered RMSNorm to match HuggingFace exactly (no +1/-1 conversion needed)
-        layer_norm_impl = Qwen3NextRMSNorm
         moe_layer_spec = get_gpt_layer_with_transformer_engine_spec(
             num_experts=config.num_moe_experts,
             moe_grouped_gemm=config.moe_grouped_gemm,
@@ -544,34 +542,29 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
             use_kitchen=config.use_kitchen,
         )
         layer_specs = []
-        for is_linear_attention in self.config.linear_attention_freq:
+        for is_linear_attention in config.linear_attention_freq:
             layer_spec = deepcopy(moe_layer_spec)
             if is_linear_attention:
                 layer_spec.submodules.self_attention.module = self.gated_delta_net
             else:
                 layer_spec.submodules.self_attention.submodules.linear_qkv = TEColumnParallelLinear
                 layer_spec.submodules.self_attention.module = Qwen3NextSelfAttention
-            # Replace ALL layernorms with Qwen3NextRMSNorm (Zero-Centered)
-            layer_spec.submodules.input_layernorm = layer_norm_impl
+            # Zero-centered RMSNorm everywhere, matching HuggingFace without a +1/-1 conversion.
+            layer_spec.submodules.input_layernorm = Qwen3NextRMSNorm
             if hasattr(layer_spec.submodules, 'pre_mlp_layernorm'):
-                layer_spec.submodules.pre_mlp_layernorm = layer_norm_impl
+                layer_spec.submodules.pre_mlp_layernorm = Qwen3NextRMSNorm
             # qwen3.5 dense
             if config.hf_model_type == 'qwen3_5':
                 layer_spec.submodules.mlp.submodules.linear_fc1 = TEColumnParallelLinear
             # Replace qk_layernorm if present
             if hasattr(layer_spec.submodules.self_attention.submodules, 'q_layernorm'):
-                layer_spec.submodules.self_attention.submodules.q_layernorm = layer_norm_impl
+                layer_spec.submodules.self_attention.submodules.q_layernorm = Qwen3NextRMSNorm
             if hasattr(layer_spec.submodules.self_attention.submodules, 'k_layernorm'):
-                layer_spec.submodules.self_attention.submodules.k_layernorm = layer_norm_impl
+                layer_spec.submodules.self_attention.submodules.k_layernorm = Qwen3NextRMSNorm
             layer_specs.append(layer_spec)
+        return layer_specs
 
-        local_layer_specs = get_local_layer_specs(config, layer_specs, vp_stage=vp_stage)
-        block_spec = TransformerBlockSubmodules(layer_specs=local_layer_specs, layer_norm=layer_norm_impl)
-
-        return block_spec
-
-    def get_mtp_block_spec(self, *args, **kwargs):
-        mtp_block_spec = super().get_mtp_block_spec(*args, **kwargs)
+    def _patch_mtp_block_spec(self, mtp_block_spec):
         if mtp_block_spec is not None:
             for layer_spec in mtp_block_spec.layer_specs:
                 layer_spec.submodules.enorm = Qwen3NextRMSNorm

src/mcore_bridge/model/gpts/qwen3_next_gdn.py

@@ -16,6 +16,14 @@
 from ..register import ModelLoader, ModelMeta, register_model
 from .qwen3_next import Qwen3NextBridge
 
+try:
+    from megatron.core.models.gpt.experimental_attention_variant_module_specs import \
+        get_transformer_layer_with_experimental_attention_variant_spec as _build_all_layer_specs
+except ImportError:  # megatron-core 0.18: block-level entry point only
+    _build_all_layer_specs = None
+from megatron.core.models.gpt.experimental_attention_variant_module_specs import \
+    get_transformer_block_with_experimental_attention_variant_spec as _build_block_spec
+
 
 class Qwen3NextGDNBridgeMixin(GPTBridge):
     hf_mtp_prefix = 'mtp.layers'
@@ -115,11 +123,23 @@ def _set_linear_decoupled_in_proj(self, mg_attn, hf_state_dict, to_mcore: bool):
 class Qwen3NextLoader(ModelLoader):
     gated_delta_net = GatedDeltaNet
 
+    def _build_decoder_layer_specs(self):
+        if _build_all_layer_specs is None:
+            raise NotImplementedError('An MTP-only stage needs the full layer list, which requires the layer-level '
+                                      'experimental entry point (megatron-core >= 0.19). Put `m` on a stage that also '
+                                      'holds decoder layers, or upgrade megatron-core.')
+        layer_specs = [copy.deepcopy(spec) for spec in _build_all_layer_specs(self.config)]
+        return self._patch_gdn_layer_specs(layer_specs)
+
     def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
-        from megatron.core.models.gpt.experimental_attention_variant_module_specs import \
-            get_transformer_block_with_experimental_attention_variant_spec
-        layer_specs = get_transformer_block_with_experimental_attention_variant_spec(self.config, vp_stage)
-        for layer_spec in layer_specs.layer_specs:
+        if _build_all_layer_specs is not None:
+            return super().get_transformer_layer_spec(vp_stage)
+        block = _build_block_spec(self.config, vp_stage)
+        block.layer_specs = self._patch_layer_specs(self._patch_gdn_layer_specs(block.layer_specs))
+        return block
+
+    def _patch_gdn_layer_specs(self, layer_specs):
+        for layer_spec in layer_specs:
             attn_module = layer_spec.submodules.self_attention.module
             if issubclass(attn_module, SelfAttention):
                 layer_spec.submodules.self_attention.module = GatedSelfAttention

src/mcore_bridge/model/gpts/qwen4_exp.py

@@ -8,7 +8,7 @@
 from contextlib import contextmanager
 from copy import deepcopy
 from megatron.core.extensions.transformer_engine import TEColumnParallelLinear, TENorm, TERowParallelLinear
-from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec, get_gpt_mtp_block_spec
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.ssm.gated_delta_net import GatedDeltaNetSubmodules
 from megatron.core.tensor_parallel import gather_from_sequence_parallel_region
@@ -17,13 +17,12 @@
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.spec_utils import ModuleSpec
-from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
 from megatron.core.utils import make_viewless_tensor
 from torch import nn
 from transformers.utils import is_torch_npu_available
 from typing import List, Optional
 
-from mcore_bridge.utils import get_env_args, get_local_layer_specs, get_logger
+from mcore_bridge.utils import get_env_args, get_logger
 from mcore_bridge.utils.megatron_utils import reconstruct_tensor_cp
 
 from ..modules import (QSA_SPARSE_KERNEL_ENV, GatedDeltaNet, QSAIndexer, QSASparseCoreAttention,
@@ -680,6 +679,9 @@ def _convert_mtp_layer(self, lm_model, hf_state_dict, hf_prefix: str, layer_idx:
 
 class Qwen4ExpLoader(ModelLoader):
     transformer_block = Qwen4ExpTransformerBlock
+    # No final layernorm in this model; keep the slot so the HC mixer stage logic
+    # (has_final_layernorm_in_this_stage) still triggers.
+    final_layer_norm_cls = IdentityOp
 
     def _get_moe_layer_pattern(self) -> List[bool]:
         config = self.config
@@ -689,7 +691,7 @@ def _get_moe_layer_pattern(self) -> List[bool]:
         # int N: one MoE every N layers (mcore convention: i % N == N - 1).
         return [i % freq == freq - 1 for i in range(config.num_layers)]
 
-    def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
+    def _build_decoder_layer_specs(self):
         config = self.config
         config.hetereogenous_dist_checkpoint = True
         if config.context_parallel_size > 1 and getattr(config, 'cp_comm_type', None) in (None, 'p2p'):
@@ -741,17 +743,13 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
                 layer_spec.submodules.pre_mlp_layernorm = IdentityOp
             layer_specs.append(layer_spec)
 
-        local_layer_specs = get_local_layer_specs(config, layer_specs, vp_stage=vp_stage)
-        # No final layernorm in this model; keep the slot so the HC mixer
-        # stage logic (has_final_layernorm_in_this_stage) still triggers.
-        block_spec = TransformerBlockSubmodules(layer_specs=local_layer_specs, layer_norm=IdentityOp)
-        return block_spec
+        return layer_specs
 
     def _set_transformer_layer(self, transformer_layer_spec):
         for layer_spec in transformer_layer_spec.layer_specs:
             layer_spec.module = Qwen4ExpLayer
 
-    def get_mtp_block_spec(self, transformer_layer_spec, vp_stage: Optional[int] = None):
+    def _patch_mtp_block_spec(self, mtp_block_spec):
         if (self.config.mtp_num_layers or 0) > 1 and not self.config.mtp_shared_weights:
             # HF exposes a single top-level `mtp.*` head; multiple independent heads would all map
             # to the same `mtp.*` keys on export and overwrite each other (no lossless round-trip).
@@ -760,15 +758,14 @@ def get_mtp_block_spec(self, transformer_layer_spec, vp_stage: Optional[int] = N
                 'MTP heads (mtp_num_layers>1 without mtp_shared_weights) cannot round-trip through HF. '
                 'Use mtp_shared_weights=True to reuse one head across prediction depths, or '
                 'mtp_num_layers=1.')
-        mtp_block_spec = get_gpt_mtp_block_spec(
-            self.config, transformer_layer_spec, use_transformer_engine=True, vp_stage=vp_stage)
         if mtp_block_spec is not None:
             for layer_spec in mtp_block_spec.layer_specs:
                 sub = layer_spec.submodules
                 if not (hasattr(sub, 'e_proj') and hasattr(sub, 'h_proj')):
-                    raise NotImplementedError('Qwen3.8-Flash-Next MTP requires a Megatron dev build exposing separate '
-                                              f'e_proj/h_proj, got {megatron.core.__version__}.'
-                                              'run pip install "git+https://github.com/NVIDIA/Megatron-LM.git@dev"')
+                    raise NotImplementedError(
+                        'Qwen3.8-Flash-Next MTP requires megatron-core >= 0.20, whose MTP layer spec '
+                        'exposes separate e_proj/h_proj (added in the 0.20 release); '
+                        f'got {megatron.core.__version__}.')
                 layer_spec.module = Qwen4ExpMultiTokenPredictionLayer
                 # residual_linear_shared head: separate e_proj/h_proj (fc_embedding/fc_hidden), a joint
                 # multi-stream hnorm (pre_fc_norm_hidden over n*H), and no fused eh_proj. layer_norm