fix(model): keep full rotary_percent for DeepSeek-V4 MLA rope

[Meirtz]

What

One-line fix: force rotary_percent = 1.0 in the DeepSeek-V4 bridge, plus a regression test.

Why — root cause of the DSv4-Flash SFT NaN (NVIDIA/Megatron-LM#4468)

HF's partial_rotary_factor (0.125) is relative to head_dim (512). The generic partial_rotary_factor → rotary_percent mapping applies it on top of qk_pos_emb_head_dim (64) — which already encodes the rope split (448 nope + 64 rope). The double-applied factor shrinks the rope cos/sin cache to 64 × 0.125 = 8 dims, and then:

• the unfused path silently rotates only 8 of 64 rope dims (no crash — masked correctness issue on every DSv4 bridge run);
• the fused MLA yarn-rope kernel (fused_mla_rope_inplace, called with emb_dim=64) reads cos/sin out of bounds, producing huge garbage activations (q jumps ~5 → ~1e14, captured in-situ) → the iteration-2 found NaN in local forward loss calculation in DeepSeek-V4-Flash SFT, reproduced on both H100 and GB300.

Debug trail (single-variable isolations, in-situ capture, and replay) is in NVIDIA/Megatron-LM#4468. Pretrain never hit it because the mock data path uses fixed-length sequences and the validated pretrain configs didn't exercise the same cache shape.

Validation (GB300, proxy DSv4-Flash SFT, EP4)

Run	rotary_percent	apply_rope_fusion	Result
control	0.125 (status quo)	False	trains, but rotates 8/64 dims
repro	0.125	True	NaN at iteration 2
fixed	1.0	False	clean, 10/10 iters
fixed	1.0	True	clean, matches unfused control to 4 decimals (12.46808→7.554 vs 12.46882→7.5535)

Unit test: mocks the parent provider_bridge to inject rotary_percent=0.125 (what the generic mapping produces) and asserts the DSv4 bridge forces it back to 1.0.

Notes

• Related: [recipe,docs,test] feat: add DeepSeek-V4-Flash SFT recipes, launcher and tests #4131 (DSv4-Flash SFT recipes; independent — different hunks of the same file, no stacking needed). Also unblocks re-enabling apply_rope_fusion for DSv4 SFT ([megatron] feat: add DeepSeek-V4-Flash SFT example (Megatron-Bridge backend) verl-project/verl#6603 ships it off as a workaround).
• A follow-up suggestion for Megatron-Core (tracked in docs: fix broken README model links breaking docs build #4468): fused_mla_rope_inplace should assert cos.shape[-1] == emb_dim instead of reading out of bounds.
• AI-assisted (Claude); debugging, validation runs, and review by the human author.

🤖 Generated with Claude Code

[yaoyu-33]

is this a different definition issue of HF vs. Megatron, if so, how can we avoid in the future? we can do a auto mapping still. or guard with a check

[Meirtz]

Yes — definition mismatch: HF's partial_rotary_factor is relative to head_dim (512×0.125 = the 64 rope dims), while Megatron's rotary_percent applies to qk_pos_emb_head_dim (64), which already is the rope dim — so for MLA it gets applied twice.

For the future, both of your options work: (1) keep the auto-mapping but skip it for MLA providers (multi_latent_attention=True), where the factor is always redundant — happy to send that as a follow-up; (2) guards: fused_mla_rope_inplace asserting cos.shape[-1] == emb_dim (suggested in NVIDIA/Megatron-LM#4468), so the misconfig fails loudly instead of OOB-reading.