From d203f8734746be052a4baddbe58ecf58e47b2a74 Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Mon, 21 Sep 2026 11:40:40 +0800 Subject: [PATCH 1/3] update --- src/mcore_bridge/model/gpt_model.py | 15 + src/mcore_bridge/model/gpts/qwen4_exp.py | 206 ++++++++++- src/mcore_bridge/model/mm_gpts/glm5_next.py | 219 +++++++++--- src/mcore_bridge/model/modules/mtp_layer.py | 11 + .../model/modules/transformer_block.py | 6 + .../patches/megatron_glm53_dev.patch | 18 +- src/mcore_bridge/utils/megatron_utils.py | 10 + tests/test_glm5_next.py | 140 +++++++- tests/test_qwen4_exp_mtp.py | 330 ++++++++++++++++++ 9 files changed, 893 insertions(+), 62 deletions(-) create mode 100644 tests/test_qwen4_exp_mtp.py diff --git a/src/mcore_bridge/model/gpt_model.py b/src/mcore_bridge/model/gpt_model.py index 52946102..a3a1968a 100644 --- a/src/mcore_bridge/model/gpt_model.py +++ b/src/mcore_bridge/model/gpt_model.py @@ -177,6 +177,11 @@ def _preprocess( if self.config.is_multimodal and self.config.mtp_num_layers and decoder_input is None: input_tensor = self.get_input_tensor() input_tensor, mtp_decoder_input = input_tensor.chunk(2, dim=0) + # Pipeline communication carries a single tensor, so models whose backbone keeps + # multiple hidden streams pad the H-wide embedding to the n*H transport width. + # The padding is transport-only and must not enter the MTP projection. + if mtp_decoder_input.shape[-1] != self.config.hidden_size: + mtp_decoder_input = mtp_decoder_input[..., :self.config.hidden_size].contiguous() self.set_input_tensor(input_tensor) rotary_pos_emb, rotary_pos_cos, rotary_pos_sin = self._get_rotary_pos_emb( @@ -448,6 +453,16 @@ def _postprocess(self, """ if not self.post_process: if self.config.is_multimodal and self.config.mtp_num_layers: + # Pipeline P2P sends one tensor. Hyper-connection backbones can expose n*H + # hidden states while the embedding retained for MTP is only H wide, so pad + # the embedding for transport and slice it back in _preprocess on the next stage. + hidden_width = hidden_states.shape[-1] + decoder_width = decoder_input.shape[-1] + assert decoder_width <= hidden_width, ( + f'MTP pipeline transport requires decoder width <= hidden width, got ' + f'{decoder_width} and {hidden_width}.') + if decoder_width != hidden_width: + decoder_input = F.pad(decoder_input, (0, hidden_width - decoder_width)) return torch.concat([hidden_states, decoder_input], dim=0) else: return hidden_states diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index aa0c4dce..470332d0 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/src/mcore_bridge/model/gpts/qwen4_exp.py @@ -1,29 +1,34 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import copy import math +import megatron.core import torch import torch.distributed as dist import torch.nn.functional as F 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 +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec, get_gpt_mtp_block_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 +from megatron.core.tensor_parallel.mappings import (gather_from_tensor_model_parallel_region, + scatter_to_sequence_parallel_region) 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.megatron_utils import reconstruct_tensor_cp -from ..modules import (QSA_SPARSE_KERNEL_ENV, GatedDeltaNet, QSAIndexer, QSASparseCoreAttention, - Qwen4ExpTextGatedResidual, Qwen4ExpTextPLELayer, TransformerBlock, TransformerLayer, - qsa_sparse_supported, use_qsa_sparse_kernel) +from ..modules import (QSA_SPARSE_KERNEL_ENV, GatedDeltaNet, MultiTokenPredictionLayer, QSAIndexer, + QSASparseCoreAttention, Qwen4ExpTextGatedResidual, Qwen4ExpTextPLELayer, TransformerBlock, + TransformerLayer, qsa_sparse_supported, use_qsa_sparse_kernel) from ..modules.ple import Qwen4ExpTextNGramEmbedding from ..register import ModelLoader from .qwen3_next import Qwen3NextBridge, Qwen3NextRMSNorm, Qwen3NextSelfAttention @@ -61,7 +66,7 @@ def __init__(self, config, submodules, layer_number: int = 1, **kwargs): if self.layer_number in config.ple_layer_ids: self.ple = Qwen4ExpTextPLELayer( config, config.ple_layer_ids.index(self.layer_number), pg_collection=self.pg_collection) - is_linear_attention = config.linear_attention_freq[self.layer_number - 1] + is_linear_attention = self._resolve_is_linear_attention(config) if not is_linear_attention and config.indexer_n_heads is not None: self.self_attention.indexer = QSAIndexer(config, tp_group=self.tp_group) if qsa_sparse_supported(config.kv_channels): @@ -71,6 +76,10 @@ def __init__(self, config, submodules, layer_number: int = 1, **kwargs): self.attn_hyper_connection = Qwen4ExpTextGatedResidual(config) self.mlp_hyper_connection = Qwen4ExpTextGatedResidual(config) + # override in MTP layer + def _resolve_is_linear_attention(self, config): + return config.linear_attention_freq[self.layer_number - 1] + def forward(self, hidden_states: torch.Tensor, **kwargs): attention_mask = kwargs.get('attention_mask') packed_seq_params: PackedSeqParams = kwargs.get('packed_seq_params') @@ -284,6 +293,105 @@ def __init__(self, *args, **kwargs): self.hyper_connection_mixer = Qwen4ExpTextGatedResidual(config, use_combine=False) +class Qwen4ExpMTPInnerLayer(Qwen4ExpLayer): + + def _resolve_is_linear_attention(self, config): + return False + + +class Qwen4ExpMTPStreamNorm(nn.Module): + """Zero-centered RMSNorm over the full multi-stream (``hc_count * hidden_size``). + + Qwen3.8-Flash-Next's MTP ``pre_fc_norm_hidden`` normalizes the concatenated ``hc_count`` streams + jointly (a single GemmaRMSNorm over n*H with a per-element affine), unlike Megatron's mHC MTP which + normalizes each H-sized stream independently. The MTP spec builds norms with ``hidden_size=H``, so + scale by ``hc_count`` here. + """ + + def __init__(self, config, hidden_size, eps): + super().__init__() + self.dim = config.hc_count * hidden_size + self.eps = eps + self.weight = nn.Parameter(torch.zeros(self.dim, dtype=config.params_dtype)) + self.weight.sequence_parallel = config.sequence_parallel + + def forward(self, x): + input_dtype = x.dtype + x = x.float() + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + return ((1.0 + self.weight.float()) * x).to(input_dtype) + + +class Qwen4ExpMultiTokenPredictionLayer(MultiTokenPredictionLayer): + """Qwen3.8-Flash-Next MTP head: the ``residual_linear_shared`` fusion over the gated-HC backbone. + + The backbone runs Qwen4Exp's own gated hyper-connections (``hc_count`` streams) but NOT Megatron's + ``enable_hyper_connections`` mHC. The MTP head still needs Megatron's mHC *projection* form + (separate ``e_proj``/``h_proj`` rather than a fused ``eh_proj``), because the checkpoint stores + ``fc_embedding``/``fc_hidden`` (H->H each) and adds the embedding residual to every stream. So the + subtree is built with hyper-connections enabled via a config copy -- which selects e_proj/h_proj -- + then Megatron's mHC-only contraction params (``hc_head_*``) and ``final_layernorm`` are dropped in + favour of Qwen4Exp's own ``hyper_connection_mixer``, matching the checkpoint and the vLLM draft + model. ``_concat_embeddings`` normalizes the multi-stream jointly (``Qwen4ExpMTPStreamNorm``) and + ``_postprocess`` contracts with the mixer, so the layer returns the multi-stream and the MTP block + applies ``_postprocess`` for the loss head. + """ + + def __init__(self, config, submodules, *args, **kwargs): + mtp_config = copy.copy(config) + mtp_config.enable_hyper_connections = True + super().__init__(mtp_config, submodules, *args, **kwargs) + for name in ('hc_head_fn', 'hc_head_base', 'hc_head_scale', 'final_layernorm'): + if hasattr(self, name): + delattr(self, name) + self.hyper_connection_mixer = Qwen4ExpTextGatedResidual(config, use_combine=False) + + def _concat_embeddings(self, hidden_states, decoder_input): + # hidden_states: pre-mixer multi-stream [s, b, n*H]; decoder_input: rolled-token embedding [s, b, H]. + n = self.config.hc_count + h = self.config.hidden_size + decoder_input = self.enorm(decoder_input) + decoder_input = make_viewless_tensor(inp=decoder_input, requires_grad=True, keep_graph=True) + # Qwen4Exp normalizes the full multi-stream jointly (pre_fc_norm_hidden over n*H), not per-stream. + hs = self.hnorm(hidden_states) + hs = make_viewless_tensor(inp=hs, requires_grad=True, keep_graph=True).unflatten(-1, (n, h)) + e_out, _ = self.e_proj(decoder_input) # fc_embedding -> [s, b, H/tp] + h_out, _ = self.h_proj(hs) # fc_hidden -> [s, b, n, H/tp] + out = e_out.unsqueeze(2) + h_out # add the embedding residual to every stream + out = gather_from_tensor_model_parallel_region(out, group=self.tp_group) + # Read the shape AFTER the gather: under sequence parallel the column-parallel projections + # all-gather the sequence dim, so a pre-projection `s` would be stale. + s, b, n_out, h_dim = out.shape + out = out.reshape(s, b, n_out * h_dim) + if self.sequence_parallel: + out = scatter_to_sequence_parallel_region(out, group=self.tp_group) + return out + + def _postprocess(self, hidden_states): + # Contract the multi-stream [s, b, n*H] to [s, b, H] with Qwen4Exp's gated mixer (no final norm). + return self.hyper_connection_mixer(hidden_states) + + def _get_embeddings(self, + input_ids, + position_ids, + embedding, + hidden_states, + packed_seq_params=None, + decoder_input=None): + input_ids, position_ids, decoder_input, hidden_states = super()._get_embeddings( + input_ids, position_ids, embedding, hidden_states, packed_seq_params, decoder_input) + # Stash the rolled ids so _proj_and_transformer_layer can forward them to the inner + # Qwen4ExpMTPInnerLayer (its QSA indexer needs position_ids under packing/CP; PLE is absent). + self._mtp_input_ids = input_ids + self._mtp_position_ids = position_ids + return input_ids, position_ids, decoder_input, hidden_states + + def _proj_and_transformer_layer(self, *args, **kwargs): + kwargs.setdefault('input_ids', getattr(self, '_mtp_input_ids', None)) + kwargs.setdefault('position_ids', getattr(self, '_mtp_position_ids', None)) + return super()._proj_and_transformer_layer(*args, **kwargs) + + class Qwen4ExpBridge(Qwen3NextBridge): hf_mixer_prefix = 'model.' @@ -434,7 +542,13 @@ def _set_layer_ple(self, mg_layer, hf_state_dict, to_mcore: bool, layer_prefix: # loop runs pp collectives (broadcast_object_list) and export_table_to_hf runs # tp ones, and stages disagreeing on whether to enter would deadlock. ple_offloaded = self._reduce_tensor_pp_group(ple is not None and ple.ple_embedding.cpu_offload, to_mcore) - skip_ngram_state = not to_mcore and not self._is_saving and (self._peft_format or ple_offloaded) + if to_mcore: + # A PEFT/adapter checkpoint carries no PLE n-gram buffers -- those come from the base + # checkpoint that the adapter is applied on top of -- so a peft-format load must skip + # them instead of KeyError-ing on `ple.ple_embedding.layer_multipliers`. + skip_ngram_state = self._peft_format + else: + skip_ngram_state = not self._is_saving and (self._peft_format or ple_offloaded) for buf in () if skip_ngram_state else self._PLE_NGRAM_BUFFERS: if to_mcore: buffer = getattr(ple.ple_embedding, buf) @@ -512,6 +626,54 @@ def _convert_post_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcor f'{self.hf_mixer_prefix}hyper_connection_mixer.{key}', to_mcore) return res + def _convert_mtp_extra(self, mtp_layer, hf_state_dict, to_mcore, origin_hf_state_dict): + # Qwen3.8-Flash-Next's MTP head lives at the `mtp.` level (not under `mtp.layers.i`): + # pre_fc_norm_embedding/pre_fc_norm_hidden -> enorm/hnorm, fc_embedding/fc_hidden -> e_proj/h_proj + # (the residual_linear_shared fusion), plus its own hyper_connection_mixer for the contraction. + # There is no fused eh_proj and no final norm (the mixer is the contraction). + sd = self._remove_prefix(origin_hf_state_dict, 'mtp.') + for mg_key, key in [('enorm.weight', 'pre_fc_norm_embedding.weight'), + ('hnorm.weight', 'pre_fc_norm_hidden.weight'), ('e_proj.weight', 'fc_embedding.weight'), + ('h_proj.weight', 'fc_hidden.weight')]: + self._set_state_dict(mtp_layer, mg_key, sd, key, to_mcore) + self._fp8_skip_modules.update({'mtp.fc_embedding', 'mtp.fc_hidden'}) + mixer = None if mtp_layer is None else getattr(mtp_layer, 'hyper_connection_mixer', None) + for key in ('hc_norm.weight', 'input_mix_weight_down.weight', 'input_mix_weight_up.weight'): + self._set_state_dict(mixer, key, sd, f'hyper_connection_mixer.{key}', to_mcore) + if not to_mcore: + origin_hf_state_dict.update(self._add_prefix(sd, 'mtp.')) + + def _convert_mtp_layer(self, lm_model, hf_state_dict, hf_prefix: str, layer_idx: int, to_mcore: bool): + mtp_layer = lm_model.mtp.layers[layer_idx] if hasattr(lm_model, 'mtp') else None + hf_prefix = f'{hf_prefix}{layer_idx}.' # 'mtp.layers.0.' + if to_mcore: + origin_hf_state_dict = hf_state_dict + hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) + if len(hf_state_dict) == 0: + logger.info(f'MTP layer {layer_idx} safetensors weights not found, ' + 'this part will be randomly initialized.') + for param in mtp_layer.parameters(): + if param.ndim == 2: + mtp_layer.config.init_method(param.data) + return {} + else: + origin_hf_state_dict = {} + hf_state_dict = {} + self._convert_mtp_extra(mtp_layer, hf_state_dict, to_mcore, origin_hf_state_dict) + # Inner block: a full-attention + MoE Qwen4ExpLayer with its own gated hyper-connections. + # layer_idx=-1 routes _set_layer_attn through linear_attention_freq[-1] (the backbone's last + # layer, full_attention), matching the MTP head, which is always full-attention. + inner = None if mtp_layer is None else mtp_layer.transformer_layer + hf_state_dict.update(self._set_layer_attn(inner, hf_state_dict, -1, to_mcore)) + hf_state_dict.update(self._set_layer_mlp(inner, hf_state_dict, -1, to_mcore, is_mtp=True)) + self._set_layer_hc(inner, hf_state_dict, to_mcore) + if to_mcore: + hf_state_dict = {} + else: + hf_state_dict = self._add_prefix(hf_state_dict, hf_prefix) + hf_state_dict.update(origin_hf_state_dict) + return hf_state_dict + class Qwen4ExpLoader(ModelLoader): transformer_block = Qwen4ExpTransformerBlock @@ -532,8 +694,6 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): "Qwen4-Exp QSA under context parallelism requires cp_comm_type='all_gather'; " f"got {getattr(config, 'cp_comm_type', None)!r} (mcore's default), promoting to 'all_gather'.") config.cp_comm_type = 'all_gather' - if getattr(config, 'mtp_num_layers', None): - raise NotImplementedError('Qwen4-Exp MTP is not supported yet') moe_spec = get_gpt_layer_with_transformer_engine_spec( num_experts=config.num_moe_experts, moe_grouped_gemm=config.moe_grouped_gemm, @@ -588,6 +748,36 @@ 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): + 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 + # The residual_linear_shared head needs Megatron's mHC *projection* form (separate + # e_proj/h_proj slots). megatron-core <= 0.18 only has the fused eh_proj slot, and + # assigning e_proj/h_proj there would silently no-op and surface later as an + # AttributeError on self.e_proj -- so reject early and name the fix. + if not (hasattr(sub, 'e_proj') and hasattr(sub, 'h_proj')): + raise NotImplementedError( + 'Qwen3.8-Flash-Next MTP requires a Megatron whose MultiTokenPredictionLayerSubmodules ' + 'exposes e_proj/h_proj (megatron-core >= 0.19 / dev); got ' + f'{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 + # (final_layernorm) is built then dropped by the layer -- the mixer is the contraction. + sub.enorm = TENorm + sub.hnorm = Qwen4ExpMTPStreamNorm + sub.eh_proj = None + sub.e_proj = TEColumnParallelLinear + sub.h_proj = TEColumnParallelLinear + sub.layer_norm = TENorm + # The MTP inner block is always full-attention (config.mtp.layer_types), independent of + # the backbone layer numbering that Qwen4ExpLayer would otherwise read. + sub.mtp_model_layer.module = Qwen4ExpMTPInnerLayer + return mtp_block_spec + def build_model( self, pre_process=True, diff --git a/src/mcore_bridge/model/mm_gpts/glm5_next.py b/src/mcore_bridge/model/mm_gpts/glm5_next.py index deced8b8..ee3dbbd3 100644 --- a/src/mcore_bridge/model/mm_gpts/glm5_next.py +++ b/src/mcore_bridge/model/mm_gpts/glm5_next.py @@ -1,4 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import copy import torch import torch.nn as nn from copy import deepcopy @@ -15,8 +16,11 @@ from .utils import HuggingFaceVit try: + from megatron.core.models.hybrid.hybrid_block import HybridStack from megatron.core.models.hybrid.hybrid_layer_allocation import parse_hybrid_pattern, select_pipeline_segment from megatron.core.transformer.module import mark_keep_in_fp32 + from megatron.core.transformer.multi_token_prediction import \ + MultiTokenPredictionLayer as McoreMultiTokenPredictionLayer from ..hybrid_model import HybridModel except ImportError: @@ -24,7 +28,7 @@ # models.hybrid). The package must still import for every other model, so fall back to a base # that lets the class definitions below succeed; a GLM config is rejected before any of this is # used, by require_glm5_hybrid(), which names the missing dev patch. - HybridModel = object + HybridModel = HybridStack = McoreMultiTokenPredictionLayer = object parse_hybrid_pattern = select_pipeline_segment = mark_keep_in_fp32 = None logger = get_logger() @@ -107,6 +111,42 @@ def __init__(self, config, hidden_size, eps): self.weight.sequence_parallel = config.sequence_parallel +class Glm5NextHybridStack(HybridStack): + """Outer decoder stack that never hands the pre-contraction multi-stream tensor to MTP. + + GLM-5.3's backbone runs mHC (4 residual streams) but its MTP head is non-mHC: it consumes the + mean-contracted single stream through a fused ``eh_proj`` (2H -> H), exactly like the HF/vLLM + draft model. Upstream ``HybridStack.forward`` returns ``(hidden_states, mhc_multistream)`` when + mHC + MTP are both active, and ``MultiTokenPredictionBlock.forward`` then feeds the multi-stream + to the MTP layer and re-applies ``_postprocess`` to its output -- which double-norms a non-mHC + MTP layer whose output is already single-stream and post-final-norm. Dropping the second element + keeps the block on its single-stream path, where the MTP layer output is used as-is. + """ + + def forward(self, *args, **kwargs): + out = super().forward(*args, **kwargs) + return out[0] if isinstance(out, tuple) else out + + +class Glm5NextMultiTokenPredictionLayer(McoreMultiTokenPredictionLayer): + """GLM-5.3 MTP layer: a non-mHC ``eh_proj`` head over a plain DSA + MoE inner block. + + The backbone sets ``enable_hyper_connections=True``, which would make Megatron build the mHC MTP + head (per-stream ``e_proj``/``h_proj`` plus an ``hc_head_*`` learned contraction) and wrap the + inner block's sublayers in hyper-connections. GLM's MTP checkpoint has none of those: it carries + a fused ``eh_proj``, plain ``enorm``/``hnorm``, a ``shared_head.norm`` final norm, and an inner + block with standard residuals. So the MTP subtree is built with hyper-connections disabled via a + shallow config copy that flips only that flag; the backbone keeps the real config. ``mhc_enabled`` + (read once in ``__init__``) then follows the copy, selecting the ``eh_proj`` branch, and the + nested ``HybridStack`` skips hyper-connection wrapping for its DSA/MoE sublayers. + """ + + def __init__(self, config, *args, **kwargs): + mtp_config = copy.copy(config) + mtp_config.enable_hyper_connections = False + super().__init__(mtp_config, *args, **kwargs) + + class Glm5NextHybridModel(HybridModel): extra_forward_keys = () @@ -120,20 +160,29 @@ def _resolve_hybrid_layer_pattern(config): 'num_layers_in_first_pipeline_stage / num_layers_in_last_pipeline_stage.') if ('|' in pattern or config.num_layers_in_first_pipeline_stage is not None or config.num_layers_in_last_pipeline_stage is not None): - return pattern - stages = config.pipeline_model_parallel_size - if config.virtual_pipeline_model_parallel_size is not None: - stages *= config.virtual_pipeline_model_parallel_size - blocks, extra = divmod(len(mapping) // 2, stages) - if blocks == 0: - raise ValueError('The current model needs at least one attention+FFN block per pipeline stage; ' - 'lower pipeline_model_parallel_size.') - segments, offset = [], 0 - for stage in range(stages): - count = 2 * (blocks + int(stage < extra)) - segments.append(pattern[offset:offset + count]) - offset += count - return '|'.join(segments) + main = pattern + else: + stages = config.pipeline_model_parallel_size + if config.virtual_pipeline_model_parallel_size is not None: + stages *= config.virtual_pipeline_model_parallel_size + blocks, extra = divmod(len(mapping) // 2, stages) + if blocks == 0: + raise ValueError('The current model needs at least one attention+FFN block per pipeline stage; ' + 'lower pipeline_model_parallel_size.') + segments, offset = [], 0 + for stage in range(stages): + count = 2 * (blocks + int(stage < extra)) + segments.append(pattern[offset:offset + count]) + offset += count + main = '|'.join(segments) + # Append the MTP segment (`
///...`, one per prediction depth). Megatron's + # parse_hybrid_pattern splits on '/' and keeps '|' inside the main pattern for PP boundaries; + # the MTP inner block (DSA + MoE, no hyper-connections) is described by mtp_hybrid_override_pattern. + mtp_pattern = getattr(config, 'mtp_hybrid_override_pattern', None) + mtp_num_layers = getattr(config, 'mtp_num_layers', None) + if mtp_pattern and mtp_num_layers and '/' not in main: + main = main + '/' + '/'.join([mtp_pattern] * mtp_num_layers) + return main def __init__(self, config, transformer_layer_spec, pre_process=True, post_process=True, vp_stage=None): # The PP boundaries have to be validated before super().__init__: an attention/FFN pair must @@ -166,24 +215,36 @@ def __init__(self, config, transformer_layer_spec, pre_process=True, post_proces for layer in self.decoder.layers: layer.hyper_connection.sinkhorn_eps = config.hc_eps layer.hyper_connection.compute_h_eps = config.hc_eps - if mapping[layer.layer_number - 1][2] == 'D': - # HF runs the KPool indexer under no_grad, so it must not be counted in DDP's - # grad-ready accounting. - layer.inner_layer.self_attention.core_attention.indexer.requires_grad_(False) - elif mapping[layer.layer_number - 1][2] == 'E': - # Experts reduce over expert-DP, and ETP can differ from attention TP even at EP=1. - # dev's TEGroupedLinear derives `allreduce` from EP alone and does not mark the ETP - # shard, which would miss expert grad reduction and norm sharding. Corrected here at - # the GLM boundary so other models keep their defaults. - for param in layer.inner_layer.mlp.experts.parameters(): - param.allreduce = False - param.tensor_model_parallel = config.expert_tensor_parallel_size > 1 - router = layer.inner_layer.mlp.router - if router.enable_expert_bias: - # BF16 cannot accumulate large integer counts exactly, and PP changes the number - # of micro-batches, which would change the router update. - mark_keep_in_fp32(router.local_tokens_per_expert) - mark_keep_in_fp32(router.expert_bias) + self._finalize_hybrid_sublayer(layer.inner_layer, mapping[layer.layer_number - 1][2], config) + # The MTP head's inner block is a nested HybridStack of raw (non-hyper-connected) sublayers + # following mtp_hybrid_override_pattern ('DE'); it needs the same DSA/MoE finalization. + mtp = getattr(self, 'mtp', None) + if mtp is not None: + mtp_pattern = getattr(config, 'mtp_hybrid_override_pattern', '') or '' + for mtp_layer in mtp.layers: + for symbol, sublayer in zip(mtp_pattern, mtp_layer.mtp_model_layer.layers): + self._finalize_hybrid_sublayer(sublayer, symbol, config) + + @staticmethod + def _finalize_hybrid_sublayer(sublayer, symbol, config): + if symbol == 'D': + # HF runs the KPool indexer under no_grad, so it must not be counted in DDP's + # grad-ready accounting. + sublayer.self_attention.core_attention.indexer.requires_grad_(False) + elif symbol == 'E': + # Experts reduce over expert-DP, and ETP can differ from attention TP even at EP=1. + # dev's TEGroupedLinear derives `allreduce` from EP alone and does not mark the ETP + # shard, which would miss expert grad reduction and norm sharding. Corrected here at + # the GLM boundary so other models keep their defaults. + for param in sublayer.mlp.experts.parameters(): + param.allreduce = False + param.tensor_model_parallel = config.expert_tensor_parallel_size > 1 + router = sublayer.mlp.router + if router.enable_expert_bias: + # BF16 cannot accumulate large integer counts exactly, and PP changes the number + # of micro-batches, which would change the router update. + mark_keep_in_fp32(router.local_tokens_per_expert) + mark_keep_in_fp32(router.expert_bias) def _get_packed_padding_mask(self, packed_seq_params, position_ids): # `seq_lens` is the logical length attached by swift's prepare_batch; it is not one of the @@ -412,26 +473,81 @@ def _set_layer_state(self, mg_layer, hf_state_dict, hf_prefix: str, layer_idx: i return {} if to_mcore else self._add_prefix(hf_state_dict, layer_prefix) def _filter_mtp_layer(self, hf_state_dict): - """Drop the extra MTP-only decoder layer (index == num_layers) from a HF checkpoint.""" - # TODO: MTP is not supported yet -- the loader rejects mtp_num_layers and the pattern resolver - # emits no MTP segment -- so the checkpoint's extra MTP layer has nowhere to go. Wire it up - # through mtp_hybrid_override_pattern / mtp_on_this_rank once that path is validated here. + """Drop the extra MTP-only decoder layer (index == num_hidden_layers) from a HF checkpoint. + + Only used when MTP training is OFF: the checkpoint always ships the MTP head, but a run with + ``mtp_num_layers=0`` builds no MTP module, so those tensors would otherwise be reported as + unexpected. When MTP is ON they are consumed by ``_convert_mtp_layer`` instead. + """ hf_num_layers = self.config.num_layers // 2 layer_prefix = f'{self.hf_layers_prefix}.{hf_num_layers}.' prefixes = (layer_prefix, layer_prefix.removeprefix('model.'), f'layers.{hf_num_layers}.') ignored = [key for key in hf_state_dict if key.startswith(prefixes)] if ignored: logger.warning_once( - f'Ignoring {len(ignored)} MTP tensors under decoder layer {hf_num_layers}: the current model ' - 'builds exactly num_hidden_layers decoder layers and does not train MTP yet.') + f'Ignoring {len(ignored)} MTP tensors under decoder layer {hf_num_layers}: this run builds no ' + 'MTP module (mtp_num_layers=0). Pass --mtp_num_layers 1 to train the MTP head.') hf_state_dict = {key: value for key, value in hf_state_dict.items() if not key.startswith(prefixes)} return hf_state_dict def _convert_hf_state_dict(self, hf_state_dict, to_mcore): - if to_mcore: + if to_mcore and not getattr(self.config, 'mtp_num_layers', None): hf_state_dict = self._filter_mtp_layer(hf_state_dict) return super()._convert_hf_state_dict(hf_state_dict, to_mcore) + def _convert_mtp_layer(self, lm_model, hf_state_dict, hf_prefix: str, layer_idx: int, to_mcore: bool): + """Map GLM's MTP head, which lives at decoder-layer index ``num_hidden_layers``. + + ``num_layers`` counts hybrid *sublayers* (two per HF block), so the MTP head's HF index is + ``num_layers // 2 + layer_idx`` and it shares the backbone's ``model.language_model.layers`` + prefix -- neither the separate-prefix nor the ``+num_layers`` assumptions of the base + ``_convert_mtp_layer`` hold. The head is non-mHC: ``enorm``/``hnorm``/fused ``eh_proj`` plus a + ``shared_head.norm`` final norm, over a nested HybridStack whose raw sublayers are a DSA + attention block (``input_layernorm`` + ``self_attn``) and a MoE FFN block + (``pre_mlp_layernorm`` + ``mlp``), with no hyper-connection weights. + """ + mtp_layer = lm_model.mtp.layers[layer_idx] if hasattr(lm_model, 'mtp') else None + hf_layer_idx = self.config.num_layers // 2 + layer_idx + layer_prefix = f'{self.hf_layers_prefix}.{hf_layer_idx}.' + if to_mcore: + origin_hf_state_dict = hf_state_dict + hf_state_dict = self._remove_prefix(hf_state_dict, layer_prefix) + if len(hf_state_dict) == 0: + logger.info(f'MTP layer {hf_layer_idx} safetensors weights not found, ' + 'this part will be randomly initialized.') + for param in mtp_layer.parameters(): + if param.ndim == 2: + mtp_layer.config.init_method(param.data) + return {} + else: + origin_hf_state_dict = {} + hf_state_dict = {} + # MTP head: enorm / hnorm / fused eh_proj (2H -> H) / final_layernorm (shared_head.norm). + for mg_key, hf_key in [('enorm.weight', 'enorm.weight'), ('hnorm.weight', 'hnorm.weight'), + ('eh_proj.weight', 'eh_proj.weight'), + ('final_layernorm.weight', 'shared_head.norm.weight')]: + self._set_state_dict(mtp_layer, mg_key, hf_state_dict, hf_key, to_mcore) + # eh_proj is kept in BF16 in the FP8 checkpoint (see modules_to_not_convert); never quantize it. + self._fp8_skip_modules.update({'eh_proj'}) + # Inner block: nested HybridStack, sublayer 0 = DSA attention, sublayer 1 = MoE FFN. + inner = None if mtp_layer is None else mtp_layer.mtp_model_layer + attn_sub = None if inner is None else inner.layers[0] + ffn_sub = None if inner is None else inner.layers[1] + attn = None if attn_sub is None else attn_sub.self_attention + hf_state_dict.update(self._set_dsa_state(attn, hf_state_dict, to_mcore)) + self._set_state_dict(attn_sub, 'input_layernorm.weight', hf_state_dict, self.hf_input_layernorm_key, to_mcore) + mlp = None if ffn_sub is None else ffn_sub.mlp + hf_state_dict.update( + self._set_moe_state(mlp, hf_state_dict, f'{self.hf_mlp_prefix}.', hf_layer_idx, to_mcore, is_mtp=True)) + self._set_state_dict(ffn_sub, 'pre_mlp_layernorm.weight', hf_state_dict, self.hf_post_attention_layernorm_key, + to_mcore) + if to_mcore: + hf_state_dict = {} + else: + hf_state_dict = self._add_prefix(hf_state_dict, layer_prefix) + hf_state_dict.update(origin_hf_state_dict) + return hf_state_dict + class Glm5NextLoader(ModelLoader): @@ -442,8 +558,6 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): from ..modules import TopKRouter config = self.config - if config.mtp_num_layers: - raise NotImplementedError('The current model builds no MTP layers; use mtp_num_layers=0') if config.fp8 or config.fp4: raise NotImplementedError('The current model is validated for BF16/FP32 only; ' 'fp8/fp4 training is not supported') @@ -461,6 +575,14 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): config.mscale_all_dim = MLATransformerConfig.mscale_all_dim config.cache_mla_latents = MLATransformerConfig.cache_mla_latents config.enable_hyper_connections = True + # GLM-5.3-Flash's MTP head (decoder layer index num_hidden_layers) is a single DSA-attention + + # MoE-FFN block with plain residuals -- it carries no hyper-connection weights, unlike every + # backbone block. This describes that inner block for Megatron's hybrid MTP path and is only + # consumed when mtp_num_layers > 0. Set here rather than in the parser: it is a Megatron + # TransformerConfig field (not a ModelConfig one), so emitting it from the parser would make + # ModelConfig(**values) raise a bare TypeError on a Megatron that lacks it -- before + # __post_init__ could run require_glm5_hybrid() and name the actual fix. + config.mtp_hybrid_override_pattern = 'DE' config.mhc_norm_eps_inside_sqrt = config.mhc_keep_mappings_in_fp32 = True config.mhc_learned_output_contract = False config.kda_two_stage_gates = True @@ -492,6 +614,19 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): moe = spec.submodules.moe_layer.submodules moe.pre_mlp_layernorm = Glm5NextHybridRMSNorm moe.mlp.keywords['submodules'].router = TopKRouter + # Outer decoder stack drops the multi-stream tensor so MTP stays on its non-mHC path. + spec.module = Glm5NextHybridStack + # MTP head (built only when mtp_num_layers > 0): non-mHC eh_proj over a plain DSA + MoE inner + # block. Glm5NextMultiTokenPredictionLayer disables hyper-connections for the MTP subtree, so + # only eh_proj/enorm/hnorm/layer_norm are used; the GLM RMSNorm matches the checkpoint's plain + # (non zero-centered) enorm/hnorm/shared_head.norm. + mtp_block_spec = getattr(spec.submodules, 'mtp_block_spec', None) + if mtp_block_spec is not None: + mtp_layer_spec = mtp_block_spec.submodules.layer_specs[0] + mtp_layer_spec.module = Glm5NextMultiTokenPredictionLayer + mtp_sub = mtp_layer_spec.submodules + mtp_sub.enorm = mtp_sub.hnorm = mtp_sub.layer_norm = Glm5NextHybridRMSNorm + mtp_sub.eh_proj = TEColumnParallelLinear return spec def build_model(self, pre_process=True, post_process=True, vp_stage: Optional[int] = None): diff --git a/src/mcore_bridge/model/modules/mtp_layer.py b/src/mcore_bridge/model/modules/mtp_layer.py index 162e06d0..b6d8b306 100644 --- a/src/mcore_bridge/model/modules/mtp_layer.py +++ b/src/mcore_bridge/model/modules/mtp_layer.py @@ -60,6 +60,17 @@ def __init__(self, config: ModelConfig, submodules, *args, **kwargs): tp_group=self.tp_group, ) + @property + def transformer_layer(self): + """The MTP inner transformer block, under whichever name the running Megatron registers it. + + Megatron dev renamed the attribute to ``mtp_model_layer`` (and keys checkpoints as + ``transformer_layer`` for backward compat), while main still calls the module itself + ``transformer_layer``. mcore-bridge's forward below and ``GPTBridge._convert_mtp_layer`` + reference ``.transformer_layer``, so resolve whichever exists. + """ + return self._modules.get('mtp_model_layer', None) or self._modules.get('transformer_layer', None) + def forward( self, input_ids: torch.Tensor, diff --git a/src/mcore_bridge/model/modules/transformer_block.py b/src/mcore_bridge/model/modules/transformer_block.py index 05a333f0..b185132c 100644 --- a/src/mcore_bridge/model/modules/transformer_block.py +++ b/src/mcore_bridge/model/modules/transformer_block.py @@ -497,6 +497,12 @@ def forward( elif enable_gated_hc and self.has_final_layernorm_in_this_stage(): # Gated low-rank contraction (hyper_connection_mixer, use_combine=False # so forward returns only the mixed stream). + # When MTP is enabled, save the pre-contraction multi-stream [s, b, n*C] for the MTP + # head: Qwen4Exp's MTP fuses it with the rolled-token embedding (residual_linear_shared) + # and runs its own hyper_connection_mixer, so it needs the multi-stream, not the + # contracted [s, b, C] that the lm_head consumes. + if self.config.mtp_num_layers: + mhc_multistream = hidden_states # [s, b, n*C] -> [s, b, C] hidden_states = self.hyper_connection_mixer(hidden_states) diff --git a/src/mcore_bridge/patches/megatron_glm53_dev.patch b/src/mcore_bridge/patches/megatron_glm53_dev.patch index 0122d874..af40e7c2 100644 --- a/src/mcore_bridge/patches/megatron_glm53_dev.patch +++ b/src/mcore_bridge/patches/megatron_glm53_dev.patch @@ -830,18 +830,30 @@ index 9193f4e00a..70174144ba 100644 assert q is not None and k is not None and weights is not None fused_output = dsa_kernels.run_fused_dsa_attention( config=self.config, -@@ -2581,7 +2741,19 @@ class DSAttention(MegatronModule): +@@ -2581,7 +2741,31 @@ class DSAttention(MegatronModule): # =================================== # Get top-k indices # =================================== - if fused_bounds is not None: + if is_kpool: + # cuDNN's token top-k does not implement pool compression; keep dev's subsequent -+ # sparse attention path. ++ # sparse attention path. K was gathered to global sequence order above for CP; ++ # gather the per-token compression gates identically before pooling them together. ++ kpool_gate_score = self.indexer._kpool_gate_score ++ if cp_size > 1 and kpool_gate_score.size(0) in local_cp_kv_lens: ++ if kv_reorder_idx is None: ++ kv_reorder_idx = _build_kv_reorder_idx(kpool_gate_score.size(0)) ++ kpool_gate_score = gather_from_sequence_parallel_region(kpool_gate_score, group=cp_group) ++ if kpool_gate_score.size(0) != kv_reorder_idx.numel(): ++ raise RuntimeError( ++ "DSA gathered KPool gate length mismatch: " ++ f"gate_seqlen={kpool_gate_score.size(0)}, expected={kv_reorder_idx.numel()}" ++ ) ++ kpool_gate_score = kpool_gate_score.index_select(0, kv_reorder_idx) + with torch.no_grad(): + _, topk_indices = fused_qk_topk_kpool( + q, k, weights, self.index_topk, self.indexer.index_kpool, -+ self.indexer._kpool_gate_score, self.indexer.index_kpool_compress_ape, ++ kpool_gate_score, self.indexer.index_kpool_compress_ape, + mask=float_mask, varlen_starts=varlen_starts, varlen_ends=varlen_ends, + key_positions=key_positions, cu_seqlens_kv=cu_seqlens_kv if packed_thd else None, + use_relu=self.config.dsa_indexer_scoring_relu, diff --git a/src/mcore_bridge/utils/megatron_utils.py b/src/mcore_bridge/utils/megatron_utils.py index 419ed539..57dda01b 100644 --- a/src/mcore_bridge/utils/megatron_utils.py +++ b/src/mcore_bridge/utils/megatron_utils.py @@ -18,6 +18,13 @@ mcore_016 = version.parse(megatron.core.__version__) >= version.parse('0.16.0rc0') +# Megatron dev (0.19+) refactored `roll_tensor` to take a LIST of tensors and return a list of rolled +# tensors; 0.16/0.18 take a single tensor and return `(rolled, rolled.sum())`. mcore-bridge's callers +# (gpt_model / mtp_layer) use the single-tensor form, so detect the installed signature once. +import inspect # noqa: E402 + +_ROLL_TAKES_LIST = 'tensors' in inspect.signature(mcore_roll_tensor).parameters + logger = get_logger() @@ -284,6 +291,9 @@ def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=No def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=None): if mcore_016 or packed_seq_params is None: kwargs = {'packed_seq_params': packed_seq_params} if mcore_016 else {} + if _ROLL_TAKES_LIST: + rolled = mcore_roll_tensor([tensor], shifts=shifts, dims=dims, cp_group=cp_group, **kwargs)[0] + return rolled, rolled.sum() return mcore_roll_tensor(tensor, shifts=shifts, dims=dims, cp_group=cp_group, **kwargs) # mcore 0.15 & packed_seq_params return _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group) diff --git a/tests/test_glm5_next.py b/tests/test_glm5_next.py index 643f64a3..3f89e29f 100644 --- a/tests/test_glm5_next.py +++ b/tests/test_glm5_next.py @@ -17,6 +17,7 @@ from mcore_bridge.config.parser import hf_to_mcore_config from mcore_bridge.model.mm_gpts.glm5_next import Glm5NextRMSNorm, _get_physical_cu_seqlens from mcore_bridge.model.register import get_mcore_model +from mcore_bridge.utils import split_cp_inputs def _glm_config(): @@ -125,7 +126,7 @@ def _tiny_glm_config(moe=False, optimized_dsa=False): return config -def _mcore_config(hf_config, tp=1, pp=1, ep=1, sequence_parallel=False, dtype=torch.float32): +def _mcore_config(hf_config, tp=1, pp=1, ep=1, cp=1, sequence_parallel=False, dtype=torch.float32, mtp=0): values = hf_to_mcore_config(hf_config) # HF `glm5_next` resolves to the multimodal type; these are language-model-only fixtures. values['mcore_model_type'] = 'glm5_next' @@ -142,8 +143,11 @@ def _mcore_config(hf_config, tp=1, pp=1, ep=1, sequence_parallel=False, dtype=to pipeline_model_parallel_size=pp, expert_model_parallel_size=ep, expert_tensor_parallel_size=1, + context_parallel_size=cp, sequence_parallel=sequence_parallel, ) + if mtp: + values.update(mtp_num_layers=mtp, mtp_loss_scaling_factor=0.1) return ModelConfig(**values) @@ -155,11 +159,11 @@ def _distributed_session(): @contextmanager -def _parallel_context(tp=1, pp=1, ep=1): +def _parallel_context(tp=1, pp=1, ep=1, cp=1): if not torch.cuda.is_available(): pytest.skip('CUDA is required') world_size = int(os.environ.get('WORLD_SIZE', '1')) - expected_world_size = max(tp * pp, ep) + expected_world_size = tp * pp * cp * ep if world_size != expected_world_size: pytest.skip(f'requires world size {expected_world_size}') local_rank = int(os.environ.get('LOCAL_RANK', '0')) @@ -176,12 +180,12 @@ def _parallel_context(tp=1, pp=1, ep=1): pipeline_model_parallel_size=pp, expert_model_parallel_size=ep, expert_tensor_parallel_size=1, - context_parallel_size=1, + context_parallel_size=cp, ) from megatron.core.process_groups_config import ProcessGroupCollection pg = ProcessGroupCollection.use_mpu_process_groups() assert pg.tp.size() == tp and pg.pp.size() == pp and pg.ep.size() == ep - assert pg.cp.size() == parallel_state.get_context_parallel_world_size() == 1 + assert pg.cp.size() == parallel_state.get_context_parallel_world_size() == cp model_parallel_cuda_manual_seed(123) try: yield @@ -495,20 +499,138 @@ def test_glm5_uses_padded_boundaries_for_physical_thd_slices(): assert _get_physical_cu_seqlens(None) is None -def test_glm5_mtp_layer_is_filtered_for_full_and_stripped_prefixes(): +def test_glm5_mtp_layer_is_filtered_only_when_mtp_disabled(): from mcore_bridge.model.mm_gpts.glm5_next import Glm5NextBridge - bridge = object.__new__(Glm5NextBridge) - bridge.config = SimpleNamespace(num_layers=90) state = { 'model.language_model.layers.44.input_layernorm.weight': 1, 'model.language_model.layers.45.input_layernorm.weight': 2, 'language_model.layers.45.mlp.down_proj.weight': 3, 'layers.45.self_attn.q_proj.weight': 4, } - converted = bridge._convert_hf_state_dict(state, True) + + # MTP off: the extra decoder layer (index num_hidden_layers) has nowhere to go -> filtered. + bridge = object.__new__(Glm5NextBridge) + bridge.config = SimpleNamespace(num_layers=90, mtp_num_layers=0) + converted = bridge._convert_hf_state_dict(dict(state), True) assert converted == {'model.language_model.layers.44.input_layernorm.weight': 1} + # MTP on: the layer-45 tensors are kept for _convert_mtp_layer. + bridge_mtp = object.__new__(Glm5NextBridge) + bridge_mtp.config = SimpleNamespace(num_layers=90, mtp_num_layers=1) + kept = bridge_mtp._convert_hf_state_dict(dict(state), True) + assert 'model.language_model.layers.45.input_layernorm.weight' in kept + + +def _build_glm_mtp_model(mtp=1, moe=True, dtype=torch.float32, tp=1, pp=1, ep=1, cp=1, seed=5): + torch.manual_seed(seed) + hf_config = _tiny_glm_config(moe=moe, optimized_dsa=cp > 1) + config = _mcore_config(hf_config, tp=tp, pp=pp, ep=ep, cp=cp, dtype=dtype, mtp=mtp) + model = get_mcore_model(config)[0].cuda() + return model, config + + +def _lm_of(model): + return model.language_model if hasattr(model, 'language_model') else model + + +def test_glm5_mtp_builds_non_mhc_eh_proj_head(): + """GLM's MTP head is non-mHC despite the mHC backbone: fused eh_proj, no hyper-connected inner block.""" + with _parallel_context(): + model, config = _build_glm_mtp_model(mtp=1) + lm = _lm_of(model) + assert config.mtp_hybrid_override_pattern == 'DE' + assert lm.mtp_process and len(lm.mtp.layers) == 1 + mtp_layer = lm.mtp.layers[0] + assert type(mtp_layer).__name__ == 'Glm5NextMultiTokenPredictionLayer' + assert mtp_layer.mhc_enabled is False + assert mtp_layer.eh_proj is not None and mtp_layer.e_proj is None and mtp_layer.h_proj is None + inner = mtp_layer.mtp_model_layer + assert inner.is_mtp_layer and len(inner.layers) == 2 + # No hyper-connection wrapping on the MTP inner sublayers (the checkpoint carries no hc_* there). + assert all(not hasattr(sub, 'hyper_connection') for sub in inner.layers) + # The outer decoder must be the stack that suppresses the multi-stream tensor. + assert type(lm.decoder).__name__ == 'Glm5NextHybridStack' + + +def test_glm5_mtp_bridge_roundtrip_matches_checkpoint_layout(): + """mcore -> hf export of the MTP head matches the real GLM-5.3-Flash key layout and re-imports bit-exactly.""" + with _parallel_context(): + model_a, config = _build_glm_mtp_model(mtp=1, seed=5) + exported = _export_to_hf(config, model_a) + hf_idx = config.num_layers // 2 # MTP head lives at decoder-layer index num_hidden_layers + p = f'model.language_model.layers.{hf_idx}.' + for key in ('enorm.weight', 'hnorm.weight', 'eh_proj.weight', 'shared_head.norm.weight', + 'input_layernorm.weight', 'post_attention_layernorm.weight', 'mlp.gate.weight', + 'mlp.gate.e_score_correction_bias', 'self_attn.q_a_proj.weight', + 'self_attn.kv_a_proj_with_mqa.weight', 'self_attn.indexer.wq_b.weight'): + assert p + key in exported, f'MTP export missing {p}{key}' + + model_b, _ = _build_glm_mtp_model(mtp=1, seed=999) + lazy = {key: _LazyTensor(value) for key, value in exported.items()} + list(config.bridge._convert([model_b], lazy, '', True, 'Reloading test: ')) + sd_a = _lm_of(model_a).mtp.state_dict() + sd_b = _lm_of(model_b).mtp.state_dict() + assert sd_a.keys() == sd_b.keys() + for key in sd_a: + torch.testing.assert_close( + sd_b[key].cpu(), sd_a[key].cpu(), atol=0, rtol=0, msg=lambda m, k=key: f'{k}: {m}') + + +def test_glm5_mtp_forward_backward_flows_grads_to_head(): + """The MTP loss path runs and every trainable MTP parameter receives a gradient. + + BF16 (how GLM trains): the FP32 KDA backward routes through a TileLang kernel whose nvcc + toolchain is unrelated to MTP; BF16 uses the FLA/Triton chunk-KDA path. + """ + with _parallel_context(): + model, _ = _build_glm_mtp_model(mtp=1, dtype=torch.bfloat16) + lm = _lm_of(model).train() + input_ids, position_ids, attention_mask = _model_inputs(batch=2, sequence=8) + labels = torch.randint(1, 128, (2, 8), device='cuda') + loss_mask = torch.ones(2, 8, device='cuda') + out = lm(input_ids, position_ids, attention_mask, labels=labels, loss_mask=loss_mask) + loss = out if out.dim() == 0 else out.float().mean() + assert torch.isfinite(loss).item() + loss.backward() + nograd = [n for n, p in lm.mtp.named_parameters() if p.requires_grad and p.grad is None] + assert not nograd, f'MTP params without grad: {nograd}' + head = dict(lm.mtp.named_parameters()) + for probe in ('layers.0.eh_proj.weight', 'layers.0.enorm.weight', 'layers.0.hnorm.weight', + 'layers.0.final_layernorm.weight'): + assert probe in head, f'MTP head missing {probe}' + assert head[probe].grad is not None and torch.isfinite(head[probe].grad.float()).all() + + +def test_glm5_mtp_cp2_packed_forward_backward(): + """KPool gates must follow the same CP gather/reorder as their indexer keys.""" + if os.environ.get('GLM5_MTP_PARALLEL_TEST') != 'cp2': + pytest.skip('run with GLM5_MTP_PARALLEL_TEST=cp2 torchrun --nproc-per-node=2') + with _parallel_context(cp=2): + model, _ = _build_glm_mtp_model(mtp=1, dtype=torch.bfloat16, cp=2) + lm = _lm_of(model).train() + input_ids, position_ids, packed_seq_params = _packed_inputs([16]) + labels = torch.randint(1, 128, input_ids.shape, device='cuda') + loss_mask = torch.ones_like(labels, dtype=torch.bool) + cu_seqlens = packed_seq_params.cu_seqlens_q + position_ids = split_cp_inputs(position_ids, cu_seqlens, -1) + labels = split_cp_inputs(labels, cu_seqlens, -1) + loss_mask = split_cp_inputs(loss_mask, cu_seqlens, -1) + + out = model( + input_ids, + position_ids, + None, + labels=labels, + loss_mask=loss_mask, + packed_seq_params=packed_seq_params, + ) + loss = out if out.dim() == 0 else out.float().mean() + assert torch.isfinite(loss).item() + loss.backward() + missing = [name for name, param in lm.mtp.named_parameters() if param.requires_grad and param.grad is None] + assert not missing, f'MTP params without grad: {missing}' + def test_glm5_kda_head_sharding_matches_world8(): if int(os.environ.get('WORLD_SIZE', '1')) != 8 or not torch.cuda.is_available(): diff --git a/tests/test_qwen4_exp_mtp.py b/tests/test_qwen4_exp_mtp.py new file mode 100644 index 00000000..41736009 --- /dev/null +++ b/tests/test_qwen4_exp_mtp.py @@ -0,0 +1,330 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""MTP (multi-token prediction) tests for Qwen3.8-Flash-Next (mcore type `qwen4_exp`). + +The MTP head is the `residual_linear_shared` fusion: separate fc_embedding/fc_hidden (e_proj/h_proj), +a joint multi-stream pre_fc_norm_hidden, and a hyper_connection_mixer contraction -- built on the +gated-HC backbone but NOT Megatron's enable_hyper_connections mHC. Fixtures are tiny synthetic +configs (last backbone layer is full-attention so the MTP inner block, which is always +full-attention, matches). The QSA indexer's hard top-k selection is non-differentiable in the +sbhd/bool-mask path, so its params legitimately carry no gradient there -- exactly as the backbone's +full-attention indexer does -- and are excluded from the grad-flow assertion. +""" +import os +import pytest +import torch +import uuid +from contextlib import contextmanager +from megatron.core import parallel_state +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + +from mcore_bridge.config import ModelConfig +from mcore_bridge.config.parser import hf_to_mcore_config +from mcore_bridge.model.register import get_mcore_model + + +def _tiny_qwen4exp_config(): + cfg = pytest.importorskip('transformers.models.qwen4_exp.configuration_qwen4_exp') + text = { + 'model_type': 'qwen4_exp_text', + 'vocab_size': 512, + 'hidden_size': 128, + 'head_dim': 64, + 'num_attention_heads': 4, + 'num_key_value_heads': 4, + 'num_hidden_layers': 4, + # Last layer full-attention so layer_specs[-1] (the MTP inner block) is full-attention + QSA. + 'layer_types': ['linear_attention', 'linear_attention', 'linear_attention', 'qwen_sparse_attention'], + 'full_attention_interval': 4, + 'num_experts': 4, + 'num_experts_per_tok': 2, + 'moe_intermediate_size': 64, + 'shared_expert_intermediate_size': 64, + 'hc_count': 4, + 'hc_lowrank': 32, + 'output_gate_type': 'sigmoid', + 'partial_rotary_factor': 0.25, + 'rope_parameters': { + 'mrope_interleaved': True, + 'mrope_section': [11, 11, 10], + 'partial_rotary_factor': 0.25, + 'rope_theta': 10000000, + 'rope_type': 'default' + }, + 'linear_num_key_heads': 16, + 'linear_num_value_heads': 48, + 'linear_key_head_dim': 128, + 'linear_value_head_dim': 128, + 'linear_conv_kernel_dim': 4, + 'indexer_n_heads': 4, + 'indexer_kv_heads': 1, + 'indexer_head_dim': 128, + 'indexer_budget': 2048, + 'indexer_compress_ratio': 4, + 'ple_layer_ids': [2], + 'ple_embed_dim': 32, + 'ple_conv_kernel_size': 4, + 'ngram_size': 3, + 'heads_per_ngram': 8, + 'ngram_vocab_size_base': 256, + 'make_ngram_vocab_size_divisible_by': 128, + 'split_ngram_parts': 2, + 'rms_norm_eps': 1e-6, + 'eos_token_id': 248044, + 'seed': 1234, + 'mtp_num_hidden_layers': 1, + 'mtp_use_dedicated_embeddings': False, + 'mtp': { + 'hybrid': True, + 'layer_types': ['full_attention'], + 'mtp_use_hidden_state_from_layer': None, + 'num_hidden_layers': 1, + 'rope_theta': 10000000 + }, + } + vision = { + 'model_type': 'qwen4_exp', + 'depth': 2, + 'hidden_size': 128, + 'intermediate_size': 256, + 'num_heads': 4, + 'out_hidden_size': 128, + 'in_channels': 3, + 'patch_size': 16, + 'spatial_merge_size': 2, + 'temporal_patch_size': 2, + 'num_position_embeddings': 2304, + } + return cfg.Qwen4ExpConfig( + text_config=text, + vision_config=vision, + image_token_id=248056, + video_token_id=248057, + vision_start_token_id=248053, + vision_end_token_id=248054, + tie_word_embeddings=False, + language_model_only=True, + ) + + +def _mcore_config(hf_config, mtp=1, dtype=torch.bfloat16, tp=1, pp=1, ep=1): + values = hf_to_mcore_config(hf_config) + values['mcore_model_type'] = 'qwen4_exp' + values['hf_config'] = hf_config + values.update( + params_dtype=dtype, + pipeline_dtype=dtype, + bf16=dtype == torch.bfloat16, + perform_initialization=True, + use_cpu_initialization=False, + language_model_only=True, + moe_grouped_gemm=True, + overlap_p2p_comm=False, + tensor_model_parallel_size=tp, + pipeline_model_parallel_size=pp, + expert_model_parallel_size=ep, + expert_tensor_parallel_size=1, + sequence_parallel=False, + recompute_granularity=None, + ) + if mtp: + values.update(mtp_num_layers=mtp, mtp_loss_scaling_factor=0.1) + return ModelConfig(**values) + + +@pytest.fixture(scope='session', autouse=True) +def _distributed_session(): + yield + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +@contextmanager +def _parallel_context(tp=1, pp=1, ep=1): + if not torch.cuda.is_available(): + pytest.skip('CUDA is required') + world_size = int(os.environ.get('WORLD_SIZE', '1')) + if world_size != max(tp * pp, ep): + pytest.skip(f'requires world size {max(tp * pp, ep)}') + local_rank = int(os.environ.get('LOCAL_RANK', '0')) + torch.cuda.set_device(local_rank) + # Another test module in the same pytest process may already have initialized (and not torn + # down) the default group and/or the model-parallel state, so repair both instead of asserting. + created_pg = False + if not torch.distributed.is_initialized(): + if world_size == 1: + torch.distributed.init_process_group( + 'nccl', init_method=f'file:///tmp/q4e-mtp-{uuid.uuid4().hex}', rank=0, world_size=1) + else: + torch.distributed.init_process_group('nccl') + created_pg = True + if parallel_state.model_parallel_is_initialized(): + parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=tp, + pipeline_model_parallel_size=pp, + expert_model_parallel_size=ep, + expert_tensor_parallel_size=1, + context_parallel_size=1, + ) + model_parallel_cuda_manual_seed(123) + try: + yield + finally: + if world_size > 1: + torch.distributed.barrier() + parallel_state.destroy_model_parallel() + from mcore_bridge.bridge import gpt_bridge + gpt_bridge.EP_PP_GROUP = gpt_bridge.EP_PP_RANK = gpt_bridge.EP_PP_SIZE = None + if created_pg: + torch.distributed.destroy_process_group() + + +class _Lazy: + + def __init__(self, t): + self.t = t + + def load(self): + return self.t + + +def _build(mtp=1, seed=5, dtype=torch.bfloat16, tp=1, pp=1, ep=1): + torch.manual_seed(seed) + hf_config = _tiny_qwen4exp_config() + config = _mcore_config(hf_config, mtp=mtp, dtype=dtype, tp=tp, pp=pp, ep=ep) + model = get_mcore_model(config)[0].cuda() + return model, config + + +def _lm(model): + return model.language_model if hasattr(model, 'language_model') else model + + +def _export(config, model): + out = {} + for k, v in config.bridge._convert([model], {}, '', False, 'Exporting test: '): + out[k] = v.load() if hasattr(v, 'load') else v + return out + + +def test_qwen4exp_mtp_builds_residual_linear_shared_head(): + """The MTP head uses e_proj/h_proj + joint hnorm + mixer, not a fused eh_proj/final_layernorm.""" + with _parallel_context(): + model, config = _build(mtp=1) + lm = _lm(model) + assert lm.mtp_process and len(lm.mtp.layers) == 1 + ml = lm.mtp.layers[0] + assert type(ml).__name__ == 'Qwen4ExpMultiTokenPredictionLayer' + assert ml.mhc_enabled is True # selects the e_proj/h_proj branch + assert ml.eh_proj is None and not hasattr(ml, 'final_layernorm') and not hasattr(ml, 'hc_head_fn') + assert type(ml.e_proj).__name__ == 'TEColumnParallelLinear' + assert type(ml.h_proj).__name__ == 'TEColumnParallelLinear' + # joint multi-stream norm: hc_count * hidden_size + assert ml.hnorm.weight.shape[0] == config.hc_count * config.hidden_size + assert type(ml.hyper_connection_mixer).__name__ == 'Qwen4ExpTextGatedResidual' + inner = ml.transformer_layer + assert type(inner).__name__ == 'Qwen4ExpMTPInnerLayer' + assert getattr(inner.self_attention, 'indexer', None) is not None # full-attention + QSA + assert inner.ple is None # no PLE in the MTP block + + +def test_qwen4exp_mtp_bridge_roundtrip_matches_checkpoint_layout(): + """mcore -> hf export matches the real Qwen3.8-Flash-Next mtp.* layout and re-imports bit-exactly.""" + with _parallel_context(): + model_a, config = _build(mtp=1, seed=5) + exported = _export(config, model_a) + for key in ('mtp.fc_embedding.weight', 'mtp.fc_hidden.weight', 'mtp.pre_fc_norm_embedding.weight', + 'mtp.pre_fc_norm_hidden.weight', 'mtp.hyper_connection_mixer.hc_norm.weight', + 'mtp.hyper_connection_mixer.input_mix_weight_down.weight', + 'mtp.hyper_connection_mixer.input_mix_weight_up.weight', 'mtp.layers.0.self_attn.q_proj.weight', + 'mtp.layers.0.self_attn.indexer.index_qk_proj.weight', 'mtp.layers.0.mlp.gate.weight', + 'mtp.layers.0.attn_hyper_connection.hc_norm.weight', + 'mtp.layers.0.mlp_hyper_connection.hc_norm.weight'): + assert key in exported, f'MTP export missing {key}' + assert not any(k.startswith('mtp.') and 'eh_proj' in k for k in exported) + + model_b, _ = _build(mtp=1, seed=999) + list( + config.bridge._convert([model_b], { + k: _Lazy(v) + for k, v in exported.items() + }, '', True, 'Reloading test: ')) + sd_a = _lm(model_a).mtp.state_dict() + sd_b = _lm(model_b).mtp.state_dict() + assert sd_a.keys() == sd_b.keys() + for key in sd_a: + torch.testing.assert_close( + sd_b[key].cpu(), sd_a[key].cpu(), atol=0, rtol=0, msg=lambda m, k=key: f'{k}: {m}') + + +def test_qwen4exp_mtp_forward_backward_flows_grads_to_head(): + """The MTP loss path runs; the head and inner block get gradients (except the non-differentiable indexer).""" + with _parallel_context(): + model, _ = _build(mtp=1) + lm = _lm(model).train() + b, s = 2, 16 + g = torch.Generator(device='cuda').manual_seed(23) + input_ids = torch.randint(1, 500, (b, s), device='cuda', generator=g) + position_ids = torch.arange(s, device='cuda').unsqueeze(0).expand(b, -1) + attention_mask = torch.triu(torch.ones(b, 1, s, s, device='cuda', dtype=torch.bool), diagonal=1) + labels = torch.randint(1, 500, (b, s), device='cuda', generator=g) + loss_mask = torch.ones(b, s, device='cuda', dtype=torch.bool) + out = lm(input_ids, position_ids, attention_mask, labels=labels, loss_mask=loss_mask) + loss = out if out.dim() == 0 else out.float().mean() + assert torch.isfinite(loss).item() + loss.backward() + # The QSA indexer's hard top-k selection is non-differentiable in the bool-mask path (the + # backbone's full-attention indexer carries no gradient there either), so exclude it. + nograd = [n for n, p in lm.mtp.named_parameters() if p.requires_grad and p.grad is None and 'indexer' not in n] + assert not nograd, f'MTP params without grad: {nograd}' + head = dict(lm.mtp.named_parameters()) + for probe in ('layers.0.e_proj.weight', 'layers.0.h_proj.weight', 'layers.0.enorm.weight', + 'layers.0.hnorm.weight', 'layers.0.hyper_connection_mixer.input_mix_weight_down.weight'): + assert head[probe].grad is not None and torch.isfinite(head[probe].grad.float()).all() + + +def test_qwen4exp_mtp_pp2_forward_backward(): + """PP transports both n*H gated-HC states and the H-wide embedding needed by MTP.""" + if os.environ.get('QWEN4EXP_MTP_PARALLEL_TEST') != 'pp2': + pytest.skip('run with QWEN4EXP_MTP_PARALLEL_TEST=pp2 torchrun --nproc-per-node=2') + with _parallel_context(pp=2): + from megatron.core.pipeline_parallel.schedules import get_forward_backward_func + + model, _ = _build(mtp=1, pp=2) + b, s = 2, 16 + generator = torch.Generator(device='cuda').manual_seed(23) + input_ids = torch.randint(1, 500, (b, s), device='cuda', generator=generator) + position_ids = torch.arange(s, device='cuda').unsqueeze(0).expand(b, -1) + attention_mask = torch.triu(torch.ones(b, 1, s, s, device='cuda', dtype=torch.bool), diagonal=1) + labels = torch.randint(1, 500, (b, s), device='cuda', generator=generator) + loss_mask = torch.ones(b, s, device='cuda', dtype=torch.bool) + + def forward_step(data_iterator, stage_model): + ids, positions, mask, target, target_mask = next(data_iterator) + output = stage_model(ids, positions, mask, labels=target, loss_mask=target_mask) + + def loss_func(tensor): + loss = tensor if tensor.dim() == 0 else tensor.float().mean() + return loss, {'loss': loss.detach()} + + return output, loss_func + + losses = get_forward_backward_func()( + forward_step_func=forward_step, + data_iterator=iter([(input_ids, position_ids, attention_mask, labels, loss_mask)]), + model=model, + num_microbatches=1, + seq_length=s, + micro_batch_size=b, + forward_only=False, + ) + if parallel_state.is_pipeline_last_stage(): + assert len(losses) == 1 and torch.isfinite(losses[0]['loss']).all() + lm = _lm(model) + missing = [ + name for name, param in lm.mtp.named_parameters() + if param.requires_grad and param.grad is None and 'indexer' not in name + ] + assert not missing, f'MTP params without grad: {missing}' + else: + assert losses == [] From 64a004bbaec27f5aa4c40fed54fe4c3f4f58396a Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Mon, 21 Sep 2026 15:45:25 +0800 Subject: [PATCH 2/3] fix --- src/mcore_bridge/model/gpts/qwen4_exp.py | 22 ++------ src/mcore_bridge/model/mm_gpts/glm5_next.py | 37 +++++++++++++ src/mcore_bridge/model/modules/mtp_layer.py | 5 ++ tests/test_glm5_next.py | 49 ++++++++++++++++-- tests/test_qwen4_exp_mtp.py | 57 +++++++++++++++++++-- 5 files changed, 144 insertions(+), 26 deletions(-) diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index e64894c9..bb30b161 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/src/mcore_bridge/model/gpts/qwen4_exp.py @@ -371,25 +371,9 @@ def _postprocess(self, hidden_states): # Contract the multi-stream [s, b, n*H] to [s, b, H] with Qwen4Exp's gated mixer (no final norm). return self.hyper_connection_mixer(hidden_states) - def _get_embeddings(self, - input_ids, - position_ids, - embedding, - hidden_states, - packed_seq_params=None, - decoder_input=None): - input_ids, position_ids, decoder_input, hidden_states = super()._get_embeddings( - input_ids, position_ids, embedding, hidden_states, packed_seq_params, decoder_input) - # Stash the rolled ids so _proj_and_transformer_layer can forward them to the inner - # Qwen4ExpMTPInnerLayer (its QSA indexer needs position_ids under packing/CP; PLE is absent). - self._mtp_input_ids = input_ids - self._mtp_position_ids = position_ids - return input_ids, position_ids, decoder_input, hidden_states - - def _proj_and_transformer_layer(self, *args, **kwargs): - kwargs.setdefault('input_ids', getattr(self, '_mtp_input_ids', None)) - kwargs.setdefault('position_ids', getattr(self, '_mtp_position_ids', None)) - return super()._proj_and_transformer_layer(*args, **kwargs) + def _get_inner_layer_kwargs(self, input_ids, position_ids): + # Pass per-depth state through the checkpoint boundary instead of mutable module attributes. + return {'input_ids': input_ids, 'position_ids': position_ids} class Qwen4ExpBridge(Qwen3NextBridge): diff --git a/src/mcore_bridge/model/mm_gpts/glm5_next.py b/src/mcore_bridge/model/mm_gpts/glm5_next.py index ee3dbbd3..9aceec23 100644 --- a/src/mcore_bridge/model/mm_gpts/glm5_next.py +++ b/src/mcore_bridge/model/mm_gpts/glm5_next.py @@ -18,6 +18,7 @@ try: from megatron.core.models.hybrid.hybrid_block import HybridStack from megatron.core.models.hybrid.hybrid_layer_allocation import parse_hybrid_pattern, select_pipeline_segment + from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region from megatron.core.transformer.module import mark_keep_in_fp32 from megatron.core.transformer.multi_token_prediction import \ MultiTokenPredictionLayer as McoreMultiTokenPredictionLayer @@ -146,6 +147,42 @@ def __init__(self, config, *args, **kwargs): mtp_config.enable_hyper_connections = False super().__init__(mtp_config, *args, **kwargs) + def _get_embeddings(self, + input_ids, + position_ids, + embedding, + hidden_states, + packed_seq_params=None, + padding_mask=None, + sequence_roll_context=None, + roll_depth=0): + # Decoder/MoE routing consumes the sequence-parallel shard, but MTP must roll the mask + # beside the unsharded input_ids. Reconstruct it only around the roll, then restore the + # shard expected by the nested HybridStack. + sequence_parallel_mask = ( + padding_mask is not None and self.config.sequence_parallel and self.tp_group.size() > 1) + if sequence_parallel_mask: + padding_mask = gather_from_sequence_parallel_region( + padding_mask.transpose(0, 1).contiguous(), + tensor_parallel_output_grad=False, + group=self.tp_group, + ).transpose(0, 1).contiguous() + + input_ids, position_ids, padding_mask, decoder_input, hidden_states = super()._get_embeddings( + input_ids=input_ids, + position_ids=position_ids, + embedding=embedding, + hidden_states=hidden_states, + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, + sequence_roll_context=sequence_roll_context, + roll_depth=roll_depth, + ) + + if sequence_parallel_mask: + padding_mask = padding_mask.chunk(self.tp_group.size(), dim=-1)[self.tp_group.rank()].contiguous() + return input_ids, position_ids, padding_mask, decoder_input, hidden_states + class Glm5NextHybridModel(HybridModel): extra_forward_keys = () diff --git a/src/mcore_bridge/model/modules/mtp_layer.py b/src/mcore_bridge/model/modules/mtp_layer.py index b6d8b306..1dd1eeb7 100644 --- a/src/mcore_bridge/model/modules/mtp_layer.py +++ b/src/mcore_bridge/model/modules/mtp_layer.py @@ -71,6 +71,10 @@ def transformer_layer(self): """ return self._modules.get('mtp_model_layer', None) or self._modules.get('transformer_layer', None) + def _get_inner_layer_kwargs(self, input_ids, position_ids): + """Return model-specific rolled inputs consumed by the inner transformer layer.""" + return {} + def forward( self, input_ids: torch.Tensor, @@ -102,6 +106,7 @@ def forward( hidden_states=hidden_states, decoder_input=decoder_input, ) + kwargs.update(self._get_inner_layer_kwargs(input_ids, position_ids)) assert not self.transformer_layer.self_attention.config.apply_rope_fusion packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' if self.config.position_embedding_type == 'rope' and packed_seq: diff --git a/tests/test_glm5_next.py b/tests/test_glm5_next.py index 3f89e29f..0b210050 100644 --- a/tests/test_glm5_next.py +++ b/tests/test_glm5_next.py @@ -522,10 +522,19 @@ def test_glm5_mtp_layer_is_filtered_only_when_mtp_disabled(): assert 'model.language_model.layers.45.input_layernorm.weight' in kept -def _build_glm_mtp_model(mtp=1, moe=True, dtype=torch.float32, tp=1, pp=1, ep=1, cp=1, seed=5): +def _build_glm_mtp_model(mtp=1, moe=True, dtype=torch.float32, tp=1, pp=1, ep=1, cp=1, sequence_parallel=False, seed=5): torch.manual_seed(seed) - hf_config = _tiny_glm_config(moe=moe, optimized_dsa=cp > 1) - config = _mcore_config(hf_config, tp=tp, pp=pp, ep=ep, cp=cp, dtype=dtype, mtp=mtp) + hf_config = _tiny_glm_config(moe=moe, optimized_dsa=cp > 1 or sequence_parallel) + config = _mcore_config( + hf_config, + tp=tp, + pp=pp, + ep=ep, + cp=cp, + sequence_parallel=sequence_parallel, + dtype=dtype, + mtp=mtp, + ) model = get_mcore_model(config)[0].cuda() return model, config @@ -632,6 +641,40 @@ def test_glm5_mtp_cp2_packed_forward_backward(): assert not missing, f'MTP params without grad: {missing}' +def test_glm5_mtp_tp2_sequence_parallel_packed_forward_backward(): + """MTP rolls a TP-full padding mask, then restores the SP shard used by its MoE router.""" + if os.environ.get('GLM5_MTP_PARALLEL_TEST') != 'tp2_sp': + pytest.skip('run with GLM5_MTP_PARALLEL_TEST=tp2_sp torchrun --nproc-per-node=2') + pytest.importorskip('tilelang') + with _parallel_context(tp=2): + model, _ = _build_glm_mtp_model( + mtp=1, + dtype=torch.bfloat16, + tp=2, + sequence_parallel=True, + ) + lm = _lm_of(model).train() + input_ids, position_ids, packed_seq_params = _packed_inputs([16]) + packed_seq_params.seq_lens = torch.tensor([13], device='cuda', dtype=torch.int32) + labels = torch.randint(1, 128, input_ids.shape, device='cuda') + loss_mask = torch.ones_like(labels, dtype=torch.bool) + loss_mask[:, 13:] = False + + out = model( + input_ids, + position_ids, + None, + labels=labels, + loss_mask=loss_mask, + packed_seq_params=packed_seq_params, + ) + loss = out if out.dim() == 0 else out.float().mean() + assert torch.isfinite(loss).item() + loss.backward() + missing = [name for name, param in lm.mtp.named_parameters() if param.requires_grad and param.grad is None] + assert not missing, f'MTP params without grad: {missing}' + + def test_glm5_kda_head_sharding_matches_world8(): if int(os.environ.get('WORLD_SIZE', '1')) != 8 or not torch.cuda.is_available(): pytest.skip('run with torchrun --nproc-per-node=8 to validate eight-way KDA head sharding') diff --git a/tests/test_qwen4_exp_mtp.py b/tests/test_qwen4_exp_mtp.py index 41736009..532920e3 100644 --- a/tests/test_qwen4_exp_mtp.py +++ b/tests/test_qwen4_exp_mtp.py @@ -106,7 +106,7 @@ def _tiny_qwen4exp_config(): ) -def _mcore_config(hf_config, mtp=1, dtype=torch.bfloat16, tp=1, pp=1, ep=1): +def _mcore_config(hf_config, mtp=1, dtype=torch.bfloat16, tp=1, pp=1, ep=1, recompute=False, shared=False): values = hf_to_mcore_config(hf_config) values['mcore_model_type'] = 'qwen4_exp' values['hf_config'] = hf_config @@ -124,7 +124,10 @@ def _mcore_config(hf_config, mtp=1, dtype=torch.bfloat16, tp=1, pp=1, ep=1): expert_model_parallel_size=ep, expert_tensor_parallel_size=1, sequence_parallel=False, - recompute_granularity=None, + recompute_granularity='full' if recompute else None, + recompute_method='uniform' if recompute else None, + recompute_num_layers=1 if recompute else None, + mtp_shared_weights=shared, ) if mtp: values.update(mtp_num_layers=mtp, mtp_loss_scaling_factor=0.1) @@ -188,10 +191,10 @@ def load(self): return self.t -def _build(mtp=1, seed=5, dtype=torch.bfloat16, tp=1, pp=1, ep=1): +def _build(mtp=1, seed=5, dtype=torch.bfloat16, tp=1, pp=1, ep=1, recompute=False, shared=False): torch.manual_seed(seed) hf_config = _tiny_qwen4exp_config() - config = _mcore_config(hf_config, mtp=mtp, dtype=dtype, tp=tp, pp=pp, ep=ep) + config = _mcore_config(hf_config, mtp=mtp, dtype=dtype, tp=tp, pp=pp, ep=ep, recompute=recompute, shared=shared) model = get_mcore_model(config)[0].cuda() return model, config @@ -283,6 +286,52 @@ def test_qwen4exp_mtp_forward_backward_flows_grads_to_head(): assert head[probe].grad is not None and torch.isfinite(head[probe].grad.float()).all() +def test_mtp_inner_layer_kwargs_are_model_opt_in(): + from mcore_bridge.model.modules.mtp_layer import MultiTokenPredictionLayer + layer = object.__new__(MultiTokenPredictionLayer) + assert layer._get_inner_layer_kwargs(None, None) == {} + + +def test_qwen4exp_shared_mtp_recompute_restores_each_depth_ids(monkeypatch): + """Checkpoint recompute must replay each shared MTP depth with that depth's rolled IDs.""" + from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper + monkeypatch.setattr(MTPLossLoggingHelper, 'tracker', {}) + with _parallel_context(): + model, config = _build(mtp=2, recompute=True, shared=True) + lm = _lm(model).train() + assert config.mtp_num_layers == 1 and config.mtp_unroll_steps == 2 + assert len(lm.mtp.layers) == 1 + + calls = [] + layer = lm.mtp.layers[0] + original = layer._proj_and_transformer_layer + + def record_ids(*args, **kwargs): + calls.append((kwargs['input_ids'].detach().clone(), kwargs['position_ids'].detach().clone())) + return original(*args, **kwargs) + + layer._proj_and_transformer_layer = record_ids + b, s = 1, 16 + generator = torch.Generator(device='cuda').manual_seed(29) + input_ids = torch.randint(1, 500, (b, s), device='cuda', generator=generator) + position_ids = torch.arange(s, device='cuda').unsqueeze(0) + attention_mask = torch.triu(torch.ones(b, 1, s, s, device='cuda', dtype=torch.bool), diagonal=1) + labels = torch.randint(1, 500, (b, s), device='cuda', generator=generator) + loss_mask = torch.ones(b, s, device='cuda', dtype=torch.bool) + + output = lm(input_ids, position_ids, attention_mask, labels=labels, loss_mask=loss_mask) + loss = output if output.dim() == 0 else output.float().mean() + loss.backward() + + assert len(calls) == 4, f'expected two forwards and two recomputes, got {len(calls)}' + assert not torch.equal(calls[0][0], calls[1][0]) + for field in range(2): + torch.testing.assert_close(calls[2][field], calls[1][field], atol=0, rtol=0) + torch.testing.assert_close(calls[3][field], calls[0][field], atol=0, rtol=0) + assert not hasattr(layer, '_mtp_input_ids') + assert not hasattr(layer, '_mtp_position_ids') + + def test_qwen4exp_mtp_pp2_forward_backward(): """PP transports both n*H gated-HC states and the H-wide embedding needed by MTP.""" if os.environ.get('QWEN4EXP_MTP_PARALLEL_TEST') != 'pp2': From fe6d75ceceb9cda4dcc0546ba367dd57885115df Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Mon, 21 Sep 2026 17:17:19 +0800 Subject: [PATCH 3/3] fix --- src/mcore_bridge/bridge/gpt_bridge.py | 2 ++ src/mcore_bridge/model/gpts/qwen4_exp.py | 17 ++++++++++++++++- src/mcore_bridge/model/mm_gpts/glm5_next.py | 16 ++++++++++++++++ 3 files changed, 34 insertions(+), 1 deletion(-) diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index b9ca26aa..7493fa62 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/src/mcore_bridge/bridge/gpt_bridge.py @@ -1975,6 +1975,8 @@ def _convert_mtp_layer(self, lm_model, hf_state_dict, hf_prefix: str, layer_idx: origin_hf_state_dict = hf_state_dict hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) if len(hf_state_dict) == 0: + if self._peft_format: + return {} logger.info(f'MTP Layer {mtp_layer.layer_number} safetensors weights not found, ' 'this part will be randomly initialized.') for param in mtp_layer.parameters(): diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index bb30b161..4414371d 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/src/mcore_bridge/model/gpts/qwen4_exp.py @@ -378,6 +378,7 @@ def _get_inner_layer_kwargs(self, input_ids, position_ids): class Qwen4ExpBridge(Qwen3NextBridge): hf_mixer_prefix = 'model.' + additional_dim0_keys = {'e_proj', 'h_proj'} def _save_missing_weights(self, saver, saved_keys, source_model_dir=None) -> None: # PLE export emits every shard. If it omits the scale, these are already @@ -642,8 +643,14 @@ def _convert_mtp_layer(self, lm_model, hf_state_dict, hf_prefix: str, layer_idx: hf_prefix = f'{hf_prefix}{layer_idx}.' # 'mtp.layers.0.' if to_mcore: origin_hf_state_dict = hf_state_dict + mtp_present = len(self._remove_prefix(origin_hf_state_dict, 'mtp.')) > 0 hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) - if len(hf_state_dict) == 0: + if not mtp_present: + if self._peft_format: + # A PEFT/adapter checkpoint carries no base MTP weights -- those were already + # loaded from the base checkpoint the adapter sits on top of. Re-initializing + # here would wipe the (frozen) base MTP head + inner block. + return {} logger.info(f'MTP layer {layer_idx} safetensors weights not found, ' 'this part will be randomly initialized.') for param in mtp_layer.parameters(): @@ -743,6 +750,14 @@ def _set_transformer_layer(self, transformer_layer_spec): layer_spec.module = Qwen4ExpLayer def get_mtp_block_spec(self, transformer_layer_spec, vp_stage: Optional[int] = None): + 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). + raise NotImplementedError( + 'Qwen3.8-Flash-Next exposes a single top-level `mtp.*` head, so multiple independent ' + '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: diff --git a/src/mcore_bridge/model/mm_gpts/glm5_next.py b/src/mcore_bridge/model/mm_gpts/glm5_next.py index 9aceec23..5582f852 100644 --- a/src/mcore_bridge/model/mm_gpts/glm5_next.py +++ b/src/mcore_bridge/model/mm_gpts/glm5_next.py @@ -550,6 +550,11 @@ def _convert_mtp_layer(self, lm_model, hf_state_dict, hf_prefix: str, layer_idx: origin_hf_state_dict = hf_state_dict hf_state_dict = self._remove_prefix(hf_state_dict, layer_prefix) if len(hf_state_dict) == 0: + if self._peft_format: + # A PEFT/adapter checkpoint carries no base MTP weights -- those came from the + # base checkpoint the adapter sits on top of. Re-initializing here would wipe + # the (frozen) base MTP head + inner HybridStack. + return {} logger.info(f'MTP layer {hf_layer_idx} safetensors weights not found, ' 'this part will be randomly initialized.') for param in mtp_layer.parameters(): @@ -601,6 +606,17 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): if config.dsa_indexer_loss_coeff: raise NotImplementedError('The current model has no KPool indexer auxiliary loss; ' 'use dsa_indexer_loss_coeff=0') + if config.mtp_num_layers and ((config.mtp_unroll_steps or 1) > 1 + or getattr(config, 'mtp_use_repeated_layer', False)): + # Only single-step MTP is wired for GLM's hybrid path. The hybrid pattern is built from + # mtp_num_layers, and the upstream execution/loss loop iterates mtp_num_layers and reads + # neither mtp_unroll_steps (World-1 shared-weights depth) nor mtp_use_repeated_layer, so a + # multi-step request would silently run one step; the weight converter would also index + # physical mtp.layers[depth] that a repeated build never creates. + raise NotImplementedError( + 'GLM-5.3-Flash MTP is validated only for single-step (mtp_num_layers=1). Multi-step ' + 'MTP via mtp_shared_weights / mtp_use_repeated_layer / mtp_num_layers>1 is not yet ' + 'wired into the hybrid execution path; use mtp_num_layers=1.') if config.context_parallel_size > 1: if config.cp_comm_type != 'all_gather': logger.warning_once("GLM5-Next under context parallelism requires cp_comm_type='all_gather'; "