bridge.data.packing.global_batch

Data-side contract for global-batch online packing.

Offline, in-batch, and Energon packing form bins from a local candidate pool (the dataset, one microbatch, or one worker’s buffer). Global-batch packing forms them once per training step inside Megatron-Core’s sequence-packing scheduler, after a data-parallel all-gather of sample lengths, so the candidate pool is the whole global batch across DP x CP ranks. That pool is also what lets dynamic context parallelism size a CP group for each packed bin.

The scheduler consumes unpacked per-sample dicts, one sequence per sample, delivered as a list rather than stacked into a batch. GPT-SFT datasets produce them when enable_global_batch_packing is set; this module holds that sample contract and the helper that builds one sample.

Module Contents

Functions

fold_alignment_padding	Report the padded length as the sequence length so bins carry no padding between sequences.
build_unpacked_sequence_sample	Build one scheduler sample from an unpadded sequence.

Data

REQUIRED_SAMPLE_KEYS

Keys every unpacked sample must carry for the Megatron-Core scheduler.

API

bridge.data.packing.global_batch.REQUIRED_SAMPLE_KEYS: tuple[str, ...]

(‘tokens’, ‘labels’, ‘loss_mask’, ‘position_ids’, ‘original_seq_len’, ‘padded_seq_len’)

Keys every unpacked sample must carry for the Megatron-Core scheduler.

tokens/labels/position_ids are int64 [L], loss_mask is float32 [L], and the two lengths are int32 [1] tensors (L is the padded length).

bridge.data.packing.global_batch.fold_alignment_padding(

sample: dict[str, torch.Tensor],

) → dict[str, torch.Tensor]

Report the padded length as the sequence length so bins carry no padding between sequences.

Set it when attention runs on FlashAttention under context parallelism: FlashAttention rejects THD bins with padding between sequences, and Transformer Engine picks it for training whenever cuDNN fused attention is unavailable, for example at head dimension 256 on Blackwell (Qwen3.5). Leave it off otherwise: cuDNN fused attention skips the padding natively. The padding stays loss-masked and sits after each sequence’s last token, so causal attention and recurrent layers compute the same outputs for real tokens; only the boundaries reported to the attention backend change. Megatron-Core must also derive pad_between_seqs from the actual cu_seqlens (https://github.com/NVIDIA/Megatron-LM/pull/7417).

bridge.data.packing.global_batch.build_unpacked_sequence_sample(

tokens: collections.abc.Sequence[int] | torch.Tensor,

labels: collections.abc.Sequence[int] | torch.Tensor,

loss_mask: collections.abc.Sequence[int] | collections.abc.Sequence[float] | torch.Tensor,

*,

pad_to_multiple_of: int,

pad_token_id: int,

ignore_label_id: int | None = None,

) → dict[str, torch.Tensor]

Build one scheduler sample from an unpadded sequence.

Parameters:

• tokens – Input token ids for one sequence.

• labels – Next-token targets aligned with tokens.

• loss_mask – 1 where labels are supervised, 0 elsewhere.

• pad_to_multiple_of – Alignment multiple for CP THD slicing.

• pad_token_id – Token written into alignment padding.

• ignore_label_id – Label written into alignment padding (defaults to pad_token_id).

Returns:

A dict with :data:REQUIRED_SAMPLE_KEYS; padding positions have loss_mask == 0.