nemo_automodel.components.models.mimo_v2_flash.parallelization

MiMo-specific Transformer Engine context-parallel setup.

Module Contents

Functions

Name	Description
`_unwrap_checkpoint_module`	Return the attention module beneath any activation-checkpoint wrappers.
`ensure_mimo_te_context_parallel`	Install MiMo’s TE a2a transport once for the active CP process group.
`setup_mimo_te_context_parallel`	Configure every MiMo attention layer through its model-owned CP hook.

API

nemo_automodel.components.models.mimo_v2_flash.parallelization._unwrap_checkpoint_module(
    module: torch.nn.Module
) -> torch.nn.Module

Return the attention module beneath any activation-checkpoint wrappers.

nemo_automodel.components.models.mimo_v2_flash.parallelization.ensure_mimo_te_context_parallel(
    model: torch.nn.Module,
    cp_mesh
) -> None

Install MiMo’s TE a2a transport once for the active CP process group.

nemo_automodel.components.models.mimo_v2_flash.parallelization.setup_mimo_te_context_parallel(
    model: torch.nn.Module,
    cp_mesh
) -> None

Configure every MiMo attention layer through its model-owned CP hook.

The attention module owns TE backend validation, head-partition validation, and the a2a transport choice. This traversal only finds pipeline-local attention layers and shares one CUDA communication stream across them.

Parameters:

model

torch.nn.Module

A complete MiMo model or one pipeline-local model part.

cp_mesh

One-dimensional context-parallel device mesh.