nemo_automodel.components.models.mimo_v2_flash

Submodules

Package Contents

Classes

Name	Description
`MiMoV2Config`	Configuration for MiMo-V2.5 and MiMo-V2.6 checkpoints.
`MiMoV2FlashConfig`	Configuration for XiaomiMiMo/MiMo-V2-Flash.
`MiMoV2FlashForCausalLM`	Causal LM wrapper for MiMo-V2-Flash with Automodel checkpoint adapters.
`MiMoV2FlashModel`	Backbone model for Xiaomi MiMo-V2-Flash.
`MiMoV2FlashStateDictAdapter`	Convert MiMo-V2-Flash HF checkpoints to Automodel’s grouped MoE layout.
`MiMoV2ForCausalLM`	MiMo-V2.5/V2.6 wrapper using the `mimo_v2` checkpoint config.

API

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

Bases: 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'

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

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

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.

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.

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.

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.

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.

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

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

classmethod

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

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

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.

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

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.

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

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

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

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]

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.

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

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.

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

Bases: MiMoV2FlashForCausalLM

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

ModelCapabilities

= MiMoV2FlashForCausalLM.ModelCapabilities

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.