Is this page helpful?

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.

UNIFIED: No module_to_grid_map. All modules share same ranks and parallelism. Uses the unified forward path (_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.

COLOCATED: (future) module_to_grid_map is set with overlapping rank ranges. Encoder(s) and language model share ranks but have different parallelism configs. Uses role-based dispatch but allows both module types on the same rank.

Initialization

UNIFIED#

‘unified’

NON_COLOCATED#

‘non_colocated’

COLOCATED#

‘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 unified(
module_names: List[str],
) → core.models.mimo.config.role.RankRole#

Create a role for the unified case: every module, first+last stage.

classmethod from_grid_map(
module_to_grid_map: Dict[str, megatron.core.hyper_comm_grid.HyperCommGrid],
modality_module_names: List[str],
) → core.models.mimo.config.role.RankRole#

Create a role from a module-to-grid mapping for non-colocated PP.

Determines which modules the current rank participates in and its pipeline stage position within each module.

Parameters:

  • module_to_grid_map – Dict mapping module names to HyperCommGrid objects. Must contain keys matching modality_module_names + MIMO_LANGUAGE_MODULE_KEY.

  • modality_module_names – List of modality module names (e.g., [“images”, “audio”]).

Returns:

RankRole for the current rank.

Raises:

  • ValueError – If grid map keys don’t match expected module names.

  • 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.