nemo_automodel.components.models.mimo_v2_flash.model

Module Contents

Classes

Name	Description
`MiMoV2FlashAttention`	MiMo-V2 attention with a Transformer Engine THD context-parallel path.
`MiMoV2FlashBlock`	Decoder block that alternates dense MLP and routed-MoE layers.
`MiMoV2FlashForCausalLM`	Causal LM wrapper for MiMo-V2-Flash with Automodel checkpoint adapters.
`MiMoV2FlashModel`	Backbone model for Xiaomi MiMo-V2-Flash.
`MiMoV2FlashRotaryEmbedding`	Rotary embedding module matching MiMo-V2-Flash partial-RoPE behavior.
`MiMoV2ForCausalLM`	MiMo-V2.5/V2.6 wrapper using the `mimo_v2` checkpoint config.
`MiMoV2RMSNorm`	RMSNorm used by MiMo-V2-Flash decoder blocks.

Functions

Name	Description
`_apply_rotary_pos_emb`	-
`_convert_bool_4d_mask_to_additive`	-
`_derive_padding_mask`	-
`_eager_attention_forward`	-
`_ensure_additive_mask`	-
`_fallback_additive_mask`	-
`_initialize_mimo_rms_norm`	Construct a MiMo RMSNorm from the configured backend.
`_intersect_with_sliding_window`	Apply causal sliding-window bounds to an existing 4D document mask.
`_normalize_image_grid`	Accept the THW grids emitted by Qwen processors and PP’s HW alias.
`_repeat_kv`	-
`_replace_modal_embeddings`	Replace one kind of multimodal placeholder while preserving gradients.
`_rotate_half`	-
`_validate_te_thd_sink_cudnn`	Reject the known pre-9.26 cuDNN ragged dSink kernel bug.

Data

ModelClass

API

class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention(
    config: nemo_automodel.components.models.mimo_v2_flash.config.MiMoV2FlashConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    is_swa: bool,
    layer_idx: int
)

Bases: Module

MiMo-V2 attention with a Transformer Engine THD context-parallel path.

_te_v_head_dim

= self.v_head_dim

attention_dropout

= float(config.attention_dropout or 0.0)

head_dim

= int(config.swa_head_dim)

k_proj

num_attention_heads

= int(config.swa_num_attention_heads)

num_key_value_groups

num_key_value_heads

= int(config.swa_num_key_value_heads)

o_proj

q_proj

rope_dim

= int(self.head_dim * config.partial_rotary_factor)

scaling

= self.head_dim ** -0.5

sliding_window

= int(config.sliding_window) if is_swa else None

v_head_dim

= int(config.swa_v_head_dim)

v_proj

v_scale

= getattr(config, 'attention_value_scale', None)

nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention._apply_te_rope(
    query_states: torch.Tensor,
    key_states: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    qkv_format: str
) -> tuple[torch.Tensor, torch.Tensor]

staticmethod

Apply partial RoPE to BSHD or flattened THD query/key tensors.

nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention._forward_te(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None,
    kwargs: typing.Any = {}
) -> torch.Tensor

Run TE attention with native V dimensions, padding for legacy TE a2a CP.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden] for BSHD, or [tokens, hidden] for THD; a singleton batch is also accepted.

position_embeddings

tuple[torch.Tensor, torch.Tensor]

Cosine and sine tensors aligned to the input sequence, each ending in the rotary feature dimension.

attention_mask

torch.Tensor | None

Optional mask of shape [batch, sequence] for BSHD.

**kwargs

AnyDefaults to {}

TE attention metadata, including one-dimensional cumulative sequence-length tensors for THD inputs.

Returns: torch.Tensor

Attention output with the same shape as hidden_states.

nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention._validate_a2a_cp_size(
    cp_size: int
) -> None

Validate the head partition required by TE’s a2a CP transport.

nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor | None]

nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention.setup_cp_attention(
    cp_mesh,
    cp_stream: torch.cuda.Stream | None = None
) -> None

Validate and attach TE all-to-all context parallelism.

class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashBlock(
    layer_idx: int,
    config: nemo_automodel.components.models.mimo_v2_flash.config.MiMoV2FlashConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

Decoder block that alternates dense MLP and routed-MoE layers.

attention_type

input_layernorm

mlp

= MoE(moe_config, backend)

post_attention_layernorm

self_attn

nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashBlock.forward(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    padding_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor

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

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.model.MiMoV2FlashRotaryEmbedding(
    rope_theta: float,
    head_dim: int,
    partial_rotary_factor: float = 1.0,
    dtype: torch.dtype = torch.bfloat16
)

Bases: Module

Rotary embedding module matching MiMo-V2-Flash partial-RoPE behavior.

_inv_freq_initialized

= not inv_freq.is_meta

attention_scaling

= 1.0

inv_freq

Tensor

nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashRotaryEmbedding._build_inv_freq(
    device: torch.device | None = None
) -> torch.Tensor

Build the FP32 inverse frequencies on the requested device.

nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashRotaryEmbedding._materialize_inv_freq(
    device: torch.device
) -> None

Restore a non-persistent RoPE buffer after meta-model materialization.

nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashRotaryEmbedding.forward(
    x: torch.Tensor,
    position_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

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.

class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2RMSNorm(
    hidden_size: int,
    eps: float = 1e-06,
    dtype: torch.dtype = torch.bfloat16
)

Bases: Module

RMSNorm used by MiMo-V2-Flash decoder blocks.

weight

= nn.Parameter(torch.ones(hidden_size, dtype=dtype))

nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2RMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2RMSNorm.reset_parameters() -> None

nemo_automodel.components.models.mimo_v2_flash.model._apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.models.mimo_v2_flash.model._convert_bool_4d_mask_to_additive(
    attention_mask: torch.Tensor,
    dtype: torch.dtype
) -> torch.Tensor

nemo_automodel.components.models.mimo_v2_flash.model._derive_padding_mask(
    attention_mask: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.mimo_v2_flash.model._eager_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    scaling: float,
    dropout: float = 0.0
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.models.mimo_v2_flash.model._ensure_additive_mask(
    mask: torch.Tensor | None,
    batch_size: int,
    seq_len: int,
    dtype: torch.dtype,
    device: torch.device,
    attention_mask: torch.Tensor | None,
    sliding_window: int | None
) -> torch.Tensor

nemo_automodel.components.models.mimo_v2_flash.model._fallback_additive_mask(
    batch_size: int,
    seq_len: int,
    dtype: torch.dtype,
    device: torch.device,
    attention_mask: torch.Tensor | None = None,
    sliding_window: int | None = None
) -> torch.Tensor

nemo_automodel.components.models.mimo_v2_flash.model._initialize_mimo_rms_norm(
    backend: nemo_automodel.components.models.common.BackendConfig,
    hidden_size: int,
    eps: float,
    dtype: torch.dtype
) -> torch.nn.Module

Construct a MiMo RMSNorm from the configured backend.

nemo_automodel.components.models.mimo_v2_flash.model._intersect_with_sliding_window(
    attention_mask: torch.Tensor,
    sliding_window: int,
    dtype: torch.dtype
) -> torch.Tensor

Apply causal sliding-window bounds to an existing 4D document mask.

nemo_automodel.components.models.mimo_v2_flash.model._normalize_image_grid(
    grid: torch.Tensor | None
) -> torch.Tensor | None

Accept the THW grids emitted by Qwen processors and PP’s HW alias.

nemo_automodel.components.models.mimo_v2_flash.model._repeat_kv(
    hidden_states: torch.Tensor,
    n_rep: int
) -> torch.Tensor

nemo_automodel.components.models.mimo_v2_flash.model._replace_modal_embeddings(
    input_ids: torch.Tensor,
    inputs_embeds: torch.Tensor,
    token_id: int | None,
    modal_embeds: torch.Tensor | None
) -> torch.Tensor

Replace one kind of multimodal placeholder while preserving gradients.

nemo_automodel.components.models.mimo_v2_flash.model._rotate_half(
    x: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.mimo_v2_flash.model._validate_te_thd_sink_cudnn(
    head_dim: int
) -> None

Reject the known pre-9.26 cuDNN ragged dSink kernel bug.

nemo_automodel.components.models.mimo_v2_flash.model.ModelClass = MiMoV2FlashForCausalLM