Skip to content

--allgather-cp gives wrong log-probs and NaN grads on stock Qwen3 attention #4469

Open

Description

@HJSang

--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

  1. added 3 commits that reference this issue on Oct 27, 2026
    fafcada
    434c5eb
    04c0d7e
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions