Skip to navigation
DevelopmentAPI ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsMimo V2 Flashnemo_automodel.components.models.mimo_v2_flash.state_dict_adapter

nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter

View as Markdown

Module Contents

Classes

NameDescription
MiMoV2FlashStateDictAdapterConvert MiMo-V2-Flash HF checkpoints to Automodel’s grouped MoE layout.

Functions

NameDescription
_dequantize_mxfp4Dequantize one OCP MXFP4 matrix from the MiMo-V2.6 checkpoint.
_fused_qkv_sizesReturn per-checkpoint-shard Q, K, and V row counts for one layer.
_read_index_tp_sizeRead the fused-QKV storage TP degree recorded in a checkpoint’s safetensors index.
_resolve_checkpoint_tp_sizeReturn the fused-QKV checkpoint TP degree from the index of the checkpoint config was loaded from.
_should_quantize_key-
_split_fused_qkvConvert MiMo’s checkpoint-TP-interleaved fused QKV matrix to split projections.

Data

NON_QUANTIZED_KEY_PATTERNS

_E2M1_VALUES

_FUSED_QKV_WEIGHT

_MXFP4_BLOCK_SIZE

_SAFETENSORS_INDEX

_SPLIT_Q_WEIGHT

logger

API

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

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.

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.

nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter._create_scale_inv_for_hf_key(
    key: str,
    weight: torch.Tensor
) -> torch.Tensor
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter._dequantize(
    state_dict: dict[str, typing.Any]
) -> dict[str, typing.Any]
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.

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.

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

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

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.

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

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

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.

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 | NoneDefaults to 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.
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.

nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._should_quantize_key(
    key: str
) -> bool
 
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],

nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.NON_QUANTIZED_KEY_PATTERNS = ['input_layernorm.weight', 'post_attention_layernorm.weight', 'norm.weight', 'lm...
 
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)
 
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._FUSED_QKV_WEIGHT = re.compile('^(.*\\.layers\\.(\\d+)\\.self_attn)\\.qkv_proj\\.weight$')
 
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._MXFP4_BLOCK_SIZE = 32
 
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._SAFETENSORS_INDEX = 'model.safetensors.index.json'
 
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._SPLIT_Q_WEIGHT = re.compile('^(.*\\.layers\\.(\\d+)\\.self_attn)\\.q_proj\\.weight$')
 
nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.logger = logging.getLogger(__name__)