Repository navigation
[QUESTION] Support for Xiaomi MiMo-V2-Flash Architecture (Hybrid Attention + Fine-Grained MoE) #2976
Description
Is your feature request related to a problem? Please describe.
I would like to train/fine-tune the Xiaomi MiMo-V2-Flash model using Megatron-LM.
While analyzing its config.json, I noticed it employs a unique architecture that combines Hybrid Attention (mixing Full Attention and Sliding Window Attention) and a Fine-Grained Mixture-of-Experts (MoE) mechanism.
This architecture does not seem to fit into the standard LlamaModel or MixtralModel classes currently available in Megatron-LM core.
Architecture Analysis (from config.json)
The model config reveals several non-standard features that would likely require specific implementation support:
-
Hybrid Attention Strategy:
The config includes ahybrid_layer_pattern(sequence of 0s and 1s) and specific parameters for Sliding Window Attention (SWA), suggesting distinct layer behaviors:hybrid_layer_pattern:[0, 1, 1, 1, 1, 0, ...]sliding_window_size: 128- Separate heads config for SWA:
swa_num_attention_heads,swa_num_key_value_heads.
-
Fine-Grained MoE:
Unlike standard Mixtral (8 experts), this model uses a much larger number of smaller experts:n_routed_experts: 256num_experts_per_tok: 8moe_intermediate_size: 2048moe_layer_freq: Defined per layer (mostly enabled).
-
Non-Standard Head Dimensions:
head_dim: 192v_head_dim: 128 (Value head dimension differs from Query/Key, similar to MLA or specific optimizations).
Describe the solution you'd like
I would like to request support for the MiMoV2FlashForCausalLM architecture in Megatron-LM.
Given the complexity (Hybrid Attention + High Expert Count), could you provide guidance on:
- Is there any existing branch or work-in-progress that supports Hybrid Layer Patterns or Fine-Grained MoE (200+ experts)?
- If I were to implement this, which model class would be the best starting point to inherit from? (e.g., Is
DeepSeekV2support planned? As this shares some similarities with Fine-Grained MoE).
Additional Context
- Model Link: https://huggingface.co/XiaomiMiMo/MiMo-V2-Flash
- Config Snippet:
{ "architectures": ["MiMoV2FlashForCausalLM"], "model_type": "mimo_v2_flash", "n_routed_experts": 256, "num_experts_per_tok": 8, "hybrid_layer_pattern": [0, 1, 1, 1, 1, 0, ...], "head_dim": 192, "v_head_dim": 128 }
Thank you!
Activity
BoxiangW commentedon Jan 21, 2026
@sbhavani, do you know whether we have plan for MiMo-V2-Flash?
sbhavani commentedon Jan 21, 2026
CC @yaoyu-33
beccohov commentedon Mar 29, 2026
Hello, anyone working on that ?
I'd implement support for Xiaomi MiMo-V2-Flash. Please let me know if this task available for contributions.
sbhavani commentedon Mar 30, 2026
@beccohov we'd love to have you work on this! Could you make the implementation and PR in Megatron Bridge?
Megatron Bridge adds the model specific wiring support. You can add checkpoint conversion (e.g. MiMoV2FlashBridge), model configs (e.g. MiMoV2FlashProvider), and training recipes there.
beccohov commentedon Mar 30, 2026
Awesome, thanks @sbhavani!
I'll look at Megatron Bridge first then
beccohov commentedon Apr 5, 2026
I started working on this here. Right now it's very early stage, I expect it will take about 10 days for me to get first solid solution. I'll notify maintainers once the PR will be ready for review. Thanks!
cc @sbhavani
Eisenhower commentedon Jun 1, 2026
Hi @beccohov , thanks for adding MiMo-V2-Flash support.
I’m also working on this model adaptation, currently implemented directly in VeRL rather than Megatron-Bridge. We have already validated it with real training runs on both general-purpose and internal domain data, but we have not migrated the code to Megatron-Bridge yet.
Have you tested this PR with actual end-to-end training workloads? If so, could you share whether the training stability, numerical accuracy, and final training quality are in line with expectations?
If the implementation has already been validated for training and the results look good, we can consider switching our follow-up work to this Megatron-Bridge implementation.
beccohov commentedon Jun 2, 2026
Hi, @Eisenhower -- I tested e2e but not the full model. I took only 1/4 of its layers to fit to my GPUs just to check correctness. So I can't say about training stability, but should be correct.