bridge.training.global_batch_packing

Training-loop glue for global-batch online packing and dynamic context parallelism.

Megatron-Core’s sequence-packing scheduler (model.sequence_packing_scheduler, driven by dataset.enable_global_batch_packing) packs unpacked variable-length samples into THD microbatches once per training step, with the whole global batch across data-parallel ranks as its candidate pool. With model.dynamic_context_parallel it also runs every packed bin on a context-parallel group sized for its longest sequence (Megatron-Core dev).

func:

probe_global_batch_packing_support checks the pinned Megatron-Core before training starts, and the entry points below adapt to the scheduler API of the pinned release.

Configuration rules live in ConfigContainer.validate next to the other packing modes; the data contract lives in megatron.bridge.data.packing.global_batch.

Module Contents

Functions

global_batch_packing_enabled	Return whether a Megatron-Core online packing scheduler is configured on the model.
dynamic_context_parallel_enabled	Return whether dynamic context parallelism is configured (an explicit True, not any truthy value).
_scheduler_api	Import the Megatron-Core scheduler entry points.
_accepted_parameters	Return the parameter names fn accepts (cached: called once per microbatch).
_thd_cp_route_prebuilder	Return Megatron-Core’s THD CP partition-route prebuild, or None if the release has no CP layout routes.
probe_global_batch_packing_support	Return why the pinned Megatron-Core cannot run scheduler, or None if it can.
_iterator_for_tp_rank	Return the iterator on TP rank 0 and None elsewhere, as the scheduler requires.
wrap_data_iterator_for_global_batch_packing	Pack this step’s global batch into THD microbatches.
get_batch_for_global_batch_packing	Fetch one packed microbatch and its PackedSeqParams on every rank.
finalize_packed_seq_params	Bind the CP group to one microbatch and prebuild its THD CP partition route.

Data

logger
__all__
DYNAMIC_CP_SCHEDULERS	Scheduler names that form dynamic context-parallel groups per microbatch.
_DEV_PIN_HINT

API

bridge.training.global_batch_packing.logger

‘getLogger(…)’

bridge.training.global_batch_packing.__all__

[‘DYNAMIC_CP_SCHEDULERS’, ‘dynamic_context_parallel_enabled’, ‘finalize_packed_seq_params’, ‘get_bat…

bridge.training.global_batch_packing.DYNAMIC_CP_SCHEDULERS: tuple[str, ...]

(‘default_dynamic_cp’,)

Scheduler names that form dynamic context-parallel groups per microbatch.

bridge.training.global_batch_packing._DEV_PIN_HINT

‘ Dynamic context parallelism needs the Megatron-Core dev pin: `./scripts/switch_mcore.sh dev && uv s…’

bridge.training.global_batch_packing.global_batch_packing_enabled(model_config: Any) → bool

Return whether a Megatron-Core online packing scheduler is configured on the model.

Checks for a scheduler name rather than any non-None value so that stand-in configs (mocks, partial namespaces) do not route a step through the packing scheduler.

bridge.training.global_batch_packing.dynamic_context_parallel_enabled(model_config: Any) → bool

Return whether dynamic context parallelism is configured (an explicit True, not any truthy value).

bridge.training.global_batch_packing._scheduler_api() → tuple[Callable[..., Any], Callable[..., Any]]

Import the Megatron-Core scheduler entry points.

bridge.training.global_batch_packing._accepted_parameters(

fn: Callable[..., Any],

) → frozenset[str]

Return the parameter names fn accepts (cached: called once per microbatch).

bridge.training.global_batch_packing._thd_cp_route_prebuilder() → Callable[..., None] | None

Return Megatron-Core’s THD CP partition-route prebuild, or None if the release has no CP layout routes.

bridge.training.global_batch_packing.probe_global_batch_packing_support(

scheduler: str,

*,

dynamic_cp: bool = False,

) → str | None

Return why the pinned Megatron-Core cannot run scheduler, or None if it can.

Probing the registered schedulers and the entry-point signatures lets validation fail with an actionable message instead of a KeyError or TypeError at the first step. dynamic_cp additionally requires the packed batch fetch to accept the dynamic-CP arguments.

bridge.training.global_batch_packing._iterator_for_tp_rank(

data_iterator: Any,

pg_collection: megatron.core.process_groups_config.ProcessGroupCollection,

) → Any

Return the iterator on TP rank 0 and None elsewhere, as the scheduler requires.

bridge.training.global_batch_packing.wrap_data_iterator_for_global_batch_packing(

data_iterator: Any,

model_config: Any,

num_microbatches: int,

pg_collection: megatron.core.process_groups_config.ProcessGroupCollection,

) → tuple[Any, int, float, float]

Pack this step’s global batch into THD microbatches.

Pulls num_microbatches samples per data-parallel rank, gathers their lengths, schedules bins (and CP groups under dynamic CP), reroutes samples, and returns the packed iterator with the number of microbatches it yields. The packed iterator is a RerunDataIterator on TP rank 0 and None on other TP ranks, so callers must track “already wrapped” with an explicit flag.

Returns:

(packed_iterator, packed_num_microbatches, seqlen_sum, seqlen_squared_sum). The sums are data-parallel-global padded token statistics for FLOPs accounting.

bridge.training.global_batch_packing.get_batch_for_global_batch_packing(

data_iterator: Any,

model_config: Any,

*,

pg_collection: megatron.core.process_groups_config.ProcessGroupCollection,

vp_stage: int | None,

) → tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, None, torch.Tensor | None, megatron.core.packed_seq_params.PackedSeqParams | None, torch.Tensor | None]

Fetch one packed microbatch and its PackedSeqParams on every rank.

Megatron-Core slices the packed bin for this rank’s (possibly dynamic) CP group, broadcasts it over TP, and returns THD metadata.

Returns:

(tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask); attention_mask is always None for THD.

bridge.training.global_batch_packing.finalize_packed_seq_params(

packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None,

pg_collection: megatron.core.process_groups_config.ProcessGroupCollection,

) → megatron.core.packed_seq_params.PackedSeqParams | None

Bind the CP group to one microbatch and prebuild its THD CP partition route.

Uses pg_collection.cp as the static group so attention runs on the group the model was built with unless the scheduler attached a dynamic group. Releases with context-parallel layout routes get the route built here, off the forward path; older releases partition without one.