> 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_v25.model

## Module Contents

### Classes

| Name                                                                                              | Description                                         |
| ------------------------------------------------------------------------------------------------- | --------------------------------------------------- |
| [`MiMoV2Attention`](#nemo_automodel-components-models-mimo_v25-model-MiMoV2Attention)             | MiMoV2 hybrid attention (full or sliding-window).   |
| [`MiMoV2DecoderLayer`](#nemo_automodel-components-models-mimo_v25-model-MiMoV2DecoderLayer)       | -                                                   |
| [`MiMoV2ForCausalLM`](#nemo_automodel-components-models-mimo_v25-model-MiMoV2ForCausalLM)         | NeMo AutoModel causal LM wrapper for MiMo-V2.5-Pro. |
| [`MiMoV2Model`](#nemo_automodel-components-models-mimo_v25-model-MiMoV2Model)                     | -                                                   |
| [`MiMoV2RotaryEmbedding`](#nemo_automodel-components-models-mimo_v25-model-MiMoV2RotaryEmbedding) | -                                                   |

### Functions

| Name                                                                                                                      | Description |
| ------------------------------------------------------------------------------------------------------------------------- | ----------- |
| [`_convert_bool_4d_mask_to_additive`](#nemo_automodel-components-models-mimo_v25-model-_convert_bool_4d_mask_to_additive) | -           |
| [`_derive_padding_mask`](#nemo_automodel-components-models-mimo_v25-model-_derive_padding_mask)                           | -           |
| [`_ensure_additive_mask`](#nemo_automodel-components-models-mimo_v25-model-_ensure_additive_mask)                         | -           |
| [`_fallback_additive_mask`](#nemo_automodel-components-models-mimo_v25-model-_fallback_additive_mask)                     | -           |
| [`apply_rotary_pos_emb`](#nemo_automodel-components-models-mimo_v25-model-apply_rotary_pos_emb)                           | -           |
| [`eager_attention_forward`](#nemo_automodel-components-models-mimo_v25-model-eager_attention_forward)                     | -           |
| [`repeat_kv`](#nemo_automodel-components-models-mimo_v25-model-repeat_kv)                                                 | -           |
| [`rotate_half`](#nemo_automodel-components-models-mimo_v25-model-rotate_half)                                             | -           |

### Data

[`ModelClass`](#nemo_automodel-components-models-mimo_v25-model-ModelClass)

### API

```python
class nemo_automodel.components.models.mimo_v25.model.MiMoV2Attention(
    config: nemo_automodel.components.models.mimo_v25.config.MiMoV2Config,
    is_swa: bool,
    layer_idx: int,
    projection_layout: str,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype
)
```

**Bases:** `Module`

MiMoV2 hybrid attention (full or sliding-window).

**`attention_dropout`** `= getattr(config, 'attention_dropout', 0.0)`

---

**`head_dim`**

---

**`k_proj`**

---

**`k_size`** `= self.num_key_value_heads * self.head_dim`

---

**`num_attention_heads`**

---

**`num_key_value_groups`**

---

**`num_key_value_heads`**

---

**`o_hidden_size`** `= self.num_attention_heads * self.v_head_dim`

---

**`o_proj`**

---

**`q_proj`**

---

**`q_size`** `= self.num_attention_heads * self.head_dim`

---

**`qkv_proj`**

---

**`rope_dim`**

---

**`scaling`** `= self.head_dim ** -0.5`

---

**`sliding_window`**

---

**`v_head_dim`**

---

**`v_proj`**

---

**`v_scale`** `= getattr(config, 'attention_value_scale', None)`

---

**`v_size`** `= self.num_key_value_heads * self.v_head_dim`

---

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2Attention._forward_attention(
    query_states: torch.Tensor,
    key_states: torch.Tensor,
    value_states: torch.Tensor,
    input_shape: torch.Size,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2Attention.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.mimo_v25.model.MiMoV2DecoderLayer(
    config: nemo_automodel.components.models.mimo_v25.config.MiMoV2Config,
    layer_idx: int,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

**`attention_type`**

---

**`input_layernorm`**

---

**`mlp`**

---

**`post_attention_layernorm`**

---

**`self_attn`**

---

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2DecoderLayer.forward(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    padding_mask: torch.Tensor | None = None
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.mimo_v25.model.MiMoV2ForCausalLM(
    config: nemo_automodel.components.models.mimo_v25.config.MiMoV2Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)
```

**Bases:** [HFCheckpointingMixin](/nemo-automodel/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin), `Module`, [MoEFSDPSyncMixin](/nemo-automodel/nemo_automodel/components/moe/fsdp_mixin#nemo_automodel-components-moe-fsdp_mixin-MoEFSDPSyncMixin)

NeMo AutoModel causal LM wrapper for MiMo-V2.5-Pro.

**`_keep_in_fp32_modules_strict`**

---

**`backend`** `= backend or BackendConfig()`

---

**`lm_head`**

---

**`model`**

---

**`state_dict_adapter`**

---

**`tie_word_embeddings_support`** `TieSupport = TieSupport.UNTIED_ONLY`

---

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2ForCausalLM.customize_pipeline_stage_modules(
    module_names_per_stage: list[list[str]],
    layers_prefix: str,
    text_model: torch.nn.Module | None = None
) -> list[list[str]]
```

Keep the SWA rotary embedding on every PP stage.

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2ForCausalLM.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    position_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    logits_to_keep: typing.Union[int, torch.Tensor] = 0,
    output_hidden_states: bool | None = None,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast
```

Compute logits or pass hidden states to the next pipeline stage.

**Parameters:**

**`input_ids`** `torch.Tensor | None` — default: None

Token IDs of shape \[batch, sequence] on the first stage,
or hidden states of shape \[batch, sequence, hidden] thereafter.

---

**`inputs_embeds`** `torch.FloatTensor | None` — default: None

Optional embeddings of shape \[batch, sequence, hidden].

---

**`position_ids`** `torch.LongTensor | None` — default: None

Optional positions of shape \[batch, sequence] or
\[1, sequence] for broadcasting across the batch.

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional padding mask of shape \[batch, sequence] or
additive attention mask of shape \[batch, 1, sequence, sequence].

---

**`padding_mask`** `torch.Tensor | None` — default: None

Optional padding indicators of shape \[batch, sequence].

---

**`logits_to_keep`** `Union[int, torch.Tensor]` — default: 0

Number of trailing token positions to retain, or
indices of shape \[retained\_sequence]. Zero retains all positions.

---

**`output_hidden_states`** `bool | None` — default: None

Whether to include the decoder hidden states.

---

**`**kwargs`** `Any` — default: \{}

Additional decoder arguments, including optional
cache\_position of shape \[sequence].

---

**Returns:** `CausalLMOutputWithPast`

A model output containing logits of shape \[batch, retained\_sequence,

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2ForCausalLM.from_config(
    config: nemo_automodel.components.models.mimo_v25.config.MiMoV2Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
) -> nemo_automodel.components.models.mimo_v25.model.MiMoV2ForCausalLM
```

classmethod

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2ForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args = (),
    kwargs = {}
) -> nemo_automodel.components.models.mimo_v25.model.MiMoV2ForCausalLM
```

classmethod

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2ForCausalLM.get_input_embeddings() -> torch.nn.Embedding
```

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2ForCausalLM.get_output_embeddings() -> torch.nn.Linear
```

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2ForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None
```

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2ForCausalLM.set_input_embeddings(
    value: torch.nn.Embedding
) -> None
```

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2ForCausalLM.set_output_embeddings(
    new_embeddings: torch.nn.Linear
) -> None
```

```python
class nemo_automodel.components.models.mimo_v25.model.MiMoV2Model(
    config: nemo_automodel.components.models.mimo_v25.config.MiMoV2Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

**`embed_tokens`**

---

**`has_sliding_layers`**

---

**`layers`**

---

**`norm`**

---

**`rotary_emb`** `= MiMoV2RotaryEmbedding(config=config, is_swa=False)`

---

**`swa_rotary_emb`** `= MiMoV2RotaryEmbedding(config=config, is_swa=True)`

---

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2Model._build_causal_mask_mapping(
    inputs_embeds: torch.Tensor,
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None,
    position_ids: torch.Tensor
) -> dict[str, torch.Tensor]
```

Build full and sliding masks for the uncached training sequence.

**Parameters:**

**`inputs_embeds`** `torch.Tensor`

Embeddings of shape \[batch, sequence, hidden].

---

**`attention_mask`** `torch.Tensor | dict[str, torch.Tensor] | None`

Padding mask of shape \[batch, sequence], an additive
mask of shape \[batch, 1, sequence, sequence], or a mapping of
attention types to masks with that four-dimensional layout.

---

**`position_ids`** `torch.Tensor`

Positions of shape \[batch, sequence] or \[1, sequence].

---

**Returns:** `dict[str, torch.Tensor]`

Full and sliding additive masks of shape \[batch, 1, sequence, sequence].

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2Model.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    position_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    cache_position: torch.LongTensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Run the decoder layers owned by this pipeline stage.

**Parameters:**

**`input_ids`** `torch.Tensor | None` — default: None

Token IDs of shape \[batch, sequence] on the embedding
stage, or hidden states of shape \[batch, sequence, hidden] on
later stages.

---

**`inputs_embeds`** `torch.FloatTensor | None` — default: None

Optional embeddings of shape \[batch, sequence, hidden].

---

**`position_ids`** `torch.LongTensor | None` — default: None

Optional positions of shape \[batch, sequence] or
\[1, sequence] for broadcasting across the batch.

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional padding mask of shape \[batch, sequence] or
additive attention mask of shape \[batch, 1, sequence, sequence].

---

**`padding_mask`** `torch.Tensor | None` — default: None

Optional padding indicators of shape \[batch, sequence].

---

**`cache_position`** `torch.LongTensor | None` — default: None

Optional token positions of shape \[sequence].

---

**`**kwargs`** `Any` — default: \{}

Unused compatibility arguments.

---

**Returns:** `torch.Tensor`

Hidden states of shape \[batch, sequence, hidden], normalized only

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2Model.init_weights(
    buffer_device: torch.device | None = None
) -> None
```

```python
class nemo_automodel.components.models.mimo_v25.model.MiMoV2RotaryEmbedding(
    config: nemo_automodel.components.models.mimo_v25.config.MiMoV2Config,
    is_swa: bool,
    device: torch.device | None = None
)
```

**Bases:** `Module`

**`config`** `= copy(config)`

---

**`inv_freq`** `Tensor`

---

**`max_seq_len_cached`** `= config.max_position_embeddings`

---

**`original_inv_freq`** `= self.inv_freq`

---

**`original_max_seq_len`** `= config.max_position_embeddings`

---

**`rope_init_fn`**

---

**`rope_type`**

---

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2RotaryEmbedding.compute_default_rope_parameters(
    config: nemo_automodel.components.models.mimo_v25.config.MiMoV2Config,
    device: torch.device | None = None,
    seq_len: int | None = None,
    layer_type: str | None = None
) -> tuple[torch.Tensor, float]
```

staticmethod

```python
nemo_automodel.components.models.mimo_v25.model.MiMoV2RotaryEmbedding.forward(
    x: torch.Tensor,
    position_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

```python
nemo_automodel.components.models.mimo_v25.model._convert_bool_4d_mask_to_additive(
    mask: torch.Tensor,
    dtype: torch.dtype
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v25.model._derive_padding_mask(
    attention_mask: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v25.model._ensure_additive_mask(
    mask: torch.Tensor | None,
    batch_size: int,
    seq_len: int,
    dtype: torch.dtype,
    device: torch.device,
    attention_mask: torch.Tensor | None,
    sliding_window: int | None
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v25.model._fallback_additive_mask(
    batch_size: int,
    seq_len: int,
    dtype: torch.dtype,
    device: torch.device,
    attention_mask: torch.Tensor | None = None,
    sliding_window: int | None = None
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v25.model.apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    unsqueeze_dim: int = 1
) -> tuple[torch.Tensor, torch.Tensor]
```

```python
nemo_automodel.components.models.mimo_v25.model.eager_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    scaling: float,
    dropout: float = 0.0,
    sinks: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor]
```

```python
nemo_automodel.components.models.mimo_v25.model.repeat_kv(
    hidden_states: torch.Tensor,
    n_rep: int
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v25.model.rotate_half(
    x: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v25.model.ModelClass = MiMoV2ForCausalLM
```