Is this page helpful?

core.datasets.data_schedule#

Module Contents#

Classes#

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.

DefaultDynamicCPScheduler

Balance packed samples while selecting a runtime CP size per microbatch.

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_dynamic_cp_group

Select a pre-created runtime group from the caller’s parallel grid.

_sequence_parallel_tp_cp_group

Return the TP x CP group sequence parallelism needs, or None without it.

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.BasePackingScheduler(
max_seqlen_per_dp_cp_rank: int,
cp_size: int,
dp_size: int,
microbatch_group_size_per_vp_stage: Optional[int],
pad_sequences: bool = False,
max_num_seqs: Optional[int] = None,
)#

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)

  • pad_sequences – Account for per-sequence zigzag padding when packing.

  • max_num_seqs – Bound packed documents for fixed CUDA-graph metadata.

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)

  • pad_sequences – Account for per-sequence zigzag padding when packing.

  • max_num_seqs – Bound packed documents for fixed CUDA-graph metadata.

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.

Every PP stage owns the packed dataset on TP rank zero. Each stage runs the same schedule locally, then keeps only the data fields required by its pipeline/VPP position. This avoids a PP metadata broadcast and lets arbitrary PP layouts (including middle-stage MTP) consume the same per-microbatch runtime CP metadata.

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 real (unpadded) tokens for FLOPs calculation. seqlen_squared_sum_this_global_batch: Sum of squared real seqlens for FLOPs.

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’

DEFAULT_DYNAMIC_CP#

‘default_dynamic_cp’

class core.datasets.data_schedule.DefaultDynamicCPScheduler(
*args,
min_cp_size: int = 1,
cp_group_sizes: Sequence[int] | None = None,
**kwargs,
)#

Bases: core.datasets.data_schedule.DpBalancedScheduler

Balance packed samples while selecting a runtime CP size per microbatch.

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)

  • pad_sequences – Account for per-sequence zigzag padding when packing.

  • max_num_seqs – Bound packed documents for fixed CUDA-graph metadata.

get_groups_and_subsamples(sample_id_seqlens)#
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: megatron.core.process_groups_config.ProcessGroupCollection,
)#

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

  • pg_collection – Caller-owned process groups; must define dp_cp, dp, tp, and pp.

core.datasets.data_schedule._get_dynamic_cp_group(
group_size: int,
groups: Dict[int, torch.distributed.ProcessGroup],
name: str,
) → torch.distributed.ProcessGroup#

Select a pre-created runtime group from the caller’s parallel grid.

core.datasets.data_schedule._sequence_parallel_tp_cp_group(
config,
tp_group: torch.distributed.ProcessGroup,
dynamic_cp: bool,
local_cp_size: int | None,
dynamic_tp_cp_group_func: Optional[Callable[..., torch.distributed.ProcessGroup]],
tp_cp_group: Optional[torch.distributed.ProcessGroup],
) → Optional[torch.distributed.ProcessGroup]#

Return the TP x CP group sequence parallelism needs, or None without it.

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,
dynamic_cp: bool = False,
pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None,
config=None,
return_context_parallel_batch: bool = False,
dynamic_cp_group_func: Optional[Callable[..., torch.distributed.ProcessGroup]] = None,
dynamic_tp_cp_group_func: Optional[Callable[..., torch.distributed.ProcessGroup]] = None,
tp_cp_group: Optional[torch.distributed.ProcessGroup] = 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.

  • return_context_parallel_batch (bool) – Return layout-keyed batch views instead of the legacy GPT tuple.

  • dynamic_cp_group_func (Callable, optional) – Overrides the collection’s runtime CP group lookup for group_size ranks containing this rank.

  • dynamic_tp_cp_group_func (Callable, optional) – Returns the TP x runtime-DPxCP group of group_size ranks. Overrides pg_collection.dynamic_tp_cp.

  • tp_cp_group (torch.distributed.ProcessGroup, optional) – Static TP x CP group for sequence parallelism. Defaults to pg_collection.tp_cp.

  • pg_collection (ProcessGroupCollection) – Caller-owned parallel grid. Runtime groups default to its dynamic_cp and dynamic_tp_cp maps; callbacks may override them.

Returns:

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