> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.mimo_v2_flash.parallelization

MiMo-specific Transformer Engine context-parallel setup.

## Module Contents

### Functions

| Name                                                                                                                                 | Description                                                             |
| ------------------------------------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------- |
| [`_unwrap_checkpoint_module`](#nemo_automodel-components-models-mimo_v2_flash-parallelization-_unwrap_checkpoint_module)             | Return the attention module beneath any activation-checkpoint wrappers. |
| [`ensure_mimo_te_context_parallel`](#nemo_automodel-components-models-mimo_v2_flash-parallelization-ensure_mimo_te_context_parallel) | Install MiMo's TE a2a transport once for the active CP process group.   |
| [`setup_mimo_te_context_parallel`](#nemo_automodel-components-models-mimo_v2_flash-parallelization-setup_mimo_te_context_parallel)   | Configure every MiMo attention layer through its model-owned CP hook.   |

### API

```python
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.

```python
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.

```python
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.

---