bridge.data.megatron_mimo.dp_utils

Data parallel utilities for MegatronMIMO data loading.

Module Contents

Functions

_batch_dim_for_tensor	Return the batch dimension for known MegatronMIMO batch tensors.
_is_patch_packed_visual_dict	Detect a patch-packed visual encoder input layout.
_slice_patch_packed_visual_dict	Joint-slice a patch-packed {hidden_states, grid_thw, ...} dict.
_find_rank_module	Find which module grid the current rank belongs to.
_needs_data_for_module	Determine if the current rank needs to load data for the given module.
get_megatron_mimo_dp_info	Get module-local DP rank, size, data-loading flag, and module name.
get_megatron_mimo_sampling_info	Get sampler DP rank, size, and data-loading flag for MegatronMIMO.
slice_batch_for_megatron_mimo	Slice a global micro-batch for this rank’s module-local DP shard.
real_token_lengths	Per-sample real (non-pad) token length for a [B, S] batch.

API

bridge.data.megatron_mimo.dp_utils._batch_dim_for_tensor(key: str, value: torch.Tensor) → int

Return the batch dimension for known MegatronMIMO batch tensors.

bridge.data.megatron_mimo.dp_utils._is_patch_packed_visual_dict(value: Any) → bool

Detect a patch-packed visual encoder input layout.

Some VLM adapters pack all image patches for a microbatch into one flat tensor while keeping per-image grid metadata:

• hidden_states: [sum(patches_across_images), patch_feature_dim]

• grid_thw: [num_images, 3]

Dim 0 of these two tensors means different things (patches vs images), so they cannot be sliced independently by DP. They must be sliced jointly along the per-image boundary.

TODO(mimo): This is very model-specific. We should revisit this to make it more generic. maybe it should be handled by the model itself that defines how to slice the data.

bridge.data.megatron_mimo.dp_utils._slice_patch_packed_visual_dict(

value: Dict[str, Any],

dp_rank: int,

dp_size: int,

) → Dict[str, Any]

Joint-slice a patch-packed {hidden_states, grid_thw, ...} dict.

Shards by image count, then derives the corresponding patch range from cumsum(grid_thw.prod(dim=-1)). Other keys in the dict are passed through unchanged because they are treated as global encoder metadata.

Constraints:

• num_images (rows of grid_thw) must be divisible by dp_size. For the encoder rank, num_images_per_microbatch is set by the training-time MIMO MBS and the dataset’s images-per-sample. With single-image-per-sample data and MIMO MICRO_BATCH_SIZE divisible by encoder DP, this always holds.

TODO(mimo): This convention slices by image row. Multi-image samples with uneven image counts need explicit per-sample image split metadata so DP slices remain sample-aligned instead of only image-aligned.

bridge.data.megatron_mimo.dp_utils._find_rank_module(

grids: Dict[str, megatron.core.hyper_comm_grid.HyperCommGrid],

) → Tuple[HyperCommGrid | None, str | None]

Find which module grid the current rank belongs to.

bridge.data.megatron_mimo.dp_utils._needs_data_for_module(

grid: megatron.core.hyper_comm_grid.HyperCommGrid,

module_name: str,

) → bool

Determine if the current rank needs to load data for the given module.

LLM: all PP stages need batch metadata. First PP stages consume input_ids, last PP stages consume labels/loss_mask, and models with position-dependent decoder blocks (for example Qwen3-VL MRoPE) need position_ids on intermediate PP stages as well. Encoders: only the first PP stage needs raw modality inputs.

bridge.data.megatron_mimo.dp_utils.get_megatron_mimo_dp_info(

megatron_mimo_cfg: megatron.bridge.models.megatron_mimo.megatron_mimo_config.MegatronMIMOParallelismConfig,

grids: Dict[str, megatron.core.hyper_comm_grid.HyperCommGrid],

) → Tuple[int, int, bool, str]

Get module-local DP rank, size, data-loading flag, and module name.

Returns the DP settings for the module that the current rank participates in. These are used by :func:slice_batch_for_megatron_mimo to sub-shard a global micro-batch into per-module DP shards.

.. note:: Do not use these values to construct a DistributedSampler. For sampler construction use :func:get_megatron_mimo_sampling_info instead, which returns settings that keep all data-loading ranks synchronised on the same sample order.

Parameters:

• megatron_mimo_cfg – MegatronMIMO parallelism configuration.

• grids – Module name to HyperCommGrid mapping from build_hypercomm_grids().

Returns:

Tuple of (dp_rank, dp_size, needs_data, loader_module).

bridge.data.megatron_mimo.dp_utils.get_megatron_mimo_sampling_info(

megatron_mimo_cfg: megatron.bridge.models.megatron_mimo.megatron_mimo_config.MegatronMIMOParallelismConfig,

grids: Dict[str, megatron.core.hyper_comm_grid.HyperCommGrid],

*,

scalable_dp: bool = False,

) → Tuple[int, int, bool]

Get sampler DP rank, size, and data-loading flag for MegatronMIMO.

In heterogeneous MegatronMIMO, modules may have different DP sizes.

Default (full-batch reads). Returns dp_size=1, dp_rank=0 for all ranks, so the data loader gives every data-loading rank the same global micro-batch; per-module DP sharding is deferred to :func:slice_batch_for_megatron_mimo in the forward step, consistent with the

Class:

BridgeCommunicator fan-in / fan-out routing.

Scalable data parallelism (scalable_dp=True). Returns this rank’s module-local (dp_rank, dp_size), from which the caller derives its canonical groups and builds the read-sharded loader (see

Func:

megatron.bridge.data.megatron_mimo.canonical_sampler.build_canonical_group_batch_sampler); the forward step does not slice again. Each rank then processes its natural, unbalanced 1/dp shard: this is the read-sharding win on its own, with no cross-rank exchange.

Parameters:

• megatron_mimo_cfg – MegatronMIMO parallelism configuration.

• grids – Module name to HyperCommGrid mapping.

• scalable_dp – When True, shard reads at the sampler (module-local DP).

Returns:

Tuple of (sampler_dp_rank, sampler_dp_size, needs_data).

bridge.data.megatron_mimo.dp_utils.slice_batch_for_megatron_mimo(

batch: Dict[str, Any],

dp_rank: int,

dp_size: int,

) → Dict[str, Any]

Slice a global micro-batch for this rank’s module-local DP shard.

In the default (non-scalable) mode all data-loading ranks receive the same global micro-batch (the sampler uses dp_size=1) and this function contiguously slices it so that each module-local DP replica processes the correct subset. The slicing is contiguous to match the

Class:

BridgeCommunicator’s batch-dimension split / concatenate logic for fan-out and fan-in routing. With megatron_mimo_scalable_dp the loader already delivers the shard and forward_step skips this call.

Handles nested dicts (e.g. modality_inputs) by recursing.

Parameters:

• batch – Global batch dictionary with tensors of shape [global_batch, …], except known layouts such as Qwen-VL MRoPE position_ids shaped [3, global_batch, seq]. May contain nested dicts (e.g. modality_inputs -> encoder -> kwargs).

• dp_rank – This rank’s position in its module-local DP group.

• dp_size – Size of the module-local DP group.

Returns:

Dict with tensors sliced to shape [global_batch // dp_size, …].

.. rubric:: Example

global_batch = {‘tokens’: torch.randn(12, 2048)} local_batch = slice_batch_for_megatron_mimo(global_batch, dp_rank=1, dp_size=3) local_batch[‘tokens’].shape # torch.Size([4, 2048])

bridge.data.megatron_mimo.dp_utils.real_token_lengths(

input_ids: torch.Tensor,

*,

attention_mask: torch.Tensor | None,

) → torch.Tensor

Per-sample real (non-pad) token length for a [B, S] batch.

The length comes only from attention_mask (1 = real, 0 = pad), the padding mask the tokenizer ships with the batch. Comparing input_ids against pad_token_id is not used: tokenizers without a dedicated pad token alias pad_token = eos_token, so a value-based check would also drop genuine EOS tokens. loss_mask is not used either: it is a supervision mask (prompt tokens are zero) and would under-count the real length.

Parameters:

• input_ids – Padded token ids [B, S].

• attention_mask – [B, S] padding mask matching input_ids.

Returns:

An int64 tensor of shape [B] with each sample’s real length.

Raises:

ValueError – If attention_mask is missing or does not match input_ids.