bridge.data.megatron_mimo.canonical_sampler

Canonical-grid batch sampling for MegatronMIMO scalable data-parallel reads.

With megatron_mimo_scalable_dp every module’s loaders must materialize the same ordered global micro-batch, because the BridgeCommunicator routes modality embeddings to language ranks by contiguous position along the batch dim. A sampler’s shard assignment depends on its (data_parallel_rank, data_parallel_size, micro_batch_size), so modules with different DP sizes cannot shard by their own DP: a shuffling sampler would hand them different sample sets. Instead, all loaders shard on one shared canonical grid — the least common multiple of the module DP sizes. A rank whose module has DP size d covers grid // d consecutive canonical groups and concatenates their windows, so concatenating any module’s rank batches reproduces the identical ordered global micro-batch under any deterministic sampler (single and cyclic included; the shuffle seed derives from the epoch alone).

Module Contents

Classes

CanonicalGroupBatchSampler

Concatenate per-canonical-group Megatron samplers into one batch sampler.

Functions

canonical_grid_size	Return the shared sampler grid size: the LCM of every module’s DP size.
covered_canonical_groups	Return the canonical groups this rank reads.
build_canonical_group_batch_sampler	Build this rank’s canonical-group batch sampler for scalable MIMO reads.
build_canonical_mimo_data_loader	Build this rank’s read-sharded DataLoader for scalable MegatronMIMO reads.

API

bridge.data.megatron_mimo.canonical_sampler.canonical_grid_size(module_dp_sizes: list[int]) → int

Return the shared sampler grid size: the LCM of every module’s DP size.

bridge.data.megatron_mimo.canonical_sampler.covered_canonical_groups(

dp_rank: int,

dp_size: int,

grid_size: int,

) → list[int]

Return the canonical groups this rank reads.

The groups are the grid_size // dp_size consecutive slots matching the contiguous batch-dim chunk the BridgeCommunicator routes to this DP rank.

class bridge.data.megatron_mimo.canonical_sampler.CanonicalGroupBatchSampler(samplers: list)

Concatenate per-canonical-group Megatron samplers into one batch sampler.

Holds one flat sampler per covered canonical group, all built with the same (micro_batch // grid, grid) geometry and the same global consumed_samples. Each yield emits the groups’ current windows concatenated in group order. Group streams are functions of (group, grid, micro_batch, consumed, epoch seed) only, so every rank covering group g — in any module — sees the identical stream.

Initialization

__len__() → int

Match the flat Megatron samplers’ convention of reporting total samples.

__iter__() → Iterator[list[int]]

Yield one concatenated index window per micro-batch.

bridge.data.megatron_mimo.canonical_sampler.build_canonical_group_batch_sampler(

*,

dataloader_type: str,

dataset: torch.utils.data.Dataset,

consumed_samples: int,

micro_batch_size: int,

grid_size: int,

groups: list[int],

data_sharding: bool,

drop_last: bool = True,

) → bridge.data.megatron_mimo.canonical_sampler.CanonicalGroupBatchSampler

Build this rank’s canonical-group batch sampler for scalable MIMO reads.

Parameters:

• dataloader_type – "single" or "cyclic"; other types have no shard assignment that is consistent across modules and are rejected.

• dataset – The dataset the loader reads.

• consumed_samples – Global consumed-sample count, exactly as the flat samplers expect (used for resume / epoch derivation).

• micro_batch_size – The global micro-batch size (not the per-rank share).

• grid_size – Canonical grid size from :func:canonical_grid_size.

• groups – This rank’s groups from :func:covered_canonical_groups.

• data_sharding – Passed through to the cyclic sampler.

• drop_last – Must stay True: a partial final window gives the groups unequal shares and breaks the positional routing.

Returns:

The merged batch sampler for torch.utils.data.DataLoader(batch_sampler=...).

bridge.data.megatron_mimo.canonical_sampler.build_canonical_mimo_data_loader(

dataset: torch.utils.data.Dataset | None,

*,

consumed_samples: int,

dataloader_type: str,

micro_batch_size: int,

module_dp_sizes: list[int],

dp_rank: int,

dp_size: int,

data_sharding: bool,

drop_last: bool,

num_workers: int,

pin_memory: bool,

collate_fn: Callable | None,

persistent_workers: bool,

) → torch.utils.data.DataLoader | None

Build this rank’s read-sharded DataLoader for scalable MegatronMIMO reads.

Computes the canonical grid from module_dp_sizes, derives this rank’s covered groups from its module-local (dp_rank, dp_size), and wraps the merged batch sampler in a DataLoader. Returns None when dataset is None.