core.datasets.data_schedule_utils

Module Contents

Functions

_unpack_batch	Unpacks the packed samples into a list of sub-samples. Since each sub-sample may be routed to different DPxCP ranks, we unpack the sample here to avoid unnecessarily transferring the entire packed sample.
_get_global_seqlens_and_ids	Gathers the sequence lengths of all subsamples from all DP ranks and calculates global IDs.
_pack_sequences	Pack multiple samples into a single packed sample.
broadcast_tensor	Broadcast a tensor from src_rank to all ranks in the group.
broadcast_to_pp_group	Broadcast num_micro_batches, seqlen_sum_this_global_batch, seqlen_squared_sum_this_global_batch and metadata to middle PP stages. Before this broadcast, the new_samples on middle PP stages are None, after this broadcast, the new_samples on middle PP stages contain the metadata but without tokens, labels, loss_mask, position_ids.
broadcast_scalars	Broadcast scalar values from rank 0 to all ranks in the group.
create_data_iterator	Handle virtual pipeline parallelism.
reroute_samples_to_dcp_ranks	Reroutes the sub-samples to the correct rank after scheduling.
build_packed_microbatches	Build packed samples for each microbatch.
get_batch_and_global_seqlens	Get the batch and global sequence lengths. Each DP rank loads the same number of sequences, so we need to gather the sequence lengths from all ranks then we can schedule the sequences into groups.

API

core.datasets.data_schedule_utils._unpack_batch(

batch: List[Dict[str, torch.Tensor]],

) → List[Dict[str, torch.Tensor]]

Unpacks the packed samples into a list of sub-samples. Since each sub-sample may be routed to different DPxCP ranks, we unpack the sample here to avoid unnecessarily transferring the entire packed sample.

Two mutually exclusive input shapes are accepted, and every sample in batch must use the same one:

• Pre-packed (e.g. :class:SFTDataset): each sample carries a cu_seqlens tensor and the tokens of multiple sub-samples concatenated together. We slice them apart and synthesize original_seq_len / padded_seq_len from the cu_seqlens deltas.

• Already unpacked (e.g. :class:VarlenDataset): each sample is a single sub-sample that already carries padded_seq_len (and usually original_seq_len). We just normalize the leading batch dimension introduced by the default collate_fn and return as-is.

The shape is decided once for the whole batch and asserted per sample, so a dataset that emits both keys cannot silently bypass the cu_seqlens slicing below.

core.datasets.data_schedule_utils._get_global_seqlens_and_ids(subsample_seqlens: torch.Tensor, dp_group)

Gathers the sequence lengths of all subsamples from all DP ranks and calculates global IDs.

core.datasets.data_schedule_utils._pack_sequences(

samples: List,

padded_lengths: torch.Tensor,

original_lengths: torch.Tensor,

dev: torch.device,

) → Dict[str, torch.Tensor]

Pack multiple samples into a single packed sample.

core.datasets.data_schedule_utils.broadcast_tensor(item, src_rank, group) → None

Broadcast a tensor from src_rank to all ranks in the group.

core.datasets.data_schedule_utils.broadcast_to_pp_group(

new_samples,

num_micro_batches,

seqlen_sum_this_global_batch,

seqlen_squared_sum_this_global_batch,

pp_group,

dev,

)

Broadcast num_micro_batches, seqlen_sum_this_global_batch, seqlen_squared_sum_this_global_batch and metadata to middle PP stages. Before this broadcast, the new_samples on middle PP stages are None, after this broadcast, the new_samples on middle PP stages contain the metadata but without tokens, labels, loss_mask, position_ids.

Who needs what:

• PP rank 0 and the last PP rank both own a data iterator (only TP rank 0 on the first and last PP stage does), so both run the whole schedule -> reroute -> pack pipeline on the same input samples and independently end up with complete new_samples: tokens, labels, loss_mask, position_ids and the packing metadata. Neither takes anything from this broadcast; the last stage in particular must keep its own labels / loss_mask.

• Middle PP stages have no data iterator, so new_samples is None on entry. They only need the packing metadata (max_seqlen / cu_seqlens / cu_seqlens_padded) to rebuild the packed-sequence params, never the token tensors.

The last PP rank still takes part in the transfer because torch.distributed.broadcast is a collective over pp_group: every member has to call it or the group deadlocks. It therefore receives the payload and drops it, which is what the pp_group.rank() != pp_group.size() - 1 guard below implements. Filtering it out of the transfer itself would require a separate “first + middle” process group, which is not worth an extra process group for a payload of a few hundred bytes per global batch.

core.datasets.data_schedule_utils.broadcast_scalars(

values: List,

group,

dev,

dtype=torch.float32,

) → List

Broadcast scalar values from rank 0 to all ranks in the group.

Parameters:

• values – List of scalar values to broadcast (only used on rank 0).

• group – The process group to broadcast within.

• dev – The device to use for the tensor.

• dtype – The data type for the tensor.

Returns:

List of broadcasted values.

core.datasets.data_schedule_utils.create_data_iterator(new_samples, pp_group, tp_group, config)

Handle virtual pipeline parallelism.

core.datasets.data_schedule_utils.reroute_samples_to_dcp_ranks(

batch,

global_ids_this_rank,

global_id_seqlens,

sample_id_groups,

offsets,

dp_group,

tp_group,

dp_cp_group,

total_dcp_gpus,

)

Reroutes the sub-samples to the correct rank after scheduling.

For each key in the batch dict, we perform an all-to-all communication to transfer the data to the correct ranks.

core.datasets.data_schedule_utils.build_packed_microbatches(

grouped_samples: List[List[Dict[str, torch.Tensor]]],

dev: torch.device,

) → List[Dict[str, torch.Tensor]]

Build packed samples for each microbatch.

core.datasets.data_schedule_utils.get_batch_and_global_seqlens(

data_iterator,

num_microbatches,

dp_group,

)

Get the batch and global sequence lengths. Each DP rank loads the same number of sequences, so we need to gather the sequence lengths from all ranks then we can schedule the sequences into groups.

Parameters:

• data_iterator – The data iterator.

• num_microbatches – The number of microbatches.

• dp_group – The data parallel group.

Returns:

The sub-samples pulled from this rank’s data_iterator over num_microbatches steps, flattened and unpacked (see :func:_unpack_batch). Every dict carries tokens / labels / loss_mask / position_ids plus the original_seq_len and padded_seq_len scalars used for scheduling. global_id_seqlens (List[Tuple[int, int]]): (global_id, padded_seq_len) for every sub-sample in the DP group, ordered by DP rank and then by local index. Identical on all ranks; this is the scheduler’s input. global_ids_this_rank (torch.Tensor): int32 CUDA tensor holding the global IDs of the sub-samples loaded by this rank, i.e. batch[i] has global ID global_ids_this_rank[i]. offsets (torch.Tensor): int32 CPU tensor of shape [dp_size + 1] with the exclusive prefix sum of the per-rank sub-sample counts, so DP rank r owns global IDs offsets[r]:offsets[r + 1]. Used by :func:reroute_samples_to_dcp_ranks to map a global ID back to its source rank. seqlens_gathered (List[int]): Padded sequence length of every sub-sample in the DP group, indexed by global ID (seqlens_gathered[gid] equals global_id_seqlens[gid][1]). Handy for global-batch token counts such as the FLOPs accounting.

Return type:

batch (List[Dict[str, torch.Tensor]])