bridge.data.mimo.dp_utils#

Data parallel utilities for MIMO data loading.

Module Contents#

Classes#

MimoDpInfo

Data-parallel loader metadata for the current rank in MIMO training.

Functions#

get_mimo_dp_info

Get DP rank, size, data-loading responsibility, and loader module for MIMO.

API#

class bridge.data.mimo.dp_utils.MimoDpInfo#

Data-parallel loader metadata for the current rank in MIMO training.

dp_rank: int#

None

dp_size: int#

None

needs_data: bool#

None

loader_module: str#

None

bridge.data.mimo.dp_utils.get_mimo_dp_info(
grids: Dict[str, megatron.core.hyper_comm_grid.HyperCommGrid],
) → bridge.data.mimo.dp_utils.MimoDpInfo#

Get DP rank, size, data-loading responsibility, and loader module for MIMO.

Determines which module’s DP settings to use for data loading based on current rank’s participation in heterogeneous deployment.

In heterogeneous mode, each rank uses its own module’s DP settings.

Parameters:

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

Returns:

  • dp_rank: This rank’s position in DP group.

  • dp_size: Size of DP group for data sharding.

  • needs_data: Whether this rank needs to load data (first/last PP stage).

  • loader_module: Which module’s DP settings are being used.

Return type:

MimoDpInfo with

.. rubric:: Example

from megatron.bridge.models.mimo.mimo_builder import build_hypercomm_grids grids = build_hypercomm_grids(mimo_cfg) dp_info = get_mimo_dp_info(grids) if dp_info.needs_data: … # Build data loader with dp_rank and dp_size … sampler = DistributedSampler(dataset, num_replicas=dp_info.dp_size, rank=dp_info.dp_rank)