nemo_automodel.components.models.mimo_v2_flash
nemo_automodel.components.models.mimo_v2_flash
Submodules
nemo_automodel.components.models.mimo_v2_flash.confignemo_automodel.components.models.mimo_v2_flash.cpnemo_automodel.components.models.mimo_v2_flash.modelnemo_automodel.components.models.mimo_v2_flash.parallelizationnemo_automodel.components.models.mimo_v2_flash.processingnemo_automodel.components.models.mimo_v2_flash.state_dict_adapternemo_automodel.components.models.mimo_v2_flash.vision
Package Contents
Classes
Functions
API
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.
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.
Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin
Causal LM wrapper for MiMo-V2-Flash with Automodel checkpoint adapters.
Encode global media and splice only this TE THD shard’s feature rows.
Map local TE THD placeholders to rows in globally ordered media features.
Retrieve PP media using global masks so every CP rank runs the same vision work.
Keep both rotary embeddings per stage and the vision tower on stage zero.
Run text or VLM forward with TE-local packed media placement.
Describe PP tensors using the already TE-sharded local token length.
Install MiMo’s model-local adapter around framework TE THD sharding.
Bases: Module
Backbone model for Xiaomi MiMo-V2-Flash.
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.
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.
Whether checkpoint tensors use MiMo-V2.6 fused QKV and MXFP4 experts.
TP degree interleaving fused QKV rows in the source checkpoint, resolved once.
Rename matching state-dict leaves in place without touching prefixes.
Replace checkpoint fused QKV tensors with native split projections.
Parameters:
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.
Copy a full checkpoint-layout conversion into local model storage.
Parameters:
Native Tensor or DTensor with the global projection shape.
Rank-local model storage for native_tensor.
Full deinterleaved projection in checkpoint dtype conversion output.
Expand packed expert weights into the native floating-point layout.
Parameters:
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.
Return the rank-local storage view for a Tensor or DTensor.
Create V2.6 fused-QKV checkpoint buffers for three native weights.
Parameters:
Projection prefix ending in self_attn.
Decoder layer index selecting full or sliding shapes.
Native query weight of shape [query_rows, hidden].
Native key weight of shape [key_rows, hidden].
Native value weight of shape [value_rows, hidden].
Returns: dict[str, torch.Tensor]
Fused FP8 weight and FP32 inverse-scale destinations using the exact
Create packed MXFP4 destinations for one logical expert matrix.
Parameters:
Checkpoint expert weight key.
Logical expert view of shape [output, input].
Returns: list[tuple[str, torch.Tensor]]
Packed E2M1 weight [output, input / 2] and E8M0 scales
Map one native MiMo tensor name back to its Hugging Face leaf name.
Convert native tensors to MiMo checkpoint keys and load destinations.
Bases: MiMoV2FlashForCausalLM
MiMo-V2.5/V2.6 wrapper using the mimo_v2 checkpoint config.
Resolve a MiMo-V2 config while checkpoint loading remains framework-owned.
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.
Pixel bounds in the HF size mapping shape used by token estimators.
Preprocess images into MiMo-V2 vision-tower inputs.
Parameters:
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.
Per-call override of the instance pixel bound.
Per-call override of the instance pixel bound.
Kept for Qwen2VLImageProcessor interface
compatibility; outputs are always torch tensors.
Ignored, for the same interface compatibility.
Returns: BatchFeature
BatchFeature with pixel_values of shape [total_patches,
Load one image as a float32 tensor of shape [channels, height, width], 0-255 scale.
Run one image through the sglang pixel path.
Parameters:
PIL image or local path.
Returns: torch.Tensor
Tuple of patches of shape [grid_tgrid_hgrid_w,
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__.
Number of 16x16 patches for one image, before the merge_size**2 reduction.
Called by Qwen2_5_VLProcessor to expand <|image_pad|> placeholders.
Parameters:
Source image height in pixels.
Source image width in pixels.
Optional per-call min_pixels/max_pixels overrides.
Returns: int
grid_h * grid_w for the resized image.
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:
HF model id or local checkpoint path.
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