From 7e9e79de37b2077e523dec1d32aa450ee2cd07ba Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Fri, 6 Nov 2026 16:49:54 +0800 Subject: [PATCH] fix --- src/mcore_bridge/model/gpts/deepseek_v4.py | 30 +- src/mcore_bridge/model/gpts/glm4.py | 8 +- src/mcore_bridge/model/gpts/glm_moe_dsa.py | 18 +- src/mcore_bridge/model/gpts/minimax_m2.py | 9 +- src/mcore_bridge/model/gpts/olmoe.py | 9 +- src/mcore_bridge/model/gpts/qwen3_next.py | 29 +- src/mcore_bridge/model/gpts/qwen3_next_gdn.py | 28 +- src/mcore_bridge/model/gpts/qwen4_exp.py | 27 +- src/mcore_bridge/model/mm_gpts/gemma4.py | 1 + src/mcore_bridge/model/mm_gpts/llama4.py | 15 +- .../model/mm_gpts/muse_glimmer.py | 1 + src/mcore_bridge/model/mtp.py | 8 +- src/mcore_bridge/model/register.py | 132 +++++---- src/mcore_bridge/utils/megatron_utils.py | 1 + tests/test_mtp_layout_spec.py | 267 ++++++++++++++++++ 15 files changed, 450 insertions(+), 133 deletions(-) create mode 100644 tests/test_mtp_layout_spec.py diff --git a/src/mcore_bridge/model/gpts/deepseek_v4.py b/src/mcore_bridge/model/gpts/deepseek_v4.py index 77e2e7c0..21903b73 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v4.py +++ b/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,11 +538,15 @@ 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: @@ -542,7 +554,7 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): 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): diff --git a/src/mcore_bridge/model/gpts/glm4.py b/src/mcore_bridge/model/gpts/glm4.py index 1da91af8..a96405d4 100644 --- a/src/mcore_bridge/model/gpts/glm4.py +++ b/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( diff --git a/src/mcore_bridge/model/gpts/glm_moe_dsa.py b/src/mcore_bridge/model/gpts/glm_moe_dsa.py index 0a748578..425f03a3 100644 --- a/src/mcore_bridge/model/gpts/glm_moe_dsa.py +++ b/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( diff --git a/src/mcore_bridge/model/gpts/minimax_m2.py b/src/mcore_bridge/model/gpts/minimax_m2.py index 853e3a59..1043bc4a 100644 --- a/src/mcore_bridge/model/gpts/minimax_m2.py +++ b/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( diff --git a/src/mcore_bridge/model/gpts/olmoe.py b/src/mcore_bridge/model/gpts/olmoe.py index dd30ef36..9018ac01 100644 --- a/src/mcore_bridge/model/gpts/olmoe.py +++ b/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( diff --git a/src/mcore_bridge/model/gpts/qwen3_next.py b/src/mcore_bridge/model/gpts/qwen3_next.py index f352f4fe..2d827d26 100644 --- a/src/mcore_bridge/model/gpts/qwen3_next.py +++ b/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,8 +524,9 @@ 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 @@ -534,8 +534,6 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): 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 diff --git a/src/mcore_bridge/model/gpts/qwen3_next_gdn.py b/src/mcore_bridge/model/gpts/qwen3_next_gdn.py index e8843ef7..d42d0b52 100644 --- a/src/mcore_bridge/model/gpts/qwen3_next_gdn.py +++ b/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 diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index 451b7d52..cd1b131d 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/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 diff --git a/src/mcore_bridge/model/mm_gpts/gemma4.py b/src/mcore_bridge/model/mm_gpts/gemma4.py index 79d6fc6c..91d1607e 100644 --- a/src/mcore_bridge/model/mm_gpts/gemma4.py +++ b/src/mcore_bridge/model/mm_gpts/gemma4.py @@ -843,6 +843,7 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): for layer_spec, moe_layer_spec in zip(layer_specs.layer_specs, moe_layer_specs.layer_specs): layer_spec.submodules.experts_mlp = moe_layer_spec.submodules.mlp self._set_mlp_spec(layer_spec.submodules, Gemma4MoELayer, mlp_key='experts_mlp') + layer_specs.layer_specs = self._patch_layer_specs(layer_specs.layer_specs) return layer_specs def _set_transformer_layer(self, transformer_layer_spec): diff --git a/src/mcore_bridge/model/mm_gpts/llama4.py b/src/mcore_bridge/model/mm_gpts/llama4.py index 7c508e2e..41c817bb 100644 --- a/src/mcore_bridge/model/mm_gpts/llama4.py +++ b/src/mcore_bridge/model/mm_gpts/llama4.py @@ -1,10 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import torch -from copy import deepcopy from megatron.core.transformer.identity_op import IdentityOp -from megatron.core.transformer.transformer_layer import get_transformer_layer_offset from transformers import PretrainedConfig -from typing import Optional from mcore_bridge.bridge import GPTBridge @@ -70,16 +67,12 @@ class Llama4Bridge(GPTBridge): class Llama4Loader(ModelLoader): - def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): - layer_specs = super().get_transformer_layer_spec(vp_stage) - for i, layer_spec in enumerate(layer_specs.layer_specs): - global_i = i + get_transformer_layer_offset(self.config, vp_stage) - no_rope = self.config.no_rope_freq[global_i] - layer_spec = deepcopy(layer_spec) - if no_rope: + def _build_decoder_layer_specs(self): + layer_specs = super()._build_decoder_layer_specs() + for layer_idx, layer_spec in enumerate(layer_specs): + if self.config.no_rope_freq[layer_idx]: layer_spec.submodules.self_attention.submodules.q_layernorm = IdentityOp layer_spec.submodules.self_attention.submodules.k_layernorm = IdentityOp - layer_specs.layer_specs[i] = layer_spec return layer_specs diff --git a/src/mcore_bridge/model/mm_gpts/muse_glimmer.py b/src/mcore_bridge/model/mm_gpts/muse_glimmer.py index 5c79c9e9..f7dd3675 100644 --- a/src/mcore_bridge/model/mm_gpts/muse_glimmer.py +++ b/src/mcore_bridge/model/mm_gpts/muse_glimmer.py @@ -318,6 +318,7 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): self.config, use_transformer_engine=True, normalization=self.config.normalization, vp_stage=vp_stage) for layer_spec in layer_specs.layer_specs: layer_spec.submodules.self_attention.module = MuseGlimmerSelfAttention + layer_specs.layer_specs = self._patch_layer_specs(layer_specs.layer_specs) return layer_specs def _set_transformer_layer(self, transformer_layer_spec): diff --git a/src/mcore_bridge/model/mtp.py b/src/mcore_bridge/model/mtp.py index 6996af89..a7674883 100644 --- a/src/mcore_bridge/model/mtp.py +++ b/src/mcore_bridge/model/mtp.py @@ -35,14 +35,14 @@ class MtpCapabilities: """Signatures of the installed upstream MTP layer/block. - Releases (0.18/0.19) and dev disagree on what the layer forward accepts and + Releases (0.18-0.20) and dev disagree on what the layer forward accepts and returns, so the bridge probes the installed code instead of branching on a version number. """ has_detach_heads: bool = False - # Knobs that differ across 0.20 / main / dev; gate on these, not on version numbers. + # Knobs that differ across releases (0.18-0.20) and dev; gate on these, not on version numbers. has_mtp_hybrid_override_pattern: bool = False has_freeze_base_model_for_mtp: bool = False has_mtp_loss_early_backward: bool = False @@ -511,8 +511,8 @@ class MultiTokenPredictionBlock(McoreMultiTokenPredictionBlock): Per-forward bridge state lives on block attributes (no upstream block/layer forward takes ``**kwargs``) and is read by :class:`MultiTokenPredictionLayer` through a - back-reference. Releases (0.18/0.19) get the manual per-depth loop; on dev - (native ``roll_depth`` / ``mhc_multistream``) the loop is delegated to upstream. + back-reference. 0.18/0.19 get the bridge's manual per-depth loop; 0.20+ and dev + (native ``roll_depth`` / ``mhc_multistream``) delegate the loop to upstream. """ # Per-forward bridge state, consumed by MultiTokenPredictionLayer via _bridge_block. diff --git a/src/mcore_bridge/model/register.py b/src/mcore_bridge/model/register.py index 1773c21f..50f55aad 100644 --- a/src/mcore_bridge/model/register.py +++ b/src/mcore_bridge/model/register.py @@ -5,23 +5,32 @@ from functools import partial from megatron.core import mpu from megatron.core.enums import ModelType -from megatron.core.extensions.transformer_engine import TEGroupedLinear, TELayerNormColumnParallelLinear, TELinear +from megatron.core.extensions.transformer_engine import (TEGroupedLinear, TELayerNormColumnParallelLinear, TELinear, + TENorm) from megatron.core.models.gpt import gpt_model -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec, get_gpt_mtp_block_spec +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_layer_specs, get_gpt_mtp_block_spec from megatron.core.transformer.moe.router import TopKRouter as McoreTopKRouter from megatron.core.transformer.multi_latent_attention import MLASelfAttention as McoreMLASelfAttention +from megatron.core.transformer.multi_token_prediction import get_mtp_num_layers_to_build +from megatron.core.transformer.transformer_block import TransformerBlockSubmodules from megatron.core.transformer.transformer_layer import TransformerLayer as McoreTransformerLayer from torch import nn from typing import TYPE_CHECKING, List, Optional, Type, Union from mcore_bridge.bridge import GPTBridge from mcore_bridge.config import ModelConfig -from mcore_bridge.utils import get_logger +from mcore_bridge.utils import get_local_layer_specs, get_logger from .modules import (AbsorbedMLASelfAttention, DSAIndexer, MLASelfAttention, TopKRouter, TransformerBlock, TransformerLayer) from .mtp import adapt_mtp_block_spec, get_mtp_block_cls +try: + from megatron.core.models.gpt.experimental_attention_variant_module_specs import \ + get_transformer_layer_with_experimental_attention_variant_spec as _build_experimental_layer_specs +except ImportError: # megatron-core 0.18: layer-level entry point absent + _build_experimental_layer_specs = None + if TYPE_CHECKING: from .gpt_model import GPTModel from .mm_gpt_model import MultimodalGPTModel @@ -82,6 +91,7 @@ def get_model_meta(mcore_model_type: str) -> ModelMeta: class ModelLoader: model_cls = None transformer_block = TransformerBlock + final_layer_norm_cls = TENorm def __init__(self, config: ModelConfig): from mcore_bridge.model import GPTModel, MultimodalGPTModel @@ -99,56 +109,84 @@ def _set_mlp_spec(self, layer_submodules, mlp_module, mlp_key='mlp'): else: mlp_spec.module = mlp_module - def _replace_spec_dsa(self, layer_spec): - from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( - _get_backend_spec_provider, get_dsa_module_spec_for_backend) - backend = _get_backend_spec_provider(config=self.config) - dsa_spec = get_dsa_module_spec_for_backend(self.config, backend) - if getattr(dsa_spec.submodules.core_attention.submodules, 'indexer', None) is not None: - dsa_spec.submodules.core_attention.submodules.indexer.module = DSAIndexer - if self.config.qk_layernorm: - linear_q_up_proj = backend.column_parallel_linear() - # fix megatron-core - dsa_spec.submodules.q_layernorm = backend.layer_norm(for_qk=True) - dsa_spec.submodules.kv_layernorm = backend.layer_norm(for_qk=True) - dsa_spec.submodules.linear_q_up_proj = linear_q_up_proj - dsa_spec.submodules.linear_kv_up_proj = linear_q_up_proj - layer_spec.submodules.self_attention = dsa_spec - - @contextmanager - def _patch_experimental_attention_variant(self): - experimental_attention_variant = self.config.experimental_attention_variant - self.config.experimental_attention_variant = None - try: - yield - finally: - self.config.experimental_attention_variant = experimental_attention_variant - - def _deepcopy_layer_spec(self, transformer_layer_spec): - for i, layer_spec in enumerate(transformer_layer_spec.layer_specs): - transformer_layer_spec.layer_specs[i] = deepcopy(layer_spec) + def decoder_layer_specs(self): + """All decoder layer specs, model-patched, independent of pp/vpp slicing. + + Final entry point. Subclasses extend `_build_decoder_layer_specs` (construction) and, + rarely, `_patch_layer_specs` (class replacement) so the model-level patch below can + never be skipped. + """ + return self._patch_layer_specs(self._build_decoder_layer_specs()) + + def _build_decoder_layer_specs(self): + """Construct every decoder layer; subclasses replace this when upstream cannot.""" + variant = self.config.experimental_attention_variant + if variant is None: + layer_specs = self._plain_decoder_layer_specs() + elif _build_experimental_layer_specs is None: + raise NotImplementedError(f'{type(self).__name__} with experimental_attention_variant={variant!r} requires ' + 'megatron-core >= 0.19 (the layer-level experimental entry point).') + else: + layer_specs = _build_experimental_layer_specs(self.config) + layer_specs = [deepcopy(layer_spec) for layer_spec in layer_specs] + if variant == 'dsa': + self._replace_dsa_indexer(layer_specs) + return layer_specs + + def _plain_decoder_layer_specs(self): + return get_gpt_decoder_layer_specs( + self.config, + use_transformer_engine=True, + normalization=self.config.normalization, + qk_l2_norm=self.config.qk_l2_norm) + + def _replace_dsa_indexer(self, layer_specs): + # Upstream's indexer lacks bridge's FP8 init / SP gather / rope mscale. + for layer_spec in layer_specs: + core_attention = layer_spec.submodules.self_attention.submodules.core_attention + core_attention.submodules.indexer.module = DSAIndexer + + def _patch_layer_specs(self, layer_specs): + # The three helpers below take a block spec, so wrap the list once. + block = TransformerBlockSubmodules(layer_specs=layer_specs) + self._set_transformer_layer(block) + self._replace_mla_attention(block) + self._replace_router(block) + return layer_specs def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): - with self._patch_experimental_attention_variant(): - transformer_layer_spec = get_gpt_decoder_block_spec( - self.config, - use_transformer_engine=True, - normalization=self.config.normalization, - qk_l2_norm=self.config.qk_l2_norm, - vp_stage=vp_stage) - self._deepcopy_layer_spec(transformer_layer_spec) - if self.config.experimental_attention_variant == 'dsa': - for layer_spec in transformer_layer_spec.layer_specs: - self._replace_spec_dsa(layer_spec) - return transformer_layer_spec + """The local block of one pipeline stage (historical name: returns a block spec).""" + return TransformerBlockSubmodules( + layer_specs=get_local_layer_specs(self.config, self.decoder_layer_specs(), vp_stage=vp_stage), + layer_norm=self.final_layer_norm_cls) def get_mtp_block_spec(self, transformer_layer_spec, vp_stage: Optional[int] = None): + # Upstream returns None for stages that hold no MTP layer. Check that first so + # embedding-only / loss-only stages never build a full spec (or hit a version error). + if get_mtp_num_layers_to_build(self.config, vp_stage=vp_stage) == 0: + return None + if not transformer_layer_spec.layer_specs: + cls = type(self) + if (cls._build_decoder_layer_specs is ModelLoader._build_decoder_layer_specs + and cls.get_transformer_layer_spec is not ModelLoader.get_transformer_layer_spec): + raise NotImplementedError( + f'{cls.__name__} only overrides get_transformer_layer_spec; its extra spec ' + 'changes never reach decoder_layer_specs(), so an MTP-only stage would build a ' + 'template that differs from the normal layout. Override `_build_decoder_layer_specs` ' + '(and call super()) instead.') + transformer_layer_spec = TransformerBlockSubmodules( + layer_specs=[self.decoder_layer_specs()[-1]], layer_norm=transformer_layer_spec.layer_norm) mtp_block_spec = get_gpt_mtp_block_spec( self.config, transformer_layer_spec, use_transformer_engine=True, vp_stage=vp_stage) - return adapt_mtp_block_spec(self.config, mtp_block_spec) + return self._patch_mtp_block_spec(adapt_mtp_block_spec(self.config, mtp_block_spec)) + + def _patch_mtp_block_spec(self, mtp_block_spec): + """Model-level MTP tweaks; subclasses replace the module/norm classes here.""" + return mtp_block_spec def _set_transformer_layer(self, transformer_layer_spec): - if self.config.enable_hyper_connections: + # HC models pick their own layer class; only the vanilla one is swapped. + if getattr(self.config, 'enable_hyper_connections', False): # field exists from 0.20 return for layer_spec in transformer_layer_spec.layer_specs: if layer_spec.module is McoreTransformerLayer: @@ -184,9 +222,6 @@ def build_model( vp_stage: Optional[int] = None, ) -> Union['GPTModel', 'MultimodalGPTModel']: transformer_layer_spec = self.get_transformer_layer_spec(vp_stage=vp_stage) - self._set_transformer_layer(transformer_layer_spec) - self._replace_mla_attention(transformer_layer_spec) - self._replace_router(transformer_layer_spec) mtp_block_spec = None if self.config.mtp_num_layers is not None: mtp_block_spec = self.get_mtp_block_spec(transformer_layer_spec, vp_stage=vp_stage) @@ -202,6 +237,7 @@ def build_model( self._set_linear_is_expert(model) return model + # GPTModel picks its block class at construction time; swap it while the model is built. @contextmanager def _patch_transformer_block(self): TransformerBlock = gpt_model.TransformerBlock diff --git a/src/mcore_bridge/utils/megatron_utils.py b/src/mcore_bridge/utils/megatron_utils.py index dd2af42d..09d1be15 100644 --- a/src/mcore_bridge/utils/megatron_utils.py +++ b/src/mcore_bridge/utils/megatron_utils.py @@ -136,6 +136,7 @@ def reconstruct_tensor_cp(tensor, packed_seq_params, dim: int, cp_partition_mode return out +# Upstream slices inside get_gpt_decoder_block_spec and exposes no helper; mirror that slice here. def get_local_layer_specs(config, layer_specs, vp_stage=None): num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage) diff --git a/tests/test_mtp_layout_spec.py b/tests/test_mtp_layout_spec.py new file mode 100644 index 00000000..917da9ad --- /dev/null +++ b/tests/test_mtp_layout_spec.py @@ -0,0 +1,267 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Regression tests for MTP layer-spec resolution on stages without decoder layers. + +An MTP-only pipeline stage used to hand an empty local block to upstream's MTP spec +builder, which indexes ``layer_specs[-1]`` and raises ``IndexError`` (ms-swift#9823). +The MTP template must come from the full, patched, un-sliced layer list instead. +""" +import importlib +import inspect +import pytest +import torch +from copy import deepcopy +from megatron.core import parallel_state +from megatron.core.transformer.transformer_block import TransformerBlockSubmodules +from megatron.core.transformer.transformer_config import TransformerConfig +from unittest import mock + +import mcore_bridge # noqa: F401 (applies runtime patches at import) +from mcore_bridge.model.modules import TransformerLayer as BridgeTransformerLayer +from mcore_bridge.model.register import ModelLoader, _build_experimental_layer_specs + + +def _config(layout='Ettt|mL', num_layers=3, pp=2): + return TransformerConfig( + pipeline_dtype=torch.float32, + num_layers=num_layers, + hidden_size=128, + num_attention_heads=4, + num_query_groups=4, + ffn_hidden_size=256, + normalization='RMSNorm', + add_bias_linear=False, + gated_linear_unit=True, + pipeline_model_parallel_size=pp, + pipeline_model_parallel_layout=layout, + mtp_num_layers=1, + ) + + +def _loader(loader_cls, config): + # Bypass __init__: these tests only exercise spec resolution, not model building. + loader = loader_cls.__new__(loader_cls) + loader.config = config + return loader + + +def _rank1_last_stage(): + return mock.patch.multiple( + parallel_state, + get_pipeline_model_parallel_rank=mock.Mock(return_value=1), + get_virtual_pipeline_model_parallel_world_size=mock.Mock(return_value=None), + ) + + +def test_decoder_layer_specs_has_no_stage_argument(): + """A stage argument would turn the full layer list back into a local slice.""" + assert 'vp_stage' not in inspect.signature(ModelLoader.decoder_layer_specs).parameters + + +def test_decoder_layer_specs_returns_every_layer_regardless_of_stage(): + loader = _loader(ModelLoader, _config()) + assert len(loader.decoder_layer_specs()) == 3 + + +def test_experimental_variant_uses_upstream_layer_entry(): + """Variant models go through upstream's layer-level experimental entry (>= 0.19).""" + if _build_experimental_layer_specs is None: + pytest.skip('layer-level experimental entry requires megatron-core >= 0.19') + config = _config() + config.experimental_attention_variant = 'dsa' + config.multi_latent_attention = True + loader = _loader(ModelLoader, config) + assert len(loader.decoder_layer_specs()) == 3 + + +def test_dsa_attention_spec_is_not_dropped(): + """A plain-MLA fallback keeps the class name but loses the indexer, so check both.""" + if _build_experimental_layer_specs is None: + pytest.skip('layer-level experimental entry requires megatron-core >= 0.19') + config = _config() + config.experimental_attention_variant = 'dsa' + config.multi_latent_attention = True + config.qk_layernorm = True + loader = _loader(ModelLoader, config) + self_attention = loader.decoder_layer_specs()[0].submodules.self_attention + core = self_attention.submodules.core_attention + assert self_attention.module.__module__.startswith('mcore_bridge') + assert core.module.__name__ == 'DSAttention' + assert hasattr(core.submodules, 'indexer') + + +_MIGRATED_LOADERS = [ + ('mcore_bridge.model.gpts.glm4', 'Glm4Loader'), + ('mcore_bridge.model.gpts.olmoe', 'OLMoELoader'), + ('mcore_bridge.model.gpts.minimax_m2', 'MinimaxM2Loader'), + ('mcore_bridge.model.mm_gpts.llama4', 'Llama4Loader'), + ('mcore_bridge.model.gpts.glm_moe_dsa', 'GlmMoeDsaLoader'), + ('mcore_bridge.model.gpts.deepseek_v4', 'DeepseekV4Loader'), +] + +# Only GDN still needs the 0.18 block-level fallback; DeepSeek-V4 now requires 0.20 outright. +_MIGRATED_WITH_018_FALLBACK = [ + ('mcore_bridge.model.gpts.qwen3_next_gdn', 'Qwen3NextLoader'), +] + + +@pytest.mark.parametrize('module_name,class_name', _MIGRATED_LOADERS) +def test_migrated_loader_extends_the_full_list_only(module_name, class_name): + """A migrated loader must not keep a second slicing path.""" + loader_cls = getattr(importlib.import_module(module_name), class_name) + assert loader_cls._build_decoder_layer_specs is not ModelLoader._build_decoder_layer_specs + assert loader_cls.get_transformer_layer_spec is ModelLoader.get_transformer_layer_spec + + +@pytest.mark.parametrize('module_name,class_name', _MIGRATED_WITH_018_FALLBACK) +def test_loader_with_018_fallback_overrides_both(module_name, class_name): + loader_cls = getattr(importlib.import_module(module_name), class_name) + assert loader_cls._build_decoder_layer_specs is not ModelLoader._build_decoder_layer_specs + assert loader_cls.get_transformer_layer_spec is not ModelLoader.get_transformer_layer_spec + + +def test_glm_moe_dsa_swaps_shared_indexer_attention(): + """dsa_indexer_topk_freq > 1 must install the bridge's shared-indexer attention subclass.""" + from mcore_bridge.model.gpts.glm_moe_dsa import DSAttention, GlmMoeDsaLoader + + if getattr(DSAttention, '_HOLDER_ATTR', None) is None: + pytest.skip('shared indexer requires megatron-core >= 0.19') + config = _config() + config.experimental_attention_variant = 'dsa' + config.multi_latent_attention = True + config.qk_layernorm = True + config.dsa_indexer_topk_freq = 2 + loader = _loader(GlmMoeDsaLoader, config) + core = loader.decoder_layer_specs()[0].submodules.self_attention.submodules.core_attention + assert core.module is DSAttention + + +def test_non_mtp_empty_stage_returns_none_without_building_a_template(): + """Embedding-only stages must not build a full spec (or raise) just to find no template.""" + + class HalfLoader(ModelLoader): + + def get_transformer_layer_spec(self, vp_stage=None): + return TransformerBlockSubmodules(layer_specs=[]) + + loader = _loader(HalfLoader, _config('E|t|t|mL', num_layers=2)) + with mock.patch.multiple( + parallel_state, + get_pipeline_model_parallel_rank=mock.Mock(return_value=0), + get_virtual_pipeline_model_parallel_world_size=mock.Mock(return_value=2), + ): + assert loader.get_mtp_block_spec(TransformerBlockSubmodules(layer_specs=[]), vp_stage=0) is None + + +def test_mtp_only_stage_raises_only_when_the_stage_needs_mtp(): + """Simulated 0.18: the version error appears on MTP stages only, never on empty ones.""" + + class NoLayerLevelLoader(ModelLoader): + + def _build_decoder_layer_specs(self): + raise NotImplementedError('simulated 0.18: no layer-level entry') + + def get_transformer_layer_spec(self, vp_stage=None): + return TransformerBlockSubmodules(layer_specs=[]) + + loader = _loader(NoLayerLevelLoader, _config()) + with _rank1_last_stage(): + with pytest.raises(NotImplementedError, match='layer-level'): + loader.get_mtp_block_spec(TransformerBlockSubmodules(layer_specs=[]), vp_stage=0) + + +def test_template_uses_the_last_layer_not_the_first(): + """Heterogeneous fixture: only an implementation taking layer 0 would pass the marker.""" + + class MarkerAttention: + """Never built in the passing path; it only makes layer 0 distinguishable.""" + + class HeteroLoader(ModelLoader): + + def _build_decoder_layer_specs(self): + layer_specs = [deepcopy(spec) for spec in self._plain_decoder_layer_specs()] + layer_specs[0].submodules.self_attention.module = MarkerAttention + return layer_specs + + loader = _loader(HeteroLoader, _config()) + with _rank1_last_stage(): + local = loader.get_transformer_layer_spec(vp_stage=0) + mtp_block_spec = loader.get_mtp_block_spec(local, vp_stage=0) + inner = mtp_block_spec.layer_specs[0].submodules.mtp_model_layer + assert inner.submodules.self_attention.module is not MarkerAttention + + +def test_real_gdn_loader_mtp_only_uses_full_last_layer(): + """Real Qwen3.5 loader + MTP-only layout: the template must be the full list's last layer.""" + from mcore_bridge.model.gpts.qwen3_next import Qwen3NextLoader, Qwen3NextSelfAttention + + config = _config('Ett|t|t|mL', num_layers=4) + config.hf_model_type = 'qwen3_next' + config.linear_attention_freq = [1, 0, 1, 0] # last layer is full attention + config.hetereogenous_dist_checkpoint = False + config.layernorm_epsilon = 1e-6 + config.rms_norm_eps = 1e-6 + config.hidden_act = 'silu' + config.dtype = torch.float32 + config.params_dtype = torch.float32 + + loader = _loader(Qwen3NextLoader, config) + with mock.patch.multiple( + parallel_state, + get_pipeline_model_parallel_rank=mock.Mock(return_value=1), + get_virtual_pipeline_model_parallel_world_size=mock.Mock(return_value=2), + ): + local = loader.get_transformer_layer_spec(vp_stage=1) + assert local.layer_specs == [] + mtp_block_spec = loader.get_mtp_block_spec(local, vp_stage=1) + full_last = loader.decoder_layer_specs()[-1] + inner = mtp_block_spec.layer_specs[0].submodules.mtp_model_layer + assert inner.submodules.self_attention.module is Qwen3NextSelfAttention + assert inner.submodules.self_attention.module is full_last.submodules.self_attention.module + + +def test_mtp_only_stage_resolves_template_from_full_layer_list(): + """`Ettt|mL`: rank1 owns only MTP+loss; the local block is empty but MTP still builds.""" + loader = _loader(ModelLoader, _config()) + with _rank1_last_stage(): + local = loader.get_transformer_layer_spec(vp_stage=0) + assert local.layer_specs == [] + mtp_block_spec = loader.get_mtp_block_spec(local, vp_stage=0) + assert mtp_block_spec is not None + assert len(mtp_block_spec.layer_specs) == 1 + + +def test_mtp_template_carries_the_model_patch(): + """The template must be the patched layer, not upstream's vanilla TransformerLayer.""" + loader = _loader(ModelLoader, _config()) + with _rank1_last_stage(): + local = loader.get_transformer_layer_spec(vp_stage=0) + mtp_block_spec = loader.get_mtp_block_spec(local, vp_stage=0) + inner = mtp_block_spec.layer_specs[0].submodules.mtp_model_layer + assert inner.module is BridgeTransformerLayer + + +def test_template_is_the_last_layer_of_the_full_list(): + """The local block never participates: the template mirrors decoder_layer_specs()[-1].""" + loader = _loader(ModelLoader, _config()) + with _rank1_last_stage(): + local = loader.get_transformer_layer_spec(vp_stage=0) + assert local.layer_specs == [] + mtp_block_spec = loader.get_mtp_block_spec(local, vp_stage=0) + inner = mtp_block_spec.layer_specs[0].submodules.mtp_model_layer + full_last = loader.decoder_layer_specs()[-1] + assert inner.module is full_last.module + assert inner.submodules.self_attention.module is full_last.submodules.self_attention.module + + +def test_half_overridden_loader_fails_loud(): + """A loader that only overrides the sliced entry point must not get a silent template.""" + + class HalfLoader(ModelLoader): + + def get_transformer_layer_spec(self, vp_stage=None): + return TransformerBlockSubmodules(layer_specs=[]) + + loader = _loader(HalfLoader, _config()) + with _rank1_last_stage(): + with pytest.raises(NotImplementedError, match='_build_decoder_layer_specs'): + loader.get_mtp_block_spec(TransformerBlockSubmodules(layer_specs=[]), vp_stage=0)