From 8003b9c6f35972318b81ceb753b2d2726270373e Mon Sep 17 00:00:00 2001 From: cherry77-cloud <234872336+taking-lying-flat@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:31:17 +0800 Subject: [PATCH 1/2] fix MTP-only pipeline stage specs --- src/mcore_bridge/model/register.py | 47 ++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/src/mcore_bridge/model/register.py b/src/mcore_bridge/model/register.py index e0d6541..5272dac 100644 --- a/src/mcore_bridge/model/register.py +++ b/src/mcore_bridge/model/register.py @@ -9,6 +9,7 @@ from megatron.core.extensions.transformer_engine import TEGroupedLinear, TELayerNormColumnParallelLinear, TELinear 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.transformer.enums import LayerType 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.transformer_layer import TransformerLayer as McoreTransformerLayer @@ -135,7 +136,53 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): self._replace_spec_dsa(layer_spec) return transformer_layer_spec + def _get_mtp_transformer_layer_spec(self, transformer_layer_spec, vp_stage: Optional[int] = None): + layout = getattr(self.config, 'pipeline_model_parallel_layout', None) + if layout is None: + return transformer_layer_spec + + pp_rank = mpu.get_pipeline_model_parallel_rank() + vp_rank = vp_stage or 0 + stage = layout.layout[pp_rank][vp_rank] + if LayerType.mtp not in stage or LayerType.decoder in stage: + return transformer_layer_spec + + fallback_layout = deepcopy(layout) + stage_index = vp_rank * fallback_layout.pipeline_model_parallel_size + pp_rank + for previous_stage_index in range(stage_index - 1, -1, -1): + previous_pp_rank = previous_stage_index % fallback_layout.pipeline_model_parallel_size + previous_vp_rank = previous_stage_index // fallback_layout.pipeline_model_parallel_size + previous_stage = fallback_layout.layout[previous_pp_rank][previous_vp_rank] + decoder_count = previous_stage.count(LayerType.decoder) + if decoder_count: + break + else: + raise ValueError('Unable to build an MTP-only pipeline stage without a preceding decoder stage.') + + previous_stage[:] = [layer_type for layer_type in previous_stage if layer_type is not LayerType.decoder] + fallback_stage = fallback_layout.layout[pp_rank][vp_rank] + mtp_index = fallback_stage.index(LayerType.mtp) + fallback_stage[mtp_index:mtp_index] = [LayerType.decoder] * decoder_count + fallback_layout.flatten_layout = [ + layer_type for virtual_stages in zip(*fallback_layout.layout) for stage in virtual_stages + for layer_type in stage + ] + + self.config.pipeline_model_parallel_layout = fallback_layout + try: + transformer_layer_spec = self.get_transformer_layer_spec(vp_stage=vp_stage) + finally: + self.config.pipeline_model_parallel_layout = layout + + self._set_shared_expert_gate(transformer_layer_spec) + self._set_transformer_layer(transformer_layer_spec) + self._replace_mla_attention(transformer_layer_spec) + self._replace_router(transformer_layer_spec) + return transformer_layer_spec + def get_mtp_block_spec(self, transformer_layer_spec, vp_stage: Optional[int] = None): + transformer_layer_spec = self._get_mtp_transformer_layer_spec( + transformer_layer_spec, vp_stage=vp_stage) 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: From 7a271973a94e32270f7fa66744aef6c7d09219ac Mon Sep 17 00:00:00 2001 From: BANMA-00403 Date: Wed, 5 Aug 2026 13:27:58 +0800 Subject: [PATCH 2/2] align Qwen3.5 MTP layer spec --- src/mcore_bridge/model/gpts/qwen3_next_gdn.py | 27 +++++++++-- src/mcore_bridge/model/register.py | 47 ------------------- 2 files changed, 22 insertions(+), 52 deletions(-) diff --git a/src/mcore_bridge/model/gpts/qwen3_next_gdn.py b/src/mcore_bridge/model/gpts/qwen3_next_gdn.py index cb706a2..f659900 100644 --- a/src/mcore_bridge/model/gpts/qwen3_next_gdn.py +++ b/src/mcore_bridge/model/gpts/qwen3_next_gdn.py @@ -115,11 +115,8 @@ def _set_linear_decoupled_in_proj(self, mg_attn, hf_state_dict, to_mcore: bool): class Qwen3NextLoader(ModelLoader): gated_delta_net = GatedDeltaNet - 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: + def _replace_transformer_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 @@ -128,8 +125,28 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): if self.config.linear_decoupled_in_proj: layer_spec.submodules.input_layernorm = TENorm layer_spec.submodules.self_attention.submodules.in_proj = TEColumnParallelLinear + + 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) + self._replace_transformer_layer_specs(layer_specs.layer_specs) return layer_specs + def get_mtp_block_spec(self, transformer_layer_spec, vp_stage: Optional[int] = None): + from megatron.core.models.gpt.experimental_attention_variant_module_specs import \ + get_transformer_layer_with_experimental_attention_variant_spec + decoder_layer_specs = get_transformer_layer_with_experimental_attention_variant_spec(self.config) + self._replace_transformer_layer_specs(decoder_layer_specs) + + transformer_layer_spec = copy.deepcopy(transformer_layer_spec) + transformer_layer_spec.layer_specs = [decoder_layer_specs[-1]] + self._set_shared_expert_gate(transformer_layer_spec) + self._set_transformer_layer(transformer_layer_spec) + self._replace_mla_attention(transformer_layer_spec) + self._replace_router(transformer_layer_spec) + return super().get_mtp_block_spec(transformer_layer_spec, vp_stage=vp_stage) + def build_model( self, pre_process=True, diff --git a/src/mcore_bridge/model/register.py b/src/mcore_bridge/model/register.py index 5272dac..e0d6541 100644 --- a/src/mcore_bridge/model/register.py +++ b/src/mcore_bridge/model/register.py @@ -9,7 +9,6 @@ from megatron.core.extensions.transformer_engine import TEGroupedLinear, TELayerNormColumnParallelLinear, TELinear 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.transformer.enums import LayerType 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.transformer_layer import TransformerLayer as McoreTransformerLayer @@ -136,53 +135,7 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): self._replace_spec_dsa(layer_spec) return transformer_layer_spec - def _get_mtp_transformer_layer_spec(self, transformer_layer_spec, vp_stage: Optional[int] = None): - layout = getattr(self.config, 'pipeline_model_parallel_layout', None) - if layout is None: - return transformer_layer_spec - - pp_rank = mpu.get_pipeline_model_parallel_rank() - vp_rank = vp_stage or 0 - stage = layout.layout[pp_rank][vp_rank] - if LayerType.mtp not in stage or LayerType.decoder in stage: - return transformer_layer_spec - - fallback_layout = deepcopy(layout) - stage_index = vp_rank * fallback_layout.pipeline_model_parallel_size + pp_rank - for previous_stage_index in range(stage_index - 1, -1, -1): - previous_pp_rank = previous_stage_index % fallback_layout.pipeline_model_parallel_size - previous_vp_rank = previous_stage_index // fallback_layout.pipeline_model_parallel_size - previous_stage = fallback_layout.layout[previous_pp_rank][previous_vp_rank] - decoder_count = previous_stage.count(LayerType.decoder) - if decoder_count: - break - else: - raise ValueError('Unable to build an MTP-only pipeline stage without a preceding decoder stage.') - - previous_stage[:] = [layer_type for layer_type in previous_stage if layer_type is not LayerType.decoder] - fallback_stage = fallback_layout.layout[pp_rank][vp_rank] - mtp_index = fallback_stage.index(LayerType.mtp) - fallback_stage[mtp_index:mtp_index] = [LayerType.decoder] * decoder_count - fallback_layout.flatten_layout = [ - layer_type for virtual_stages in zip(*fallback_layout.layout) for stage in virtual_stages - for layer_type in stage - ] - - self.config.pipeline_model_parallel_layout = fallback_layout - try: - transformer_layer_spec = self.get_transformer_layer_spec(vp_stage=vp_stage) - finally: - self.config.pipeline_model_parallel_layout = layout - - self._set_shared_expert_gate(transformer_layer_spec) - self._set_transformer_layer(transformer_layer_spec) - self._replace_mla_attention(transformer_layer_spec) - self._replace_router(transformer_layer_spec) - return transformer_layer_spec - def get_mtp_block_spec(self, transformer_layer_spec, vp_stage: Optional[int] = None): - transformer_layer_spec = self._get_mtp_transformer_layer_spec( - transformer_layer_spec, vp_stage=vp_stage) 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: