bridge.models.mimo.mimo_builder

Module Contents

Functions

build_hypercomm_grids	Create HyperCommGrid objects per module from MIMO parallelism config.
_default_topology	Infer a default multi-encoder -> LLM topology.

API

bridge.models.mimo.mimo_builder.build_hypercomm_grids(

mimo_parallelism_config: megatron.bridge.models.mimo.mimo_config.MimoParallelismConfig,

) → Dict[str, megatron.core.hyper_comm_grid.HyperCommGrid]

Create HyperCommGrid objects per module from MIMO 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.

Parameters:

mimo_parallelism_config – MimoParallelismConfig specifying parallelism per module.

Returns:

Dict mapping module names to their HyperCommGrids.

bridge.models.mimo.mimo_builder._default_topology(

mimo_parallelism_config: megatron.bridge.models.mimo.mimo_config.MimoParallelismConfig,

) → Dict[str, List[str]]

Infer a default multi-encoder -> LLM topology.