Add Xiaomi MiMo-V2
Author: casinca
Pull request: yes
Labels: name: New model

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=45144)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=45144)
<!-- ci-dashboard-badge:end -->

# What does this PR do?

Hello, this PR aims to add the MiMo-V2-Flash model to the Transformers library
Fixes https://github.com/huggingface/transformers/issues/42954

MiMo-V2 is "The last of the OSS SOTAs" that isn't natively supported by the Transformers library (besides Kimi), so I hope we can make this work.

## Code Agent Policy

The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. We are currently bottlenecked by our ability to review and respond to them. As a result, 
**we ask that new users do not submit pure code agent PRs** at this time. 
You may use code agents in drafting or to help you diagnose issues. We'd also ask autonomous "OpenClaw"-like agents
not to open any PRs or issues for the moment.

PRs that appear to be fully agent-written will probably be closed without review, and we may block users who do this
repeatedly or maliciously. 

This is a rapidly-evolving situation that's causing significant shockwaves in the open-source community. As a result, 
this policy is likely to be updated regularly in the near future. For more information, please read [`CONTRIBUTING.md`](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md).

- [x] I confirm that this is not a pure code agent PR.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc
- kernels: @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->

Only pinging HF engineers who were involved in the PRs mentioned below:

Modular Sensei @vasqu , I took your comments from the other PR, to find the best candidates to inherit from. Although see https://github.com/huggingface/transformers/pull/45144#issuecomment-4162843437 for some important points.

@ArthurZucker for gpt-oss similarity and text model

For reference, there were already 2 previous PRs (1 closed and 1 non modular/abandoned):
- https://github.com/huggingface/transformers/pull/42995
- https://github.com/huggingface/transformers/pull/43020

&nbsp;

Afaik, MiMo in the codebase is a bit novel as it's the only hybrid SWA that combines both dual theta RoPE and dual head
dims + attn sinks (like gpt-oss). Most similar models are either one or the other, but not both.  
So modularity wasn't easy mode pass, pass, pass.

I added some verbose NOTES for context/choices, and make the review easier. I'll obviously remove these once things are settled.


## Additional useful info

(copy pasta from my own readme impl)
- First layer is fully dense (GA+FFN, not MoE)
- No shared experts in MoE
- SWA and GA layers have different RoPE theta bases
- SWA and GA layers have a different number of KV groups (GQA)
- Values head dim is decoupled from QK head dim
- Values are rescaled by `sqrt(hidden_size / (num_attention_heads * v_head_dim))`
- Attention sink is only applied for SWA layers
- Partial RoPE (rotating first 33% of the head dim)
- Triple dtype with 1 quant

Comments:
### Important points:

---
fixed in: https://github.com/huggingface/transformers/pull/45441https://github.com/huggingface/transformers/pull/45441

~~`MiMoV2FlashTopKRouter`:
Atm it is not a modular implementation inheriting from the existing `DeepseekV3TopKRouter` in transformers, for two reasons:~~

~~The native DSV3 router uses `masked_fill=0.0` instead of `masked_fill=-inf` when masking non-selected expert groups. The remote DSV3 repo was fixed (see https://huggingface.co/deepseek-ai/DeepSeek-V3-0324/commit/e9b33add76883f293d6bf61f6bd89b497e80e335#d2h-632685) and remote MiMo also follows the fixed version. So I had to follow the correct behavior.~~

~~If it makes sense to change this masking in the repo for DSV3, then I can try to inherit something.~~

~~The class atm combines the 2 split `DeepseekV3TopkRouter` and `DeepseekV3MoE.route_tokens_to_experts()` to serve as a drop-in replacement in `MixtralSparseMoeBlock`, and simply override `self.gate`. This follows the newer pattern used by MiniMax-M2 for fused expert etc...~~

Will try to take a look, likely next week after torch conference. Bad timing with holidays + conference sorry :cry: appreciate the work tho, just as a heads up that things are a bit delayed

> Will try to take a look, likely next week after torch conference. Bad timing with holidays + conference sorry 😢 appreciate the work tho, just as a heads up that things are a bit delayed

Ah true, np, I understand. Thanks for letting me know.

Feel free to ping me when it's ready for another review :hugs: 

Hello @vasqu I've addressed your comments.
Down to `LoC: 379 (modular) vs 637 (generated) - saved 258 LoC (40.5%)`

- We still have the DSV3 masking problem unaddressed, not sure if you discussed this internally, but you can ping me anytime after the merge, I’ll refactor if HF decides to fix it.
- Let me know your thoughts on the gemma3 optim https://github.com/huggingface/transformers/pull/45144#discussion_r3106982629 and if I should open a PR at the same time.
- We saved 1 attribute by deriving SWA kv heads from full attn kv heads by doing x2. I could do the same with `intermediate_size: int = 16384` and derive `moe_intermediate_size: int = 2048` by dividing by 8, if you want.
- Pushed a tiny-mimo to [my hub](https://huggingface.co/casinca/tiny-mimo-v2-flash), for internal testing a `class MiMoV2FlashIntegrationTest(unittest.TestCase)` but not committed as the expected text will differ in your env.

Thanks

Sorry now addressing your comments
1. Yup expecting this to be fixed, but it's not directly blocking here. It could only enable a bit more modular
2. I can't find the comment or GH is bugging - could you copy the text? My intuition tells me it's about RoPE which looks good to me atm
3. Let's keep moe intermediate size 
4. Good enough if you could still push these, I can adjust the values to our CI in the end 

Hope it answers everything, lmk if not!

> Sorry now addressing your comments
> 
>     1. Yup expecting this to be fixed, but it's not directly blocking here. It could only enable a bit more modular
> 
>     2. I can't find the comment or GH is bugging - could you copy the text? My intuition tells me it's about RoPE which looks good to me atm
> 
>     3. Let's keep moe intermediate size
> 
>     4. Good enough if you could still push these, I can adjust the values to our CI in the end
> 
> 
> Hope it answers everything, lmk if not!

Thanks for the feedback @vasqu , for point 2, my bad I marked the discussion as resolved which broke the links, it should be working now  https://github.com/huggingface/transformers/pull/45144#discussion_r3106982629


Oh wow, that is a great catch - so we go over the set instead of all layer types. Dw, about just 2 loc, those are the best fixes! Minimal effort, big gain.

I would prefer a different PR for this and then syncing it here. Can you cc me and @zucchini-nlp on that PR

Arthur's adding MTP support in https://github.com/huggingface/transformers/pull/45618 so I'll need to pay attention since MiMo has 3 modules

👋 @vasqu , I applied your suggestions and replied to questions.

- There was a sudden fix needed in the attention class, Xiaomi fixed it just a few days ago in their remote, concerning the rescaling of Values, the hparam `attention_value_scale` was in the config but never used.  
 I had recently lost remote:native logits parity and after investigating this turned out to be the cause. (I also some left comments as "unresolved" for more details)
- The integration test is copied from Minimax. It's fine on my end, but strings/logits will have to be regenerated.

Thanks

@vasqu  I applied the remaining suggestions, and re-enabled all backends so you can check on your side.  
I left comments about backends unresolved so you can see my new results with hopper gpu. Hope that helps.


Let me know and I'll re-update the attn backend flags accordingly. Thanks.

Nice, kudos to the Xiaomi team for open sourcing MiMo-V2.5-Pro 🙏 (currently 1st OSS, tied with Kimi K2.6 on the AA Intelligence index).

I've delved into the V2.5-Pro architecture a bit, some differences with V2-Flash (not really architectural, mostly hparam tuning):

- QKV linears are fused
- SWA and full attn layers have an equal number of KV heads
- SWA/full attn pattern is slightly changed (not a clean modulo, ie every 6th layers - more like every ~7th/8th) 
- rest are some hparam value changes like rope thetas or attention_value_scale and bump in size

run-slow: mimo_v2_flash

[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/25173542916)

This comment contains `run-slow`, running the specified jobs: 

models: ["models/mimo_v2_flash"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/25173542916)

### Commit Info
| Context | Commit | Description |
|---------|--------|-------------|
| RUN | [0777ea38](https://github.com/huggingface/transformers/commit/0777ea3887b0c9f6b45cb4ada50b52794f354054) | workflow commit (merge commit) |
| PR | [aff3564c](https://github.com/casinca/transformers/commit/aff3564c9d0ec642cfb5c6d5ad55aabfd2c813ae) | branch commit (from PR) |
| main | [d2b4101d](https://github.com/huggingface/transformers/commit/d2b4101d5323afa6279269e23799b359dc715e67) | base commit (on `main`) |

✅ No failing test specific to this PR 🎉 👏 !


The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_45144). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

> LGTM, I just adjusted values for our CI + small nits. Will run slow tests in a second just to confirm
> 
> @casinca would it be possible to open a PR for the model with the converted config (on the hub)? Maybe ideally, we can merge both configs so both remote code and transformers code could run at once; could you check that out?

@vasqu I appreciated your thorough review of this PR and make things as good as possible (ex: I didn't know about Laguna but it was a smart suggestion to inherit from!)

&nbsp;

Concerning the Hub PR, just to make sure I understand correctly: the idea is to keep their existing hparams for BC, and add all the native `transformers` equivalents on top, right?
Ex: if they use in their config `"hybrid_layer_pattern": [
  0,
  1,
  1,
  1` I'll add a new entry on top `"layer_types": ["full_attention", "sliding_attention", "sliding_attention", ...`
  
It can be a bit messy with all the dups no? If that's correct, I’ll try to open a PR (I feel like it'd get more chance of getting merged if it came from an HF staff member tho)

Yes, it would be a franken-merge config lol. But it makes sure we don't break existing workflows. I'm talking internally if we have contacts with the mimo team - maybe we will have a separate -hf repo instead or a dedicated community. I will keep you in the loop. 

For now, the hub PR is the simplest version. Just ping me here or on the PR, I will respond there that it is officially from us :hugs: (my hf: https://huggingface.co/AntonV)

Mind that it is golden week I think, so earliest response might be at the end of next week but unsure. Talking about internal response 

Hi @vasqu I opened the [PR](https://huggingface.co/XiaomiMiMo/MiMo-V2-Flash/discussions/41), not sure if tagging works so prefer to mention it here too.
Basically I added what's in my `convert_mimo_v2_flash_weights_to_hf.py`. Besides pure additions, I only changed their `routed_scaling_factor` from `null` to `1`, it's the convention in `transformers` for default and will be a no-op on their side too, so BC is good.

Yep, responded on the hub PR as well. Thanks! This really looks like fairly minimal changes

Some small update: We have internal contact with the team, they apparently run some tests and merge if everything is going well :hugs: @casinca 

> Some small update: We have internal contact with the team, they apparently run some tests and merge if everything is going well 🤗 @casinca

That's good to hear, hope it'll turn out fine. Thanks for keeping me up to date.

Hey @casinca, just wanted to update that I still have no news yet sadly. Just wanted to let you know that I haven't forgot about this! 

I try to ping in some other ways; if there is no response, I think we are fine to merge as is for now

tyvm for not forgetting about this and following up on it @vasqu 🙏 I’ll wait on what you decide.

@casinca Ok atp, I don't expect any response soon. Can you sync with main and the newest changes; I think we can merge afterwards and keep the revision for now :hugs: 

Updated with the recent changes, CI should be good now? Not sure that Mistral tokenizer failure is related.

Ye that's disappointing, even the V2.5 follows the same old config patterns. In any case you can ping me anytime, if you have any news I'll do a follow-up if need be @vasqu thanks again.

run-slow: mimo_v2_flash

@casinca just moved/copied the dummy model to internal testing, sanity checking with run slow but then merging :hugs: 

[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/28517464589)

This comment contains `run-slow`, running the specified jobs: 

models: ["models/mimo_v2_flash"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/28517464589)

### Commit Info
| Context | Commit | Description |
|---------|--------|-------------|
| RUN | [1df57d3f](https://github.com/huggingface/transformers/commit/1df57d3f9905b1383b300b7e8306e2207700de18) | workflow commit (merge commit) |
| PR | [97432842](https://github.com/casinca/transformers/commit/97432842e9530b79f4cf0546f162da43405ba2ac) | branch commit (from PR) |
| main | [36778030](https://github.com/huggingface/transformers/commit/36778030cba4bccfcb15366d7ad860ca30f47417) | base commit (on `main`) |

✅ No failing test specific to this PR 🎉 👏 !


**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, mimo_v2_flash

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, mimo_v2_flash

<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=45144)
**Latest run:** [28524616411:3](https://github.com/huggingface/transformers/actions/runs/28527982269)
**Result:** `failure` | **Jobs:** 1 | **Tests:** 1,134 | **Failures:** 0 | **Duration:** 36m 28s

<!-- ci-dashboard-recap:end -->

Congrats on the merge !! :hugs: This was super annoying, CI being flaky and all...