Open cp failed when use packed data

[gaojingwei]

I sft the model using NeMo with data in packed format. It runs normally when cp=1, but when cp > 1, the following error occurs:

File "/opt/megatron-lm/megatron/core/pipeline_parallel/schedules.py", line 400, in forward_step
[rank4]:     output_tensor, loss_func = forward_step_func(data_iterator, model)
[rank4]:                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "/opt/NeMo/nemo/lightning/megatron_parallel.py", line 505, in wrapped_forward_step_func
[rank4]:     batch = _data_step(dataloader_iter)
[rank4]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "/opt/NeMo/nemo/collections/llm/gpt/model/base.py", line 667, in data_step
[rank4]:     return self.config.data_step_fn(dataloader_iter)
[rank4]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "/opt/NeMo/nemo/collections/llm/gpt/model/base.py", line 111, in gpt_data_step
[rank4]:     output = get_batch_on_this_cp_rank(_batch_required_keys)
[rank4]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]:   File "/opt/megatron-lm/megatron/core/utils.py", line 1838, in get_batch_on_this_cp_rank
[rank4]:     raise e
[rank4]:   File "/opt/megatron-lm/megatron/core/utils.py", line 1825, in get_batch_on_this_cp_rank
[rank4]:     val.shape[seq_dim] // (2 * cp_size),
[rank4]:     ~~~~~~~~~^^^^^^^^^
[rank4]: IndexError: tuple index out of range

I found the source code and added logging:

def get_batch_on_this_cp_rank(batch: Dict[str, Any]):
    """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.
    cp_size = parallel_state.get_context_parallel_world_size()
    if cp_size > 1:
        cp_rank = parallel_state.get_context_parallel_rank()
        for key, val in batch.items():
            if val is not None:
                seq_dim = 1 if key != "attention_mask" else 2

                try:
                    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.zeros(2, dtype=torch.int64, device=val.device)
                    index[0].fill_(cp_rank)
                    index[1].fill_(2 * cp_size - cp_rank - 1)
                    val = val.index_select(seq_dim, index)
                    val = val.view(*val.shape[0:seq_dim], -1, *val.shape[(seq_dim + 2) :])
                    batch[key] = val
                except Exception as e:
                    print(f"Rank {cp_rank}")
                    print(f"Key: {key}")
                    print(f"Shape: {val.shape} {seq_dim}  {val}")
                    raise e

    return batch

The logs show that:
Rank 0
Key: attention_mask
Shape: torch.Size([1]) 2 tensor([1], device='cuda:0')

my data like this，dataset is GPTSFTPackedDataset

({'tokens': tensor([[167543,   8948,    198,  ..., 167545, 167545, 167545]]), 'labels': tensor([[  8948,    198,   2610,  ..., 167545, 167545, 167545]]), 'loss_mask': tensor([[0, 0, 0,  ..., 0, 0, 0]]), 'position_ids': tensor([[0, 1, 2,  ..., 0, 0, 0]]), 'attention_mask': tensor([1]), 'cu_seqlens': tensor([[    0,  9484,  9665, 10350, 10643, 11403, 12009, 12691, 13199, 13686,
         14763, 15437, 16731, 17120, 18486, 19185, 20122, 21117, 21284, 21619,
         23190, 24816, 26374, 27024, 27591, 28392, 32342, 32581, 32724, 32768,
            -1]], dtype=torch.int32), 'cu_seqlens_argmin': tensor([[30]]), 'max_seqlen': tensor([[9484]], dtype=torch.int32), 'cu_seqlens_unpadded': tensor([[    0,  9484,  9665, 10350, 10643, 11403, 12009, 12691, 13199, 13686,
         14763, 15437, 16731, 17120, 18486, 19185, 20122, 21117, 21284, 21619,
         23190, 24816, 26374, 27024, 27591, 28392, 32342, 32581, 32724, 32724,
            -1]], dtype=torch.int32), 'cu_seqlens_unpadded_argmin': tensor([[30]]), 'token_count': [32724]}, 0, 0)

I tried setting cp=2 and cp=4, and encountered the same issue in both cases. I also tested it on both NeMo 2504 and NeMo 2509 container。

[akoumpa]

Hi, please note our LLM/VLM work has migrated to megatron bridge repo and nemo automodel

Can you try our latest infra to run CP with SFT?

[gaojingwei]

@akoumpa @chtruong814
I try megatron bridge in nemo 2511 image，it raise error in the same position at function get_batch_on_this_cp_rank when cp = 2，the error key is “cu_seqlens” and ”cu_seqlens_argmin“。

I previously enabled CP and successfully sft with packed data on NeMo version 24.12. I noticed that there are the following two differences between Nemo2412 and Nemo 2.0:

1> In nemo2412，in prepare_packed_for_chat_dataset.py, packed inputs and lables should be a multiple of (cp_size * 2)，but in nemo 2.0, this piece of code has disappeared.

if cp_size > 1:

        def pre_pad_dataset(data, max_seq_length, max_length_to_pad, pad_id):
            '''
            pad each individual data point to the length of max_length
            '''
            assert max_seq_length >= max_length_to_pad
            for key, val in data.items():
                if key in {'input_ids', 'context_ids', 'mask'}:
                    if len(val) <= max_length_to_pad:
                        # because input_ids are truncated by 1 for inputs and labels,
                        # we add 1 extra padding here to make sure padded inputs and labels
                        # are is a multiple of (cp_size * 2)
                        #val = val + [pad_id] * (max_length_to_pad - len(val) + 1)
                        
                        if 'mask' != key:
                            padding_tensor = torch.full((max_length_to_pad - len(val) + 1,), pad_id, dtype=val.dtype)
                        else:
                            padding_tensor = torch.full((max_length_to_pad - len(val) + 1,), False, dtype=val.dtype)

                        data[key] = torch.cat((val, padding_tensor), dim=0)

                    elif len(val) > max_seq_length:
                        logging.info(
                            f"""The current sequence length {len(val)} for packing is
                                        larger than the max_seq_length specified ({max_seq_length}).
                                        The current seqquence length is truncated to the size of max_seq_length.
                                        Please consider increase the sequence packing size"""
                        )
                        data[key] = val[:max_seq_length]

            return

        ceil_to_nearest = lambda n, m: (n + m - 1) // m * m
        for data in dataset:
            max_length_to_pad = min(max_seq_length, ceil_to_nearest(len(data['input_ids']), pad_seq_length_to_mult))
            pre_pad_dataset(data, max_seq_length, max_length_to_pad, pad_id)
 

2> in nemo 2412, in megatron_gpt_model.py，the get_batch_on_this_context_parallel_rank funcion like this，it do nothing when cu_seqlens in batch.

def get_batch_on_this_context_parallel_rank(self, batch):
        num_valid_tokens_in_ub = None
        if 'loss_mask' in batch and batch['loss_mask'] is not None:
            num_valid_tokens_in_ub = batch['loss_mask'].sum()

        cp_size = parallel_state.get_context_parallel_world_size()
        if cp_size > 1:
            cp_rank = parallel_state.get_context_parallel_rank()
            # check if the batch is not in THD format
            if 'cu_seqlens' not in batch:
                for key, val in batch.items():
                    if val is not None and key != "context_lengths":
                        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="cpu", pin_memory=True
                        ).cuda(non_blocking=True)
                        val = val.index_select(seq_dim, index)
                        val = val.view(*val.shape[0:seq_dim], -1, *val.shape[(seq_dim + 2) :])
                        batch[key] = val
        batch['num_valid_tokens_in_ub'] = num_valid_tokens_in_ub

        return batch

and fwd_output_and_loss_func like this，it handled cu_seqlens and cu_seqlens_argmin out of get_batch_on_this_context_parallel_rank.

def fwd_output_and_loss_func(dataloader_iter, model, checkpoint_activations_all_layers=None):

            # Get data batch
            batch = self.get_batch(dataloader_iter, tuning)

            # Transfer needed data to GPU
            required_keys = set()
            max_seqlen = batch['max_seqlen'].squeeze() if 'max_seqlen' in batch else None
            cu_seqlens_argmin = batch['cu_seqlens_argmin'] if 'cu_seqlens_argmin' in batch else None
            cu_seqlens_unpadded_argmin = (
                batch['cu_seqlens_unpadded_argmin'] if 'cu_seqlens_unpadded_argmin' in batch else None
            )
            if parallel_state.get_pipeline_model_parallel_world_size() == 1:
                required_keys.update(batch.keys())
            else:
                required_keys.add('attention_mask')
                if 'cu_seqlens' in batch:
                    required_keys.add('cu_seqlens')
                if 'cu_seqlens_unpadded' in batch:
                    required_keys.add('cu_seqlens_unpadded')
                if parallel_state.is_pipeline_first_stage():
                    required_keys.update(('tokens', 'position_ids'))
                if parallel_state.is_pipeline_last_stage():
                    required_keys.update(('labels', 'loss_mask'))
            if self.get_attention_mask_from_fusion and 'attention_mask' in required_keys:
                required_keys.remove('attention_mask')
            batch = {
                key: val.cuda(non_blocking=True) if key in required_keys and isinstance(val, torch.Tensor) else None
                for key, val in batch.items()
            }

            # slice batch along sequence dimension for context parallelism
            batch = self.get_batch_on_this_context_parallel_rank(batch)

            # Model forward pass
            forward_args = {
                'input_ids': batch['tokens'],
                'position_ids': batch['position_ids'],
                'attention_mask': None if self.get_attention_mask_from_fusion else batch['attention_mask'],
                'labels': batch['labels'] if 'labels' in batch else None,
                'loss_mask': batch['loss_mask'],
            }

            if not self.mcore_gpt:
                forward_args['checkpoint_activations_all_layers'] = checkpoint_activations_all_layers
                if not self.use_loss_mask:
                    forward_args.pop('loss_mask')
            else:
                # TODO: @eharper can we add this to mcore?
                forward_args.pop('loss_mask')

                if 'cu_seqlens' in batch:  # packed sequence from GPTSFTPackedDataset
                    # these args are passed eventually into TEDotProductAttention.forward()
                    cu_seqlens = batch['cu_seqlens'].squeeze()  # remove batch size dimension (mbs=1)
                    cu_seqlens_unpadded = batch['cu_seqlens_unpadded'].squeeze()
                    # remove -1 "paddings" added in collate_fn
                    if cu_seqlens_argmin is not None:
                        cu_seqlens = cu_seqlens[: cu_seqlens_argmin.item()]
                    else:
                        cu_seqlens = cu_seqlens[: torch.argmin(cu_seqlens)]
                    if cu_seqlens_unpadded_argmin is not None:
                        cu_seqlens_unpadded = cu_seqlens_unpadded[: cu_seqlens_unpadded_argmin.item()]
                    else:
                        cu_seqlens_unpadded = cu_seqlens_unpadded[: torch.argmin(cu_seqlens_unpadded)]
                    try:
                        from megatron.core.packed_seq_params import PackedSeqParams
                    except (ImportError, ModuleNotFoundError) as e:
                        mcore_version = packaging.version.Version(version('megatron-core'))
                        logging.error(
                            f"megatron-core v{mcore_version} does not support training with packed sequence. "
                            "Please use megatron-core >= 0.5.0, or set model.data.train_ds.packed_sequence=False"
                        )
                        raise e

                    # get packed sequences for this context parallel rank
                    cp_size = parallel_state.get_context_parallel_world_size()
                    if cp_size > 1:
                        try:
                            import transformer_engine_torch as tex
                        except ModuleNotFoundError as e:
                            logging.error(
                                "Please update Transformer Engine to >= 1.10 to use Context Parallel with THD format data"
                            )
                            raise e
                        cp_rank = parallel_state.get_context_parallel_rank()
                        for key in required_keys:
                            val = batch[key]
                            if key not in {
                                "cu_seqlens",
                                "cu_seqlens_unpadded",
                                "cu_seqlens_argmin",
                                "cu_seqlens_unpadded_argmin",
                                "max_seqlen",
                                "token_count",
                            }:
                                index = tex.thd_get_partitioned_indices(cu_seqlens, val.size(1), cp_size, cp_rank)
                                val = val.index_select(1, index)
                                batch[key] = val
                        forward_args = {
                            'input_ids': batch['tokens'],
                            'position_ids': batch['position_ids'],
                            'attention_mask': None if self.get_attention_mask_from_fusion else batch['attention_mask'],
                            'labels': batch['labels'] if 'labels' in batch else None,
                        }

                    forward_args['packed_seq_params'] = PackedSeqParams(
                        cu_seqlens_q=cu_seqlens_unpadded,
                        cu_seqlens_kv=cu_seqlens_unpadded,
                        cu_seqlens_q_padded=cu_seqlens,
                        cu_seqlens_kv_padded=cu_seqlens,
                        max_seqlen_q=max_seqlen,
                        max_seqlen_kv=max_seqlen,
                        qkv_format='thd',
                    )

            output_tensor = model(**forward_args)

3> In nemo 2412 , I also change the below code, at last it run successfull，my model is qwen2.5 structure。
I commented out all assignments to cu_seqlens_unpadded and directly assigned cu_seqlens to it in GPTSFTPackedDataset：：collate_fn()

 for item in batch:
            position_ids.append([])
            cu_seqlens.append([0])
            cu_seqlens_unpadded.append([0])
            seqlens = np.array(item['seq_boundaries'][1:]) - np.array(item['seq_boundaries'][:-1])
            for l in seqlens:
                # length minus 1 because input_ids is truncated by 1 for labels
                position_ids[-1].extend(list(range(l - 1)))
                cu_seqlens[-1].append(cu_seqlens[-1][-1] + l - 1)

            # the last seq needs to be the max seq len because rope and attn kernels expect no padding
            assert cu_seqlens[-1][-1] <= max_length

            # since data is prepadded when cp_size > 1, there may be some extra padding at the end
            # of the packed sequence. In this case, we need to add the max seq len to the end.
            if cu_seqlens[-1][-1] != max_length:
                cu_seqlens[-1].append(max_length)
            
            # add by gaojingwei            
            cu_seqlens_unpadded = copy.deepcopy(cu_seqlens) 
             
            '''
            for i in range(len(item['seq_boundaries']) - 1):
                current_seq = item['input_ids'][item['seq_boundaries'][i] : item['seq_boundaries'][i + 1] - 1]

                # since the data could be prepadded with tokenizer's eos_id, we can find out the index of all the eos_id
                eos_idx = np.where(np.array(current_seq) == self.custom_eos_id)
                #print("eos_idx: ", eos_idx)

                # The second eos_id index marks the length of the original unpadded sequence if the sequence is
                # prepadded for cp_size > 1. Otherwise, there is no extra padding.
                seqlen_unpadded = eos_idx[0][0] + 1 if eos_idx[0].any() else len(current_seq)
                cu_seqlens_unpadded[-1].append(cu_seqlens_unpadded[-1][-1] + seqlen_unpadded)
            
            # if extra paddings are added in the packed sequence, they can't be counted as
            # actual tokens for training
            #print("len cu_seqlens[-1] cu_seqlens_unpadded[-1]", len(cu_seqlens[-1]), len(cu_seqlens_unpadded[-1]))
            #print("cu_seqlens[-1]: ", cu_seqlens[-1])
            #print("cu_seqlens_unpadded[-1]: ", cu_seqlens_unpadded[-1])

            if len(cu_seqlens[-1]) > len(cu_seqlens_unpadded[-1]):
                cu_seqlens_unpadded[-1].append(cu_seqlens_unpadded[-1][-1])
                print("1 cu_seqlens_unpadded[-1]: ", cu_seqlens_unpadded[-1])
            else:
                print("2 cu_seqlens_unpadded[-1]: ", cu_seqlens_unpadded[-1])
            '''

In Nemo 2509，I made the above three changes, and the program now runs successfully with cp > 1。

[gaojingwei]

New progress: Following the slice_batch_for_context_parallel function in src/megatron/bridge/utils/common_utils.py in Megatron Bridge , I have modified the get_batch function in src/megatron/bridge/training/gpt_step.py file accordingly。It’s no longer throwing errors now, but the loss looks incorrect. When I run with 4 GPUs using TP=4, PP=1, CP=1 versus TP=2, PP=1, CP=2, the losses are different. Here's how I modified it:

def get_batch(
    data_iterator: Iterable, cfg: ConfigContainer, use_mtp: bool = False
) -> tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
]:
   
    if (not parallel_state.is_pipeline_first_stage()) and (not parallel_state.is_pipeline_last_stage()):
        return None, None, None, None, None, None, None, None

    batch = get_batch_from_iterator(
        data_iterator,
        use_mtp,
        getattr(cfg.dataset, "skip_getting_attention_mask_from_dataset", True),
    )

    # slice batch along sequence dimension for context parallelism
    cp_size = parallel_state.get_context_parallel_world_size()
    
    if "cu_seqlens" in batch and  cp_size > 1:
        import transformer_engine_torch as tex

        cp_rank = parallel_state.get_context_parallel_rank()

        cu_seqlens = batch["cu_seqlens"]
        cu_seqlens_padded = cu_seqlens
        cu_seqlens_padded = cu_seqlens_padded.squeeze()
        
        for key, val in batch.items():
            if key not in {
                    "cu_seqlens",
                    "cu_seqlens_unpadded",
                    "cu_seqlens_argmin",
                    "cu_seqlens_unpadded_argmin",
                    "max_seqlen",
                    "token_count",
                } and val != None:

                index = tex.thd_get_partitioned_indices(cu_seqlens_padded, val.size(1), cp_size, cp_rank)
                val = val.index_select(1, index)
                batch[key] = val
    else:
        batch = get_batch_on_this_cp_rank(batch)

    return (
        batch["tokens"],
        batch["labels"],
        batch["loss_mask"],
        batch["attention_mask"],
        batch["position_ids"],
        batch.get("cu_seqlens"),
        batch.get("cu_seqlens_argmin"),
        batch.get("max_seqlen"),
    )

I noticed that the PackedSeqParams in Megatron Bridge is missing two parameters compared to NeMo: cu_seqlens_q_padded and cu_seqlens_kv_padded. I’m not sure if this is the reason causing the loss discrepancy between runs with and without context parallelism (CP).
With the modification I described in my previous reply, the loss is consistent between CP-enabled and CP-disabled configurations in NeMo 25.09.