From ef35ea826b44abecd94168024e748501041c57fa Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Tue, 15 Sep 2026 21:26:03 +0800 Subject: [PATCH] glm5 next cp --- src/mcore_bridge/model/mm_gpts/glm5_next.py | 23 ++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/src/mcore_bridge/model/mm_gpts/glm5_next.py b/src/mcore_bridge/model/mm_gpts/glm5_next.py index 31508da..deced8b 100644 --- a/src/mcore_bridge/model/mm_gpts/glm5_next.py +++ b/src/mcore_bridge/model/mm_gpts/glm5_next.py @@ -194,14 +194,25 @@ def _get_packed_padding_mask(self, packed_seq_params, position_ids): cu = _get_physical_cu_seqlens(packed_seq_params) lengths = lengths.to(device=cu.device, dtype=cu.dtype) assert 0 < lengths.numel() <= cu.numel() - 1 - assert cu[-1] == position_ids.shape[-1] + cp_size = self.pg_collection.cp.size() + # Under CP the data pipeline shards position_ids/labels (zigzag) but passes cu_seqlens and + # seq_lens FULL, so build the mask over the full physical sequence and then shard it with the + # same zigzag, keeping it token-aligned with the local position_ids. At cp==1 the physical + # length is just the (unsharded) position count. + total = int(cu[-1]) + if cp_size == 1: + assert total == position_ids.shape[-1] # The remaining spans are alignment dummy sequences with zero logical length; they must not # update the expert bias. lengths = torch.nn.functional.pad(lengths, (0, cu.numel() - 1 - lengths.numel())) - positions = torch.arange(position_ids.shape[-1], device=cu.device) + positions = torch.arange(total, device=cu.device) segment = torch.bucketize(positions, cu[1:], right=True) mask = (positions - cu[segment] >= lengths[segment])[None] + if cp_size > 1: + from mcore_bridge.utils import split_cp_inputs + mask = split_cp_inputs(mask, cu, 1) if self.config.sequence_parallel: + assert mask.shape[1] % self.pg_collection.tp.size() == 0, f'padding_mask.shape: {mask.shape}' mask = mask.chunk(self.pg_collection.tp.size(), dim=1)[self.pg_collection.tp.rank()] return mask.contiguous() @@ -431,9 +442,6 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): from ..modules import TopKRouter config = self.config - if config.context_parallel_size > 1: - raise NotImplementedError('The current model has no KDA/DSA context parallelism; ' - 'use context_parallel_size=1') 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: @@ -442,6 +450,11 @@ 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.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'; " + "promoting to 'all_gather'.") + config.cp_comm_type = 'all_gather' config.hetereogenous_dist_checkpoint = True config.rope_type = MLATransformerConfig.rope_type config.rotary_scaling_factor = MLATransformerConfig.rotary_scaling_factor