bridge.models.megatron_mimo.megatron_mimo_builder#

Module Contents#

Functions#

build_hypercomm_grids

Create HyperCommGrid objects per module from MegatronMIMO parallelism config.

populate_embedding_and_position_groups

Create embedding-related process groups from globally enumerated PP ranks.

is_pp_first_stage

Check if current rank is first stage in pipeline.

is_pp_last_stage

Check if current rank is last stage in pipeline.

Data#

EXPERT_VIEW_NAME

API#

bridge.models.megatron_mimo.megatron_mimo_builder.EXPERT_VIEW_NAME#

‘expert’

bridge.models.megatron_mimo.megatron_mimo_builder.build_hypercomm_grids(
megatron_mimo_parallelism_config: megatron.bridge.models.megatron_mimo.megatron_mimo_config.MegatronMIMOParallelismConfig,
) → Dict[str, megatron.core.hyper_comm_grid.HyperCommGrid]#

Create HyperCommGrid objects per module from MegatronMIMO parallelism config.

Creates grids on ALL ranks (required for consistent collective calls), but only ranks in each grid’s range will participate in its operations.

Each grid is built with a dense base view plus a registered expert view over the same rank span. Dense process groups (tp/cp/dp/pp) come from the base view; expert-parallel groups (expt_tp/ep/expt_dp) come from the expert view, matching the contract mcore’s get_mimo_optimizer expects.

Parameters:

megatron_mimo_parallelism_config – MegatronMIMOParallelismConfig specifying parallelism per module.

Returns:

Dict mapping module names to their HyperCommGrids.

bridge.models.megatron_mimo.megatron_mimo_builder.populate_embedding_and_position_groups(
pp_rank_groups: list[list[int]] | None,
) → Tuple[Optional[torch.distributed.ProcessGroup], Optional[torch.distributed.ProcessGroup]]#

Create embedding-related process groups from globally enumerated PP ranks.

Following MCore semantics:

  • pos_embd_pg: Only rank 0 of PP (first stage) - for position embeddings

  • embd_pg: Ranks 0 and -1 of PP (first and last stages) - for tied word embeddings

IMPORTANT: This calls dist.new_group which is a collective operation. Must be called on all ranks that could participate.

Parameters:

pp_rank_groups – Every pipeline-parallel rank group in global creation order.

Returns:

Tuple of process groups for the current rank. Returns (None, None) when no pipeline-parallel rank groups are provided.

bridge.models.megatron_mimo.megatron_mimo_builder.is_pp_first_stage(
pp_group: Optional[torch.distributed.ProcessGroup],
) → bool#

Check if current rank is first stage in pipeline.

bridge.models.megatron_mimo.megatron_mimo_builder.is_pp_last_stage(
pp_group: Optional[torch.distributed.ProcessGroup],
) → bool#

Check if current rank is last stage in pipeline.