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

## Module Contents

### Classes

| Name                                                                                                                            | Description                                                             |
| ------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`MiMoV2FlashStateDictAdapter`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-MiMoV2FlashStateDictAdapter) | Convert MiMo-V2-Flash HF checkpoints to Automodel's grouped MoE layout. |

### Functions

| Name                                                                                                                            | Description                                                                                          |
| ------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------- |
| [`_dequantize_mxfp4`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-_dequantize_mxfp4)                     | Dequantize one OCP MXFP4 matrix from the MiMo-V2.6 checkpoint.                                       |
| [`_fused_qkv_sizes`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-_fused_qkv_sizes)                       | Return per-checkpoint-shard Q, K, and V row counts for one layer.                                    |
| [`_read_index_tp_size`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-_read_index_tp_size)                 | Read the fused-QKV storage TP degree recorded in a checkpoint's safetensors index.                   |
| [`_resolve_checkpoint_tp_size`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-_resolve_checkpoint_tp_size) | Return the fused-QKV checkpoint TP degree from the index of the checkpoint `config` was loaded from. |
| [`_should_quantize_key`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-_should_quantize_key)               | -                                                                                                    |
| [`_split_fused_qkv`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-_split_fused_qkv)                       | Convert MiMo's checkpoint-TP-interleaved fused QKV matrix to split projections.                      |

### Data

[`NON_QUANTIZED_KEY_PATTERNS`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-NON_QUANTIZED_KEY_PATTERNS)

[`_E2M1_VALUES`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-_E2M1_VALUES)

[`_FUSED_QKV_WEIGHT`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-_FUSED_QKV_WEIGHT)

[`_MXFP4_BLOCK_SIZE`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-_MXFP4_BLOCK_SIZE)

[`_SAFETENSORS_INDEX`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-_SAFETENSORS_INDEX)

[`_SPLIT_Q_WEIGHT`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-_SPLIT_Q_WEIGHT)

[`logger`](#nemo_automodel-components-models-mimo_v2_flash-state_dict_adapter-logger)

### API

```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
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._dequantize_mxfp4(
    packed_weight: torch.Tensor,
    scale_e8m0: torch.Tensor,
    dtype: torch.dtype
) -> torch.Tensor
```

Dequantize one OCP MXFP4 matrix from the MiMo-V2.6 checkpoint.

**Parameters:**

**`packed_weight`** `torch.Tensor`

Packed E2M1 tensor of shape \[output, input / 2], with
the earlier input element in the low nibble of each byte.

---

**`scale_e8m0`** `torch.Tensor`

Unsigned E8M0 block scales of shape
\[output, input / 32].

---

**`dtype`** `torch.dtype`

Floating-point dtype for the returned training weight.

---

**Returns:** `torch.Tensor`

Dequantized tensor of shape \[output, input].

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._fused_qkv_sizes(
    config: typing.Any,
    layer_idx: int,
    checkpoint_tp_size: int
) -> tuple[int, int, int]
```

Return per-checkpoint-shard Q, K, and V row counts for one layer.

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._read_index_tp_size(
    model_path: str | None,
    revision: str | None = None
) -> int
```

Read the fused-QKV storage TP degree recorded in a checkpoint's safetensors index.

MiMo-V2.6 checkpoints record it as `metadata.tp_size` in `model.safetensors.index.json`; an index without
the field describes a checkpoint saved without TP interleaving. Local directories are read directly; Hugging Face
repository IDs are resolved from the local cache only.

**Parameters:**

**`model_path`** `str | None`

Local checkpoint directory or Hugging Face repository ID.

---

**`revision`** `str | None` — default: None

Optional Hub revision or commit hash for the cache lookup.

---

**Returns:** `int`

`metadata.tp_size`, or 1 when the index records none.

**Raises:**

* `ValueError`: If the index cannot be found or records an invalid TP degree.

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._resolve_checkpoint_tp_size(
    config: typing.Any,
    revision: str | None = None
) -> int
```

Return the fused-QKV checkpoint TP degree from the index of the checkpoint `config` was loaded from.

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._should_quantize_key(
    key: str
) -> bool
```

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._split_fused_qkv(
    weight: torch.Tensor,
    scale_inv: torch.Tensor | None,
    config: typing.Any,
    layer_idx: int,
    checkpoint_tp_size: int,
    dtype: torch.dtype,
    name: str
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Convert MiMo's checkpoint-TP-interleaved fused QKV matrix to split projections.

**Parameters:**

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

Fused tensor of shape \[checkpoint\_tp\_size \* (q\_rows + k\_rows + v\_rows), hidden].

---

**`scale_inv`** `torch.Tensor | None`

Optional FP8 inverse scales of shape
\[checkpoint\_tp\_size \* ceil((q\_rows + k\_rows + v\_rows) / 128), ceil(hidden / 128)].

---

**`config`** `Any`

MiMo model configuration.

---

**`layer_idx`** `int`

Decoder layer index used to select full or sliding dimensions.

---

**`checkpoint_tp_size`** `int`

TP degree whose shards are interleaved in `weight`.

---

**`dtype`** `torch.dtype`

Floating-point dtype for dequantized weights.

---

**`name`** `str`

Checkpoint key used in diagnostics.

---

**Returns:** `torch.Tensor`

Query, key, and value weights with shapes \[q\_total, hidden],

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.NON_QUANTIZED_KEY_PATTERNS = ['input_layernorm.weight', 'post_attention_layernorm.weight', 'norm.weight', 'lm...
```

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
```

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._FUSED_QKV_WEIGHT = re.compile('^(.*\\.layers\\.(\\d+)\\.self_attn)\\.qkv_proj\\.weight$')
```

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._MXFP4_BLOCK_SIZE = 32
```

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._SAFETENSORS_INDEX = 'model.safetensors.index.json'
```

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._SPLIT_Q_WEIGHT = re.compile('^(.*\\.layers\\.(\\d+)\\.self_attn)\\.q_proj\\.weight$')
```

```python
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.logger = logging.getLogger(__name__)
```