nemo_automodel.components.models.mimo_v25.state_dict_adapter

Module Contents

Classes

Name	Description
`MiMoV2StateDictAdapter`	Convert MiMo-V2.5-Pro HF checkpoints to Automodel’s grouped MoE layout.

Functions

Name	Description
`_should_quantize_key`	-

Data

NON_QUANTIZED_KEY_PATTERNS

logger

API

class nemo_automodel.components.models.mimo_v25.state_dict_adapter.MiMoV2StateDictAdapter(
    config: typing.Any,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.bfloat16
)

Bases: MoESplitExpertsStateDictMixin, StateDictAdapter

Convert MiMo-V2.5-Pro 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 materialising every expert on every rank.

MiMo-V2.5-Pro stores fused QKV projections as TP-interleaved shards. The adapter dequantizes each checkpoint shard independently, then restores the canonical [Q, K, V] row layout expected by the model.

nemo_automodel.components.models.mimo_v25.state_dict_adapter.MiMoV2StateDictAdapter._dequantize(
    state_dict: dict[str, typing.Any]
) -> dict[str, typing.Any]

nemo_automodel.components.models.mimo_v25.state_dict_adapter.MiMoV2StateDictAdapter._dequantize_interleaved_qkv(
    weight: torch.Tensor,
    scale_inv: torch.Tensor,
    key: str
) -> torch.Tensor

Dequantize and canonicalize a TP-interleaved fused QKV projection.

Parameters:

weight

torch.Tensor

Tensor of shape [interleaved_qkv, hidden]. Axis 0 stores checkpoint TP shards, each laid out as [Q_shard, K_shard, V_shard]. A DTensor may shard either axis; its placements and global shape are preserved in the returned tensor.

scale_inv

torch.Tensor

Tensor of shape [tp * scale_rows_per_shard, scale_columns]. Each checkpoint TP shard owns an independent 128x128 FP8 scale grid.

key

str

Fully qualified weight name containing the decoder layer index.

Returns: torch.Tensor

Tensor of shape [q_rows + k_rows + v_rows, hidden] in canonical

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

Convert one parameter into HF weights and FP8 scale tensors.

Parameters:

fqn

str

Automodel parameter name.

tensor

Any

Parameter tensor with the native layout documented by to_hf; a DTensor retains its global shape and placements.

**kwargs

Defaults to {}

Adapter options, including for_checkpoint_load when allocating destinations for the pretrained checkpoint.

Returns: list[tuple[str, Any]]

Named tensors in HF layout. For checkpoint loading, fused QKV

nemo_automodel.components.models.mimo_v25.state_dict_adapter.MiMoV2StateDictAdapter.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_v25.state_dict_adapter.MiMoV2StateDictAdapter.to_hf(
    state_dict: dict[str, typing.Any],
    exclude_key_regex: str | None = None,
    quantization: bool = False,
    kwargs = {}
) -> dict[str, typing.Any]

Export HF weights or allocate destinations for an FP8 base checkpoint.

Parameters:

state_dict

dict[str, Any]

Native tensors with their original shapes and dtypes. Grouped gate/up weights have shape [experts, hidden, 2 * intermediate] and down weights have shape [experts, intermediate, hidden]. Fused QKV weights have shape [q_rows + k_rows + v_rows, hidden] in canonical [Q, K, V] row order.

exclude_key_regex

str | NoneDefaults to None

Optional pattern for omitted parameter names.

quantization

boolDefaults to False

Allocate FP8 tensors when loading the base checkpoint. Quantized export is unsupported; ordinary export retains precision.

**kwargs

Defaults to {}

Adapter options, including for_checkpoint_load.

Returns: dict[str, Any]

HF tensors preserving input dtypes and canonical QKV order. Experts

nemo_automodel.components.models.mimo_v25.state_dict_adapter._should_quantize_key(
    key: str
) -> bool

nemo_automodel.components.models.mimo_v25.state_dict_adapter.NON_QUANTIZED_KEY_PATTERNS = ['input_layernorm.weight', 'post_attention_layernorm.weight', 'norm.weight', 'lm...

nemo_automodel.components.models.mimo_v25.state_dict_adapter.logger = logging.getLogger(__name__)