nemo_automodel.components.models.mimo_v2_flash.vision

Module Contents

Classes

Name	Description
`MiMoVisionAttention`	Grouped-query vision attention with optional local windows and sinks.
`MiMoVisionBlock`	Pre-normalized MiMo vision transformer block.
`MiMoVisionPatchEmbed`	Convert flattened spatiotemporal image patches to vision tokens.
`MiMoVisionPatchMerger`	Merge each spatial 2x2 group and project it to the text width.
`MiMoVisionRotaryEmbedding`	Two-dimensional rotary frequencies used by the MiMo vision tower.
`MiMoVisionSwiGLUMLP`	SwiGLU feed-forward network used by each vision block.
`MiMoVisionTransformer`	MiMo-V2.6 vision encoder with checkpoint-compatible parameter names.

Functions

Name	Description
`_apply_rotary_pos_emb`	-
`_rotate_half`	-

API

class nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionAttention(
    dim: int,
    num_heads: int,
    num_kv_heads: int,
    head_dim: int,
    use_sinks: bool,
    window_size: int,
    dtype: torch.dtype
)

Bases: Module

Grouped-query vision attention with optional local windows and sinks.

num_kv_groups

= num_heads // num_kv_heads

proj

qkv

= nn.Linear(dim, qkv_dim, bias=True, dtype=dtype)

scaling

= head_dim ** -0.5

sinks

nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionAttention._build_window_mask(
    seq_len: int,
    device: torch.device,
    dtype: torch.dtype
) -> torch.Tensor | None

nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionAttention.forward(
    hidden_states: torch.Tensor,
    cu_seqlens: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    full_attn: bool = False
) -> torch.Tensor

class nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionBlock(
    dim: int,
    intermediate_dim: int,
    num_heads: int,
    num_kv_heads: int,
    head_dim: int,
    hidden_act: str,
    rms_norm_eps: float,
    use_sinks: bool,
    window_size: int,
    dtype: torch.dtype
)

Bases: Module

Pre-normalized MiMo vision transformer block.

attn

mlp

norm1

= nn.RMSNorm(dim, eps=rms_norm_eps, dtype=dtype)

norm2

= nn.RMSNorm(dim, eps=rms_norm_eps, dtype=dtype)

nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionBlock.forward(
    hidden_states: torch.Tensor,
    cu_seqlens: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    full_attn: bool = False
) -> torch.Tensor

class nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionPatchEmbed(
    patch_size: int,
    temporal_patch_size: int,
    in_channels: int,
    embed_dim: int,
    dtype: torch.dtype
)

Bases: Module

Convert flattened spatiotemporal image patches to vision tokens.

proj

nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionPatchEmbed.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

class nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionPatchMerger(
    dim: int,
    context_dim: int,
    spatial_merge_size: int,
    dtype: torch.dtype
)

Bases: Module

Merge each spatial 2x2 group and project it to the text width.

hidden_size

= context_dim * spatial_merge_size ** 2

ln_q

mlp

nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionPatchMerger.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

class nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionRotaryEmbedding(
    dim: int,
    theta: float = 10000.0
)

Bases: Module

Two-dimensional rotary frequencies used by the MiMo vision tower.

_inv_freq_initialized

= not inv_freq.is_meta

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

nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionRotaryEmbedding.forward(
    seqlen: int,
    device: torch.device | None = None
) -> torch.Tensor

class nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionSwiGLUMLP(
    dim: int,
    intermediate_dim: int,
    hidden_act: str,
    dtype: torch.dtype
)

Bases: Module

SwiGLU feed-forward network used by each vision block.

act_fn

= ACT2FN[hidden_act]

down_proj

gate_proj

up_proj

nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionSwiGLUMLP.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

class nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionTransformer(
    config: dict[str, typing.Any] | typing.Any,
    dtype: torch.dtype
)

Bases: Module

MiMo-V2.6 vision encoder with checkpoint-compatible parameter names.

blocks

dtype

fullatt_block_indexes

= set(getattr(config, 'fullatt_block_indexes', []))

merger

patch_embed

rotary_pos_emb

= MiMoVisionRotaryEmbedding(head_dim // 2)

spatial_merge_unit

= spatial_merge_size ** 2

vit_window_attn_types

nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionTransformer._apply_merge_index(
    tensor: torch.Tensor,
    index: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionTransformer._get_window_index(
    grid_thw: torch.Tensor,
    column_major: bool
) -> torch.Tensor

nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionTransformer._rotary_positions(
    grid_thw: torch.Tensor,
    device: torch.device
) -> torch.Tensor

nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionTransformer.forward(
    pixel_values: torch.Tensor,
    grid_thw: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.mimo_v2_flash.vision.MiMoVisionTransformer.init_weights() -> None

Initialize a materialized vision tower for scratch training.

nemo_automodel.components.models.mimo_v2_flash.vision._apply_rotary_pos_emb(
    query: torch.Tensor,
    key: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

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