Skip to navigation

nemo_automodel.components.models.mimo_v2_flash

View as Markdown

Submodules

Package Contents

Classes

NameDescription
MiMoV2ConfigConfiguration for MiMo-V2.5 and MiMo-V2.6 checkpoints.
MiMoV2FlashConfigConfiguration for XiaomiMiMo/MiMo-V2-Flash.
MiMoV2FlashForCausalLMCausal LM wrapper for MiMo-V2-Flash with Automodel checkpoint adapters.
MiMoV2FlashModelBackbone model for Xiaomi MiMo-V2-Flash.
MiMoV2FlashStateDictAdapterConvert MiMo-V2-Flash HF checkpoints to Automodel’s grouped MoE layout.
MiMoV2ForCausalLMMiMo-V2.5/V2.6 wrapper using the mimo_v2 checkpoint config.
MimoV2ImageProcessorDrop-in image processor for MiMo-V2 checkpoints matching sglang pixels.

Functions

NameDescription
build_mimo_v2_processorBuild the MiMo-V2 chat processor with sglang-aligned image preprocessing.

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 = 129,
sliding_window_size: int | None = None,
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
sliding_window
class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM(
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(
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(
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,
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.

class nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor(
patch_size: int = 16,
temporal_patch_size: int = 2,
merge_size: int = 2,
min_pixels: int = 8192,
max_pixels: int = 8388608,
image_mean: tuple[float, float, float] = (0.485, 0.456, 0.406),
image_std: tuple[float, float, float] = (0.229, 0.224, 0.225),
kwargs = {}
)

Bases: BaseImageProcessor

Drop-in image processor for MiMo-V2 checkpoints matching sglang pixels.

Implements the pixel path of the reference MiMoProcessor: load as RGB (RGBA flattened onto white), cast to 0-255 float, resize with F.interpolate(mode="bilinear", align_corners=False), standardize with mean/std expressed in 0-255 scale, duplicate the frame along the temporal axis, and emit patches in the Qwen2VL merge-block order. Geometry (smart_resize targets and patch order) matches Qwen2VLImageProcessor except for the sub-factor short-side branch described in the module docstring; only resampling and normalization arithmetic differ otherwise.

Inheriting BaseImageProcessor provides the save/reload contract (save_pretrained/from_pretrained) that ProcessorMixin relies on when a checkpoint is written, and is required for registration with AutoImageProcessor (see _register_with_auto_image_processor) so a saved processor directory reloads through AutoProcessor.from_pretrained. Only JSON-serializable constructor arguments may be stored as instance attributes so to_dict() stays serializable.

model_input_names
= ['pixel_values', 'image_grid_thw']
size
dict[str, int]

Pixel bounds in the HF size mapping shape used by token estimators.

nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor.__call__(
min_pixels: int | None = None,
max_pixels: int | None = None,
return_tensors: str | None = 'pt',
kwargs = {}
) -> transformers.BatchFeature

Preprocess images into MiMo-V2 vision-tower inputs.

Parameters:

images
ImageInput | list[ImageInput] | list[list[ImageInput]]

One PIL image / local path, a flat list of them, or a per-sample nested list ([[img1], [img2], ...]) as produced by ProcessorMixin.apply_chat_template on batched conversations. Nested input is flattened one level, preserving order.

min_pixels
int | NoneDefaults to None

Per-call override of the instance pixel bound.

max_pixels
int | NoneDefaults to None

Per-call override of the instance pixel bound.

return_tensors
str | NoneDefaults to 'pt'

Kept for Qwen2VLImageProcessor interface compatibility; outputs are always torch tensors.

**kwargs
Defaults to {}

Ignored, for the same interface compatibility.

Returns: BatchFeature

BatchFeature with pixel_values of shape [total_patches,

nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor._load(
) -> torch.Tensor

Load one image as a float32 tensor of shape [channels, height, width], 0-255 scale.

nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor._pixel_bounds(
images_kwargs: dict | None
) -> tuple[int, int]
nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor._process_one(
min_pixels: int,
max_pixels: int
) -> tuple[torch.Tensor, list[int]]

Run one image through the sglang pixel path.

Parameters:

image
ImageInput

PIL image or local path.

Returns: torch.Tensor

Tuple of patches of shape [grid_tgrid_hgrid_w,

nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor.fetch_images(
) -> PIL.Image.Image | list

Normalize raw image inputs to PIL images.

Required by ProcessorMixin.prepare_inputs_layout (transformers >= 5.18) before the images reach __call__. Mirrors ImageProcessingMixin.fetch_images for the supported input types; RGB conversion stays in __call__.

nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor.get_number_of_image_patches(
height: int,
width: int,
images_kwargs: dict | None = None
) -> int

Number of 16x16 patches for one image, before the merge_size**2 reduction.

Called by Qwen2_5_VLProcessor to expand <|image_pad|> placeholders.

Parameters:

height
int

Source image height in pixels.

width
int

Source image width in pixels.

images_kwargs
dict | NoneDefaults to None

Optional per-call min_pixels/max_pixels overrides.

Returns: int

grid_h * grid_w for the resized image.

nemo_automodel.components.models.mimo_v2_flash.processing.build_mimo_v2_processor(
pretrained_model_name_or_path: str,
kwargs: typing.Any = {}
) -> transformers.ProcessorMixin

Build the MiMo-V2 chat processor with sglang-aligned image preprocessing.

Loads the checkpoint’s stock processor (tokenizer, chat template, video and audio handling stay exactly as released) and replaces only the image processor with MimoV2ImageProcessor, initialized from the checkpoint’s own preprocessor_config.json values.

Parameters:

pretrained_model_name_or_path
str

HF model id or local checkpoint path.

**kwargs
AnyDefaults to {}

Forwarded to AutoProcessor.from_pretrained, e.g. trust_remote_code or max_pixels (recipes use the latter to cap image token counts).

Returns: ProcessorMixin

The stock MiMo-V2 processor whose image_processor is a