support qwen4_exp & glm5_next MTP

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():

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(
@@ -449,6 +454,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

src/mcore_bridge/model/modules/mtp_layer.py

@@ -60,6 +60,21 @@ 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 _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,
@@ -91,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:

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)
 

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()
 
 
@@ -291,6 +298,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)

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,181 @@ 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, sequence_parallel=False, seed=5):
+    torch.manual_seed(seed)
+    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
+
+
+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_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():