> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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

## Submodules

* **[`nemo_automodel.components.models.mimo_v2_flash.config`](/nemo-automodel/nemo_automodel/components/models/mimo_v2_flash/config)**
* **[`nemo_automodel.components.models.mimo_v2_flash.cp`](/nemo-automodel/nemo_automodel/components/models/mimo_v2_flash/cp)**
* **[`nemo_automodel.components.models.mimo_v2_flash.model`](/nemo-automodel/nemo_automodel/components/models/mimo_v2_flash/model)**
* **[`nemo_automodel.components.models.mimo_v2_flash.parallelization`](/nemo-automodel/nemo_automodel/components/models/mimo_v2_flash/parallelization)**
* **[`nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter`](/nemo-automodel/nemo_automodel/components/models/mimo_v2_flash/state_dict_adapter)**
* **[`nemo_automodel.components.models.mimo_v2_flash.vision`](/nemo-automodel/nemo_automodel/components/models/mimo_v2_flash/vision)**

## Package Contents

### Classes

| Name                                                                                                                            | Description                                                             |
| ------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`MiMoV2Config`](#nemo_automodel-components-models-mimo_v2_flash-config-MiMoV2Config)                                           | Configuration for MiMo-V2.5 and MiMo-V2.6 checkpoints.                  |
| [`MiMoV2FlashConfig`](#nemo_automodel-components-models-mimo_v2_flash-config-MiMoV2FlashConfig)                                 | Configuration for XiaomiMiMo/MiMo-V2-Flash.                             |
| [`MiMoV2FlashForCausalLM`](#nemo_automodel-components-models-mimo_v2_flash-model-MiMoV2FlashForCausalLM)                        | Causal LM wrapper for MiMo-V2-Flash with Automodel checkpoint adapters. |
| [`MiMoV2FlashModel`](#nemo_automodel-components-models-mimo_v2_flash-model-MiMoV2FlashModel)                                    | Backbone model for Xiaomi MiMo-V2-Flash.                                |
| [`MiMoV2FlashStateDictAdapter`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-MiMoV2FlashStateDictAdapter) | Convert MiMo-V2-Flash HF checkpoints to Automodel's grouped MoE layout. |
| [`MiMoV2ForCausalLM`](#nemo_automodel-components-models-mimo_v2_flash-model-MiMoV2ForCausalLM)                                  | MiMo-V2.5/V2.6 wrapper using the `mimo_v2` checkpoint config.           |

### API

```python
class nemo_automodel.components.models.mimo_v2_flash.config.MiMoV2Config(
    kwargs: typing.Any = {}
)
```

**Bases:** [MiMoV2FlashConfig](#nemo_automodel-components-models-mimo_v2_flash-config-MiMoV2FlashConfig)

Configuration for MiMo-V2.5 and MiMo-V2.6 checkpoints.

These checkpoints retain the MiMo-V2-Flash text backbone while using the
`mimo_v2` model type, fused QKV checkpoint layout, and optional
multimodal encoder configuration.

**`model_type`** `= 'mimo_v2'`

---

```python
class nemo_automodel.components.models.mimo_v2_flash.config.MiMoV2FlashConfig(
    vocab_size: int = 152576,
    hidden_size: int = 4096,
    intermediate_size: int = 16384,
    moe_intermediate_size: int = 2048,
    num_hidden_layers: int = 48,
    num_attention_heads: int = 64,
    num_key_value_heads: int = 4,
    head_dim: int = 192,
    v_head_dim: int = 128,
    swa_num_attention_heads: int = 64,
    swa_num_key_value_heads: int = 8,
    swa_head_dim: int = 192,
    swa_v_head_dim: int = 128,
    hidden_act: str = 'silu',
    max_position_embeddings: int = 262144,
    initializer_range: float = 0.02,
    layernorm_epsilon: float = 1e-05,
    rms_norm_eps: float | None = None,
    use_cache: bool = True,
    tie_word_embeddings: bool = False,
    rope_theta: float = 5000000.0,
    swa_rope_theta: float = 10000.0,
    rope_scaling: dict | None = None,
    attention_bias: bool = False,
    attention_dropout: float = 0.0,
    attention_projection_layout: str = 'split',
    attention_value_scale: float | None = 0.707,
    add_full_attention_sink_bias: bool = False,
    add_swa_attention_sink_bias: bool = True,
    hybrid_block_size: int | None = None,
    hybrid_layer_pattern: list[int] | None = None,
    partial_rotary_factor: float = 0.334,
    sliding_window: int | None = 128,
    sliding_window_size: int | None = 128,
    attention_chunk_size: int | None = 128,
    n_routed_experts: int | None = 256,
    n_shared_experts: int | None = None,
    num_experts_per_tok: int = 8,
    scoring_func: str = 'sigmoid',
    topk_method: str = 'noaux_tc',
    n_group: int = 1,
    topk_group: int = 1,
    norm_topk_prob: bool = True,
    routed_scaling_factor: float | None = 1.0,
    apply_router_weight_after_down: bool = False,
    moe_layer_freq: list[int] | None = None,
    torch_dtype: str = 'bfloat16',
    vision_config: dict | None = None,
    audio_config: dict | None = None,
    processor_config: dict | None = None,
    kwargs = {}
)
```

**Bases:** `PretrainedConfig`

Configuration for XiaomiMiMo/MiMo-V2-Flash.

The Hugging Face remote config class currently leaves `model_type` empty.
Automodel registers this local config with the hub's JSON `model_type` so
configs can resolve without executing remote code.

**`attribute_map`** `= {'num_local_experts': 'n_routed_experts'}`

---

**`keys_to_ignore_at_inference`** `= ['past_key_values']`

---

**`layer_types`**

---

**`model_type`** `= 'mimo_v2_flash'`

---

**`moe_layer_freq`**

---

**`num_key_value_heads`**

---

**`rms_norm_eps`**

---

**`rope_parameters`**

---

**`routed_scaling_factor`**

---

```python
class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM(
    config: nemo_automodel.components.models.mimo_v2_flash.config.MiMoV2FlashConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    revision: str | 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)

Causal LM wrapper for MiMo-V2-Flash with Automodel checkpoint adapters.

**`_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_v2_flash.model.MiMoV2FlashForCausalLM._get_multimodal_embeds(
    input_ids: torch.Tensor,
    inputs_embeds: torch.Tensor,
    pixel_values: torch.Tensor | None = None,
    image_grid_thw: torch.Tensor | None = None,
    image_feature_indices: torch.Tensor | None = None,
    pixel_values_videos: torch.Tensor | None = None,
    video_grid_thw: torch.Tensor | None = None,
    video_feature_indices: torch.Tensor | None = None
) -> torch.Tensor
```

Encode global media and splice only this TE THD shard's feature rows.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM._local_modal_feature_indices(
    global_mask: torch.Tensor | None,
    local_indices: torch.Tensor | None
) -> torch.Tensor | None
```

staticmethod

Map local TE THD placeholders to rows in globally ordered media features.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM._pull_pipeline_media(
    input_ids: torch.Tensor | None,
    pixel_values: torch.Tensor | None,
    image_grid_thw: torch.Tensor | None,
    pixel_values_videos: torch.Tensor | None,
    video_grid_thw: torch.Tensor | None,
    global_image_mask: torch.Tensor | None = None,
    global_video_mask: torch.Tensor | None = None
) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]
```

Retrieve PP media using global masks so every CP rank runs the same vision work.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.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 both rotary embeddings per stage and the vision tower on stage zero.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None = None,
    padding_mask: torch.Tensor | None = None,
    pixel_values: torch.Tensor | None = None,
    image_grid_thw: torch.Tensor | None = None,
    pixel_values_videos: torch.Tensor | None = None,
    video_grid_thw: 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 | torch.Tensor
```

Run text or VLM forward with TE-local packed media placement.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.from_config(
    config: nemo_automodel.components.models.mimo_v2_flash.config.MiMoV2FlashConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)
```

classmethod

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args = (),
    kwargs = {}
)
```

classmethod

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.get_input_embeddings()
```

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.get_output_embeddings()
```

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.get_pipeline_stage_metas(
    is_first: bool,
    microbatch_size: int,
    seq_len: int,
    dtype: torch.dtype
) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]
```

Describe PP tensors using the already TE-sharded local token length.

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

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.prepare_model_inputs_for_cp(
    batch: dict[str, typing.Any],
    num_chunks: int = 1
) -> dict[str, typing.Any]
```

Install MiMo's model-local adapter around framework TE THD sharding.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.set_input_embeddings(
    value
)
```

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.set_output_embeddings(
    new_embeddings
)
```

```python
class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashModel(
    config: nemo_automodel.components.models.mimo_v2_flash.config.MiMoV2FlashConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    moe_overrides: dict | None = None
)
```

**Bases:** `Module`

Backbone model for Xiaomi MiMo-V2-Flash.

**`embed_tokens`**

---

**`has_sliding_layers`**

---

**`layers`**

---

**`moe_config`** `= moe_config or MoEConfig(**moe_defaults)`

---

**`norm`**

---

**`rotary_emb`**

---

**`swa_rotary_emb`**

---

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

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashModel.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None = None,
    padding_mask: torch.Tensor | None = None,
    cache_position: torch.Tensor | None = None,
    _packed_seq_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    seq_lens_padded: torch.Tensor | None = None,
    cu_seqlens: torch.Tensor | None = None,
    cu_seqlens_padded: torch.Tensor | None = None,
    max_seqlen: int | torch.Tensor | None = None,
    qkv_format: str | None = None,
    cp_size: int = 1,
    cp_rank: int = 0,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Run the text backbone in BSHD or genuine TE THD layout.

THD batches must already have passed through the framework TE sharder.
Their cumulative lengths delimit packed documents and are forwarded to
every attention layer, including non-first pipeline stages.

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

```python
class nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter(
    config: typing.Any,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.bfloat16,
    revision: str | None = None
)
```

**Bases:** [MoESplitExpertsStateDictMixin](/nemo-automodel/nemo_automodel/components/moe/state_dict_mixin#nemo_automodel-components-moe-state_dict_mixin-MoESplitExpertsStateDictMixin), [StateDictAdapter](/nemo-automodel/nemo_automodel/components/checkpoint/state_dict_adapter#nemo_automodel-components-checkpoint-state_dict_adapter-StateDictAdapter)

Convert MiMo-V2-Flash HF checkpoints to Automodel's grouped MoE layout.

HF stores routed experts as split per-expert projections:
`mlp.experts.{E}.{gate,up,down}_proj.weight`.  Automodel groups those
into `gate_and_up_projs` and `down_projs` so EP can shard experts
without materializing every expert on every rank.

**`_uses_fused_qkv_checkpoint`** `bool`

Whether checkpoint tensors use MiMo-V2.6 fused QKV and MXFP4 experts.

---

**`checkpoint_tp_size`** `int`

TP degree interleaving fused QKV rows in the source checkpoint, resolved once.

---

**`hf_to_internal_map`** `dict[str, str] = {}`

---

**`internal_to_hf_map`**

---

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter._apply_key_mapping(
    state_dict: dict[str, typing.Any],
    mapping: dict[str, str]
) -> dict[str, typing.Any]
```

staticmethod

Rename matching state-dict leaves in place without touching prefixes.

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter._convert_fused_qkv(
    state_dict: dict[str, typing.Any]
) -> dict[str, typing.Any]
```

Replace checkpoint fused QKV tensors with native split projections.

**Parameters:**

**`state_dict`** `dict[str, Any]`

Checkpoint tensor mapping. Fused QKV weights have shape
\[checkpoint\_tp\_size \* local\_qkv\_rows, hidden], and optional scale tensors have
shape \[checkpoint\_tp\_size \* local\_scale\_rows, hidden\_blocks].

---

**Returns:** `dict[str, Any]`

The mutated mapping with split query, key, and value weights.

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter._copy_converted_to_native(
    native_tensor: torch.Tensor,
    local_view: torch.Tensor,
    converted: torch.Tensor
) -> None
```

staticmethod

Copy a full checkpoint-layout conversion into local model storage.

**Parameters:**

**`native_tensor`** `torch.Tensor`

Native Tensor or DTensor with the global projection shape.

---

**`local_view`** `torch.Tensor`

Rank-local model storage for `native_tensor`.

---

**`converted`** `torch.Tensor`

Full deinterleaved projection in checkpoint dtype conversion output.

---

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter._create_scale_inv_for_hf_key(
    key: str,
    weight: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter._dequantize(
    state_dict: dict[str, typing.Any]
) -> dict[str, typing.Any]
```

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter._dequantize_mxfp4_experts(
    state_dict: dict[str, typing.Any]
) -> dict[str, typing.Any]
```

Expand packed expert weights into the native floating-point layout.

**Parameters:**

**`state_dict`** `dict[str, Any]`

Checkpoint tensor mapping. MXFP4 expert weights have
shape \[output, input / 2], and their `weight_scale` tensors
have shape \[output, input / 32].

---

**Returns:** `dict[str, Any]`

The mutated mapping with floating-point expert matrices.

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter._local_tensor(
    tensor: torch.Tensor
) -> torch.Tensor
```

staticmethod

Return the rank-local storage view for a Tensor or DTensor.

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter._make_fused_qkv_load_destinations(
    prefix: str,
    layer_idx: int,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor
) -> dict[str, torch.Tensor]
```

Create V2.6 fused-QKV checkpoint buffers for three native weights.

**Parameters:**

**`prefix`** `str`

Projection prefix ending in `self_attn`.

---

**`layer_idx`** `int`

Decoder layer index selecting full or sliding shapes.

---

**`query`** `torch.Tensor`

Native query weight of shape \[query\_rows, hidden].

---

**`key`** `torch.Tensor`

Native key weight of shape \[key\_rows, hidden].

---

**`value`** `torch.Tensor`

Native value weight of shape \[value\_rows, hidden].

---

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

Fused FP8 weight and FP32 inverse-scale destinations using the exact

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter._make_mxfp4_load_destinations(
    key: str,
    weight: torch.Tensor
) -> list[tuple[str, torch.Tensor]]
```

Create packed MXFP4 destinations for one logical expert matrix.

**Parameters:**

**`key`** `str`

Checkpoint expert weight key.

---

**`weight`** `torch.Tensor`

Logical expert view of shape \[output, input].

---

**Returns:** `list[tuple[str, torch.Tensor]]`

Packed E2M1 weight \[output, input / 2] and E8M0 scales

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter._model_to_hf_key(
    model_key: str
) -> str
```

Map one native MiMo tensor name back to its Hugging Face leaf name.

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: typing.Any,
    kwargs = {}
) -> list[tuple[str, typing.Any]]
```

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter.from_hf(
    hf_state_dict: dict[str, typing.Any],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    kwargs = {}
) -> dict[str, typing.Any]
```

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter.to_hf(
    state_dict: dict[str, typing.Any],
    exclude_key_regex: str | None = None,
    quantization: bool = False,
    kwargs = {}
) -> dict[str, typing.Any]
```

Convert native tensors to MiMo checkpoint keys and load destinations.

```python
class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2ForCausalLM()
```

**Bases:** [MiMoV2FlashForCausalLM](#nemo_automodel-components-models-mimo_v2_flash-model-MiMoV2FlashForCausalLM)

MiMo-V2.5/V2.6 wrapper using the `mimo_v2` checkpoint config.

**`ModelCapabilities`** `= MiMoV2FlashForCausalLM.ModelCapabilities`

---

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

classmethod

Resolve a MiMo-V2 config while checkpoint loading remains framework-owned.