You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
--allgather-cp gives wrong log-probs and NaN grads on stock Qwen3 attention #4469
--allgather-cp gives wrong log-probs on stock Qwen3, and the first training step fails with "found NaN in local grad norm". Seen on main (39b8e437ce), with the torch and fused log-prob backends.
Repro: Qwen3-1.7B, Megatron GRPO, --context-parallel-size 2 --allgather-cp, replaying rollouts recorded with sglang. The trainer's mean log-prob of the sampled tokens is about −0.82 (−0.48 with micro-batch size 1), against −0.296 from the rollout. The same run without CP (−0.2968) or with zigzag CP2 (−0.2967) matches the rollout and trains without NaN.
Likely cause (not verified line by line): --allgather-cp changes miles' data layout to one contiguous chunk of the sequence per CP rank, but only the DSA, DeepSeek-V4 and fla/hf_attention model paths read that layout. Megatron's stock attention for Qwen3 still assumes zigzag chunks, so its attention sees the wrong positions.
Suggested fix: reject --allgather-cp at startup for models whose attention doesn't support the contiguous layout, with a message naming the supported ones; or make the stock path support it.
Found while measuring #4228 (which is unaffected: the bug is the same with and without it).
Activity
fix(megatron): reject --allgather-cp when an attention layer assumes …
fix(megatron): reject --allgather-cp when an attention layer assumes …
fix(megatron): reject --allgather-cp when an attention layer assumes …