update mm_gpt_model input_ids (#117)

Author: Jintao-Huang

src/mcore_bridge/model/gpt_model.py

@@ -451,9 +451,6 @@ def _postprocess(self,
         output_weight = None
         if self.share_embeddings_and_output_weights:
             output_weight = self.shared_embedding_or_output_weight()
-        if self.config.is_multimodal and self.config.context_parallel_size > 1 and input_ids is not None:
-            # input_ids is required by MTP.
-            input_ids = split_cp_inputs(input_ids, getattr(packed_seq_params, 'cu_seqlens_q', None), 1)
 
         if self.mtp_process and labels is not None:
             hidden_states = self.mtp(

src/mcore_bridge/model/mm_gpt_model.py

@@ -8,7 +8,7 @@
 from megatron.core.transformer.spec_utils import ModuleSpec
 
 from mcore_bridge.config import ModelConfig
-from mcore_bridge.utils import split_cp_inputs
+from mcore_bridge.utils import reconstruct_tensor_cp, split_cp_inputs
 
 from .gpt_model import GPTModel
 
@@ -83,18 +83,24 @@ def forward(
         **kwargs,
     ) -> torch.Tensor:
         extra_kwargs = {k: kwargs[k] for k in self.language_model.extra_forward_keys}
+        # Compatible with legacy mcore-bridge behavior.
+        cp_size = self.config.context_parallel_size
+        needs_split = cp_size > 1 and input_ids is not None and position_ids.shape[-1] * cp_size == input_ids.shape[-1]
         if decoder_input is not None:
             pass
         elif self.pre_process:
-            kwargs.update({'input_ids': input_ids, 'packed_seq_params': packed_seq_params})
+            input_ids_ = input_ids if needs_split else reconstruct_tensor_cp(input_ids, packed_seq_params, dim=1)
+            kwargs.update({'input_ids': input_ids_, 'packed_seq_params': packed_seq_params})
             with self._patch_word_embeddings(kwargs):
-                decoder_input = self.language_model.embedding(input_ids=input_ids, position_ids=position_ids)
+                decoder_input = self.language_model.embedding(input_ids=input_ids_, position_ids=position_ids)
         else:
             # intermediate stage of pipeline
             # decoder will get hidden_states from encoder.input_tensor
             decoder_input = None
             kwargs = {}
         kwargs.update(extra_kwargs)
+        if needs_split:
+            input_ids = split_cp_inputs(input_ids, getattr(packed_seq_params, 'cu_seqlens_q', None), dim=1)
         return self.language_model(
             input_ids=input_ids,
             position_ids=position_ids,

src/mcore_bridge/utils/__init__.py

@@ -3,7 +3,8 @@
 from .env import get_dist_setting, get_node_setting, is_dist, is_last_rank, is_local_master, is_master
 from .import_utils import _LazyModule, is_flash_attn_3_available
 from .logger import get_logger
-from .megatron_utils import get_local_layer_specs, roll_tensor, set_random_seed, split_cp_inputs, unwrap_model
+from .megatron_utils import (get_local_layer_specs, reconstruct_tensor_cp, roll_tensor, set_random_seed,
+                             split_cp_inputs, unwrap_model)
 from .safetensors import SafetensorLazyLoader, StreamingSafetensorSaver
 from .torch_utils import gc_collect, get_current_device, safe_ddp_context, to_device
 from .utils import deep_getattr, get_env_args, json_parse_to_dict, patch_deepcopy

src/mcore_bridge/utils/megatron_utils.py

@@ -4,6 +4,7 @@
 import torch
 from megatron.core import mpu, tensor_parallel
 from megatron.core.distributed import DistributedDataParallel as DDP
+from megatron.core.ssm.mamba_context_parallel import _undo_attention_load_balancing
 from megatron.core.transformer.module import Float16Module
 from megatron.core.transformer.multi_token_prediction import roll_tensor as mcore_roll_tensor
 from megatron.core.transformer.transformer_block import get_num_layers_to_build
@@ -67,6 +68,44 @@ def split_cp_inputs(inputs: torch.Tensor, cu_seqlens: Optional[torch.Tensor], di
     return torch.cat(new_inputs, dim=dim)
 
 
+def reconstruct_tensor_cp(tensor, packed_seq_params, dim: int) -> torch.Tensor:
+    """In CP mode, all-gather and undo the load-balanced (zigzag) chunking
+    produced by ``split_cp_inputs``, restoring the full sequence in original
+    token order along ``dim``.
+
+    Args:
+        tensor: CP-sharded local tensor whose sequence dim is at ``dim``.
+        packed_seq_params: ``PackedSeqParams`` for THD inputs, or ``None`` for
+            regular ``[B, S, ...]`` inputs.
+        dim: Sequence dimension index of ``tensor`` (default: 1).
+
+    Returns:
+        torch.Tensor: Full-sequence tensor with the same shape as ``tensor``
+        except the size at ``dim`` is multiplied by ``cp_size``.
+    """
+
+    cp_size = mpu.get_context_parallel_world_size()
+    if cp_size <= 1:
+        return tensor
+
+    cp_rank = mpu.get_context_parallel_rank()
+    cp_group = mpu.get_context_parallel_group()
+
+    # All-gather across CP ranks (preserve local autograd graph for `tensor`).
+    output_list = [torch.empty_like(tensor) for _ in range(cp_size)]
+    torch.distributed.all_gather(output_list, tensor.contiguous(), group=cp_group)
+    output_list[cp_rank] = tensor
+    gathered = torch.cat(output_list, dim=dim)
+
+    # `_undo_attention_load_balancing` assumes sequence dim is 0; transpose if needed.
+    if dim != 0:
+        gathered = gathered.transpose(0, dim).contiguous()
+    out = _undo_attention_load_balancing(gathered, cp_size, packed_seq_params)
+    if dim != 0:
+        out = out.transpose(0, dim).contiguous()
+    return out
+
+
 def get_local_layer_specs(config, layer_specs, vp_stage=None):
     num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage)