code: def get_batch_on_this_cp_rank(batch):
      """Slice batch input along sequence dimension into multiple chunks,
      which are parallelized across GPUs in a context parallel group.
      """
  
      # With causal masking, each token only attends to its prior tokens. Simply split
      # sequence into CP chunks can result in severe load imbalance. That's to say, chunks
      # at the end of sequence have bigger workload than others. To address this issue,
      # we split sequence into 2*CP ranks. Assuming CP=2, we then get 4 chunks, chunk_0
      # and chunk_3 are assigned to GPU0, chunk_1 and chunk_2 are assigned to GPU1, so
      # that we can get balanced workload among GPUs in a context parallel group.
      args = get_args()
      cp_size = args.context_parallel_size
      if cp_size > 1:
          cp_rank = mpu.get_context_parallel_rank()
          for key, val in batch.items():
              if val is not None:
                  seq_dim = 1 if key != "attention_mask" else 2
                  val = val.view(
                      *val.shape[0:seq_dim],
                      2 * cp_size,
                      val.shape[seq_dim] // (2 * cp_size),
                      *val.shape[(seq_dim + 1) :],
                  )
                  index = torch.tensor(
                      [cp_rank, (2 * cp_size - cp_rank - 1)],
                      device=val.device,
                  )
                  val = val.index_select(seq_dim, index)
                  val = val.view(
                      *val.shape[0:seq_dim],
                      -1,
                      *val.shape[(seq_dim + 2) :],
                  )
                  batch[key] = val
  
      return batch
  
file_path: megatron-lm/megatron/utils.py
imports: import sys
  import torch
  from megatron import get_args, get_adlr_autoresume
  from megatron.core import DistributedDataParallel
  from megatron.core import mpu
  from megatron.core.tensor_parallel import param_is_not_tensor_parallel_duplicate
  from megatron.model import Float16Module
  from megatron.model.module import param_is_not_shared
label_desc: Use image node_id 8 for calling a global function with example usage: get_batch_on_this_cp_rank(batch) and returns: batch