Repository navigation
Add Xiaomi MiMo-V2 - #45144
Add Xiaomi MiMo-V2#45144
Conversation
Important points:fixed in: #45441
|
|
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. |
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Some initial comments: I think the biggest point is re attention I'd rather disable the attention types than to support all these paths - it gets too messy
|
Feel free to ping me when it's ready for another review 🤗 |
|
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 🤗 |
|
This comment contains models: ["models/mimo_v2_flash"] |
|
[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 recapDashboard: View test results in Grafana |
|
Congrats on the merge !! 🤗 This was super annoying, CI being flaky and all... |
* init commit with v14 * CI fixes: missing docstring * CI fix: removed Mixtral unused arg * CI fix: shard attn_sinks (TP) * notes * refactor: using copy of gpt-oss eager attn func for sink path huggingface#45141 * v15: polish * fix: exact MiMo layer pattern when fallback * v16: re-inheriting from `Gemma3RotaryEmbedding` * docs: applied suggestion (Qwen3MoE as template) * refactor: re-using attributes pt 1 * refactor: `qk_head_dim` to `head_dim` * refactor: dropping `layernorm_epsilon` mapping, re-use `rms_norm_eps` (same val) * refactor: `MiMoV2FlashForCausalLM` inherits from `DeepseekV3ForCausalLM` * refactor(`super` pt1): rely on parent init / remove redundant overrides * refactor: re-using attributes pt2 * refactor: super pt2 * refactor(MiMoV2FlashPreTrainedModel): removed some flags * refactor(attn): replaced dual attn with eager optional sink * cleaned tests * nit style * add conversion script pt1 * fix: broken config with hub `routed_scaling_factor` hparam * conversion script pt 2 (glm4 style) * refactor: removed redundant init.normal_(module.weight ,...) handled by super * regenerate mapping for the new dynamic auto mapping * docs: removed comment (already explained in PR comment) * cleaning tests * refactor(position_embeddings): keep optim with gemm3 style notation * style * dropped model and tokenizer from convert script * docs: comment fix Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * docs: added comment Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * docs : add suggestion comment Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * applied suggestion: removed ep_plan Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * suggestion: removed no op `self.standardize_rope_params()` Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * refactor: config inherits from glm4moe * refactor: unconditional SWA mask creation Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * fix: missing decorator Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * refactor: sync rot_emb optim with gemmas * refactor 2: unconditional SWA mask creation * refactor: `compute_default_rope_parameters` to match other model patterns * refactor: attn inherits from `Qwen2Attention` * style: switching MLP inheritance for cleaner import * suggestion: switching Experts inheritance to `DeepseekV3NaiveMoe` Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * cleaning imports after Experts suggestion * refactor: MoE rework to inherit DSV3 logic + pretrainedModel from DSV3PreTrainedModel Changes made possible from correct masking huggingface#45441 * removed `test_convert_config_from_hub_format` * suggestion(test): removed dead flag from copy paste Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * refactor: dropping `test_router_group_mask_uses_negative_infinity` since DSV3 is patched * test: cleaned config * test: removing gpt-oss copy paste test * test: fix training kwarg Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * tests: removed additional tests * refactor: renamed eager func to `eager_attention_forward_with_optional_sink` * fix: Values rescaling in attn (MiMo specific) * fix: attn sinks allow grad for training * fix suggestion * fix: dropped FA2 and flex attn * test: added mimo integration test * suggestion: apply batched suggestions Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * suggestion: additional suggestions applied * suggestion: switch `MiMoV2FlashDecoderLayer` inheritance to `Glm4MoeLiteDecoderLayer` * suggestion: always apply Values rescaling * re-enable backends: FA2/3/4 and flex * test: bumped v head dim =16, head_dim = 32 to mimic decoupling * suggestion: test removed rope dict * CI fix * partial revert: see if cause of TP CI fail * partial revert 2: see if cause of TP CI fail * suggestion: update attn flags Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * refactor: revert eager attn func name * refactor: switched `MiMoV2FlashModel` inheritance to `LagunaModel` * test: comment + cuda CC version * make fix-repo * adjust integration test for CI + small nits * style * fix repo * docs: origin of attn values rescaling * added mimo entry in TOKENIZER_MAPPING_NAMES * fix CI: upd date * fix CI: make fix-repo with updated main branch * make fix repo * fix new date format * CI fix: adding back explicit ep_plan + matching DSV3 refactor update * CI fix: matching new DSV3 gate signature * move / copy to internal testing --------- Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> Co-authored-by: vasqu <antonprogamer@gmail.com>
What does this PR do?
Hello, this PR aims to add the MiMo-V2-Flash model to the Transformers library
Fixes #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.Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
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.
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 #45144 (comment) 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):
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)
sqrt(hidden_size / (num_attention_heads * v_head_dim))