This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Drafted this WIP PR for MiMoV2 Flash support according to this issue

@sbhavani Hey,
MiMo-V2-Flash uses asymmetric attention head dimensions: Q/K use head_dim=192 but V uses v_head_dim=128. This is currently not supported by MCore's standard SelfAttention for two reasons:

1. linear_proj input size — Attention.__init__ sizes linear_proj input as kv_channels * num_attention_heads = 192 * 64 = 12288 (here), but the correct size is v_head_dim * num_attention_heads = 128 * 64 = 8192.
2. QKV split in forward pass — get_query_key_value_tensors splits the fused QKV output using hidden_size_per_attention_head (= kv_channels = 192) for both K and V, so V would be extracted with the wrong size (here).

Should I expose support similar to MLA but in standard SelfAttention for decoupled V size in megatron? Alternatively, is there a recommended pattern for bridging models with asymmetric V dims without modifying MCore?

Should I expose support similar to MLA but in standard SelfAttention for decoupled V size in megatron? Alternatively, is there a recommended pattern for bridging models with asymmetric V dims without modifying MCore?

@beccohov : usually we patch the config / implementation / fwd in bridge directly and use custom layer spec to specify the customized version of it.

Here's the GitHub link to the OLMoE provider file:
https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/olmoe/olmoe_provider.py
Layer spec builder — L41-45
Custom OLMoESelfAttention class — L98-188

Hey, @yaoyu-33, I have a few questions before I mark this PR as "ready for review":

1. Currently my implementation of MiMo-V2-Flash asserts CP because TE lacks a backend for CP + learnable softmax (attention sink bias). SWA layers use softmax_type="learnable" which is refused by all TE backends when CP is enabled. However, seems like it is relatively easy to fix in TE (here we'll go via MLA branch and we have v.shape wrong because of GQA). I tried to patch it in provider, but there are lot's of TE assertions that strictly control ability of hacks. Is asserting acceptable, or should we implement a workaround ?

2. The HF config has attention_value_scale=0.707, but the released HF modeling code does not apply it in the forward pass (i.e. the released weights rescale V back). We currently skip it too (matching HF behavior). For training, however, I think we should re-enable the scale (i.e. to do value = value * 0.707) to match MiMo-v2 Flash fully. What do you think ? This is basically minor fix and I think it plays role only for large scale pretraining with fp8. For finetunes I think it's better to skip this scaling, because naive implementation would consume more memory on activations.

3. Should mimo_v2_flash/ be merged into the existing mimo/ directory? The existing mimo/ contains both the original Xiaomi MiMo bridge (Qwen2-based) and the unrelated multimodal MIMO provider. I kept it separate to avoid confusion, but happy to merge if preferred.

@beccohov I'd recommend creating an issue in TE to track CP support. I think the incomplete CP coverage for non-vanilla softmax should be fixed eventually.

Created issue in TE.
Apart from this issue, what about v_scale and directory ?

The TE issue with CP is fixed in TE2.15.

@sbhavani hey, would it be possible to have review please?
Thanks!

CC @snowmanwwg

@sbhavani the label is removed automatically anyways for some reason

@beccohov thanks! we need to fix that automation bug

CC @yaoyu-33

@cuichenx I believe I've addressed all your comments — would appreciate a re-review!

@beccohov We are working on caching this model on CI server. I will re-run the failed CI test, when the fix is in place. No action needed from your side. Thanks.

@kamran-nvidia can you please share the PR / something where I can track the progress, so that I can understand when I'll be able to finish this PR?

@JRD971000 Can you help @beccohov please? Thanks.

/ok to test 017062b

Afaiu, the tests failed due to (2) from my comment here. So what should I do? Should I add trust_remote_code ?
cc @yaoyu-33 @kamran-nvidia

@beccohov nvm, I will force merge this. It's ci caching issue.

/ok to test 341094f

will follow up add back the functional tests after caching issue resolved.

Hi @beccohov @yaoyu-33 — small follow-up while reading this merged work.

MiMoV2FlashTEDotProductAttention.__init__ reads attention_value_scale from config and stores it as self._attention_value_scale (modeling_mimo_v2_flash.py#L256), but the value is never consumed in forward (L270-L271). For the released XiaomiMiMo/MiMo-V2-Flash checkpoint (attention_value_scale = 0.707) this silently scales the attention output by ~1.414× at every layer relative to the reference implementations.

The scale can't be folded into softmax_scale (softmax is non-linear) and linear_proj.weight is mapped 1:1 from HF o_proj.weight, so it has to land on V (or equivalently on the attention output). Both upstream references do this on V before the kernel:

• HF: https://huggingface.co/XiaomiMiMo/MiMo-V2-Flash/blob/main/modeling_mimo_v2_flash.py — value_states = value_states * self.v_scale
• vLLM: https://github.com/vllm-project/vllm/blob/e68988a24807c9dfb2bf6936eb17425ce7812c5f/vllm/model_executor/models/mimo_v2.py#L327-L329 — v = v * self.v_scale

I opened a two-line follow-up: #4155. Would appreciate a quick look when you have a moment — happy to add a focused unit test if that helps the review.

Hi @Eisenhower,
thanks for your comment. Actually, the reason why I did not scale is that HF weights seem to be already static scaled in advance (V weights). When I was implementing this I noticed that by quality degradation.
So you don't need to scale this to reproduce results if you load from retrained model because it was already trained with suitable magnitudes.
However, for training dynamics you may want to scale (this will affect gradients magnitude for V and also if you retrain from scratch this may change magnitude of attention scores). But naive scaling will increase memory consumption (and this would be notable on this large model) so ideally you need to implement this scaling to avoid memory overhead.

So, one way is no upscale back HF weights while loading them and then use this scaling in training. This is especially valuable when you train from scratch.
Another way is to skip this scaling completely (as in my implementation) and keep in mind that gradients for V would be slightly different. Since we use init from pertained model, I believe this won't significantly change model convergence (since we use smaller lr and model is more stable for post-training). But will save you memory (if you don't use full AC).

Thanks for the context. I think the HF implementation you checked was likely the old version. I raised this issue with the Xiaomi MiMo team afterwards, and they have since updated the HF modeling code to apply attention_value_scale on value_states in forward. vLLM also applies the same scale before attention.

I agree the extra allocation from a naive value = value * scale is worth optimizing. My concern is mainly RL consistency: rollout usually follows the vLLM/HF inference path, while training/logprob/KL are computed by the actor path. If rollout applies the scale but training does not, the two forward paths are no longer equivalent. For that reason I would prefer Megatron-Bridge to match the HF/vLLM semantics first, and then optimize the implementation or fold the scale only if we can keep the two paths numerically equivalent.

@beccohov Thanks for the context. I think the HF implementation you checked was likely the old version. I raised this issue with the Xiaomi MiMo team afterwards, and they have since updated the HF modeling code to apply attention_value_scale on value_states in forward. vLLM also applies the same scale before attention.

I agree the extra allocation from a naive value = value * scale is worth optimizing. My concern is mainly RL consistency: rollout usually follows the vLLM/HF inference path, while training/logprob/KL are computed by the actor path. If rollout applies the scale but training does not, the two forward paths are no longer equivalent. For that reason I would prefer Megatron-Bridge to match the HF/vLLM semantics first, and then optimize the implementation or fold the scale only if we can keep the two paths numerically equivalent.

@Eisenhower Sure, if HF weights will be updated then the situation changes. You can then add a follow-up PR with this small change!