core.models.mimo.config.role

Data classes for MIMO rank role management in multi-module pipeline parallelism.

Module Contents

Classes

ModuleLayout	Pipeline mode for MIMO multi-module parallelism.
ModuleStageInfo	Information about a rank’s stage position within a module’s pipeline.
RankRole	Describes what modules this rank participates in for multi-module PP.

Data

logger
MIMO_LANGUAGE_MODULE_KEY

API

core.models.mimo.config.role.logger

‘getLogger(…)’

core.models.mimo.config.role.MIMO_LANGUAGE_MODULE_KEY

‘language’

class core.models.mimo.config.role.ModuleLayout(*args, **kwds)

Bases: enum.Enum

Pipeline mode for MIMO multi-module parallelism.

Determines how modules are distributed across ranks and which forward path is used.

COLOCATED: All modules share the same ranks. Covers both legacy (no grid map, global parallel_state) and heterogeneous TP/DP (grid map with overlapping ranks). Uses _forward_all_modules.

NON_COLOCATED: module_to_grid_map is set with non-overlapping rank ranges. Each rank runs EITHER encoder(s) OR the language model. Uses role-based dispatch with separate forward paths.

Initialization

COLOCATED

‘colocated’

NON_COLOCATED

‘non_colocated’

class core.models.mimo.config.role.ModuleStageInfo

Information about a rank’s stage position within a module’s pipeline.

Parameters:

• is_first_stage – True if this rank is the first PP stage for this module.

• is_last_stage – True if this rank is the last PP stage for this module.

is_first_stage: bool

None

is_last_stage: bool

None

class core.models.mimo.config.role.RankRole

Describes what modules this rank participates in for multi-module PP.

This class captures the role of a specific rank in a multi-module pipeline parallel setup, tracking which modules the rank participates in and their stage positions. The language module is always identified by MIMO_LANGUAGE_MODULE_KEY.

Parameters:

• modules – Dict mapping module names to their stage info for modules this rank participates in.

• mode – Pipeline mode determining forward path dispatch.

modules: Dict[str, core.models.mimo.config.role.ModuleStageInfo]

‘field(…)’

mode: core.models.mimo.config.role.ModuleLayout

None

classmethod build(

modality_module_names: List[str],

module_to_grid_map: Optional[Dict[str, megatron.core.hyper_comm_grid.HyperCommGrid]] = None,

) → core.models.mimo.config.role.RankRole

Build a RankRole, dispatching by whether grids share ranks.

No grid map or all grids span the same ranks → COLOCATED. Grids differ → NON_COLOCATED with PP-stage info per module.

static _all_grids_colocated(

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

) → bool

classmethod _colocated(

modality_module_names: List[str],

) → core.models.mimo.config.role.RankRole

Colocated layout: every module on every rank, PP=1.

classmethod _from_grid_map(

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

) → core.models.mimo.config.role.RankRole

Non-colocated role for this rank from a module-to-grid mapping.

Grid map keys are validated by MimoModelConfig.__post_init__.

Raises:

RuntimeError – If current rank is not in any module grid.

property has_modality_modules: bool

Return True if this rank participates in any modality (non-language) module.

property has_language_module: bool

Return True if this rank participates in the language module.

property modality_module_names: List[str]

Return names of modality modules (non-language) this rank participates in.

is_first_stage(module_name: str) → bool

Check if this rank is the first stage for a given module.

is_last_stage(module_name: str) → bool

Check if this rank is the last stage for a given module.