`core.datasets.data_schedule`#

Module Contents#

Classes#

`HybridCPDataLoaderWrapper`	A wrapper class that wraps around an existing data_iterator. For every next call,
`BasePackingScheduler`	Base class for sequence packing schedulers.
`DpBalancedScheduler`	Packs sequences in their original order until reaching the max limit of sequence length.
`PackingSchedulerEnum`	Enum for supported sequence packing algorithms.

Functions#

`_build_thd_padding_mask`	Build a 1D THD padding mask from scheduler sequence metadata.
`_sanitize_thd_padding_values`	Replace padded token-like slots with safe neutral values in-place.
`wrap_data_iterator`	A wrapper function that wraps around an existing data_iterator and return the num_micro_batches for sequence packing.
`get_batch_on_this_rank_for_sequence_packing`	Get a batch of data for sequence packing.

Data#

scheduler_map

API#

core.datasets.data_schedule._build_thd_padding_mask( cu_seqlens: torch.Tensor, cu_seqlens_padded: torch.Tensor, ) → torch.Tensor#: Build a 1D THD padding mask from scheduler sequence metadata.

core.datasets.data_schedule._sanitize_thd_padding_values( batch: Dict[str, Any], padding_mask: torch.Tensor, ) → None#: Replace padded token-like slots with safe neutral values in-place.

class core.datasets.data_schedule.HybridCPDataLoaderWrapper( data_iterator, config, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, )#

A wrapper class that wraps around an existing data_iterator. For every next call,

Each DP rank pulls a batch of packed samples.
Extracts the sequence lengths of each sub-sample and all-gathers across the DP group.
Schedules the sub-samples to the DPxCP ranks using the BalancedCPScheduler.
Based on the schedule, reroutes the sub-samples to the correct rank using all-to-all.
Returns the assigned sub-samples to this rank.

Parameters:

data_iterator – The original data_iterator to wrap around
config – The config object containing the max_seqlen_per_dp_cp_rank
dp_cp_group – Data parallel context parallel group.

Initialization

__iter__()#: Return self as an iterator.

get_global_seqlens( subsample_seqlens: torch.Tensor, ) → List[int]#

Gathers the sequence lengths of all subsamples from all DP ranks. Each DP rank loads the same number of microbatches but each microbatch may have a different number of subsamples. We find the number of subsamples each rank holds and then gather the sequence lengths of all subsamples from all ranks.

Delegates to data_schedule_utils._get_global_seqlens_and_ids. The shared helper returns offsets with one extra trailing entry (the total subsample count); this method preserves the original contract of returning per-rank start offsets only.

get_global_id_seqlens(num_local_subsamples, offsets, seqlens_gathered)#

Calculates the global ID for each subsample.

We assign a unique global ID to each subsample.

Kept as a local implementation: the shared helper fuses this pure indexing step with the collective gather, so delegating here would re-run the all-gathers.

Returns: global_id_seqlens: list of (global_id, seqlen) tuples for scheduling. global_ids_this_rank: list of global IDs locally present on this rank.

_gid_to_src_rank(gid: int, offsets: List[int]) → int#

reroute_samples_to_hdp_ranks( batch, global_ids_this_rank, global_id_seqlens, sample_id_groups, offsets, )#

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. Since all CP ranks within a DP group have the same data, we only need to transfer data between matching CP ranks.

unpack_batch(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.

__next__() → Any#: Get the next item from the dataset, pull scheduling metadata and return it.

class core.datasets.data_schedule.BasePackingScheduler( max_seqlen_per_dp_cp_rank: int, cp_size: int, dp_size: int, microbatch_group_size_per_vp_stage: Optional[int], )#

Base class for sequence packing schedulers.

Initialization

Parameters:

max_seqlen_per_dp_cp_rank – The maximum sequence length per DPxCP rank.
cp_size – The context parallel size.
dp_size – The data parallel size.
microbatch_group_size_per_vp_stage – The microbatch group size per virtual
stage (pipeline)
VPP (only used when enabling)
None. (otherwise)

abstractmethod get_required_sample_keys()#: Return the required key of each batch.

abstractmethod get_groups_and_subsamples(sample_id_seqlens)#: schedule the samples into groups

abstractmethod run( data_iterator, num_microbatches, dp_group, tp_group, pp_group, dp_cp_group, dev, config, )#

Run the scheduler and return the new data_iterator.

Parameters:

data_iterator – The data iterator.
num_microbatches – The number of microbatches to fetch.
dp_group – Data parallel process group.
tp_group – Tensor parallel process group.
pp_group – Pipeline parallel process group.
dp_cp_group – Data parallel + context parallel process group.
dev – CUDA device.
config – Model parallel config.

Returns:

The new data iterator (or list for VPP). num_micro_batches: Number of micro batches after scheduling. seqlen_sum_this_global_batch: Total tokens for FLOPs calculation. seqlen_squared_sum_this_global_batch: Sum of squared seqlens for FLOPs.

Return type:

new_data_iterator

class core.datasets.data_schedule.DpBalancedScheduler(*args, **kwargs)#

Bases: core.datasets.data_schedule.BasePackingScheduler

Packs sequences in their original order until reaching the max limit of sequence length.

Initialization

Parameters:

max_seqlen_per_dp_cp_rank – The maximum sequence length per DPxCP rank.
cp_size – The context parallel size.
dp_size – The data parallel size.
microbatch_group_size_per_vp_stage – The microbatch group size per virtual
stage (pipeline)
VPP (only used when enabling)
None. (otherwise)

get_required_sample_keys()#: Return the required key of each batch.

get_groups_and_subsamples(sample_id_seqlens)#: Packs sequences in their original order until reaching the max limit of sequence length.

run( data_iterator, num_microbatches: int, dp_group, tp_group, pp_group, dp_cp_group, dev: torch.device, config, )#

Run the complete scheduling pipeline.

Steps: 1. Fetch batches and gather global sequence lengths 2. Check required sample keys 3. Schedule samples into groups 4. Reroute samples to DCP ranks 5. Build packed microbatches 6. Calculate FLOPs info 7. Broadcast to PP group (for middle PP stages) 8. Broadcast to TP group (for non-TP-0 ranks) 9. Handle VPP if enabled

Parameters:

data_iterator – The data iterator.
num_microbatches – The number of microbatches to fetch.
dp_group – Data parallel process group.
tp_group – Tensor parallel process group.
pp_group – Pipeline parallel process group.
dp_cp_group – Data parallel + context parallel process group.
dev – CUDA device.
config – Model parallel config.

Returns:

Return type:

new_data_iterator

class core.datasets.data_schedule.PackingSchedulerEnum(*args, **kwds)#

Bases: enum.Enum

Enum for supported sequence packing algorithms.

Initialization

DP_BALANCED#: ‘dp_balanced’

core.datasets.data_schedule.scheduler_map: Dict[core.datasets.data_schedule.PackingSchedulerEnum, Type[core.datasets.data_schedule.BasePackingScheduler]]#: None

core.datasets.data_schedule.wrap_data_iterator( data_iterator, config, num_microbatches, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, )#

A wrapper function that wraps around an existing data_iterator and return the num_micro_batches for sequence packing.

Parameters:

data_iterator – The original data_iterator to wrap around
config – The config object containing the max_seqlen_per_dp_cp_rank
dp_cp_group – Data parallel context parallel group.
pg_collection – The process group collection.

core.datasets.data_schedule.get_batch_on_this_rank_for_sequence_packing( data_iterator, vpp_size: Optional[int] = None, mtp_on_this_rank: bool = False, vp_stage: Optional[int] = None, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, )#

Get a batch of data for sequence packing.

Parameters:

data_iterator (Iterator) – The data iterator to get the batch from.
mtp_on_this_rank (bool) – Whether to use multi-token prediction.
vp_stage (Optional[int]) – The stage of the pipeline.

Returns:

tuple of (tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask)

core.datasets.data_schedule#

Module Contents#

Classes#

Functions#

Data#

API#

`core.datasets.data_schedule`#