Repository navigation
Fix MTP CP-boundary roll to communicate on the context-parallel process group - #6746
Conversation
…ss group Signed-off-by: Li Ding <liding@nvidia.com>
|
/ok to test 360e152 |
|
/claude fix |
Signed-off-by: svcnvidia-nemo-ci <svcnvidia-nemo-ci@nvidia.com>
|
🛠️ Claude fix commit
What changed Files changed by Claude
Why DCO Sanitized and posted by |
|
/ok to test bb6d5b2 |
|
❌ Claude fix stopped because exact-SHA CI did not complete in time. View exact-SHA CI. |
|
/ok to test 0fed350 |
|
/ok to test b11538c |
|
|
||
|
|
||
| class TestRollTensorWithCPSubgroup: | ||
| """Roll-tensor CP paths on a CP group that is a strict subgroup of WORLD. |
There was a problem hiding this comment.
The existing packed-sequence test uses TP1 × CP2 and checks numerical correctness without checking communicator selection. When run with two ranks, CP and WORLD contain the same ranks, hiding the communicator mismatch from numerical checks.
This test uses TP2 × CP2 to guarantee CP is a strict subgroup of WORLD, and explicitly checks that boundary exchanges use cp_group. The explicit group assertion detects the regression
| full = torch.arange(2 * 16, dtype=torch.float32, device="cuda").reshape(2, 16) | ||
| # Oracle: the documented single-rank semantics on the full sequence. | ||
| expected_full, _ = roll_tensor(full.clone(), shifts=-1, dims=-1) | ||
| expected_local = self._cp_shard(expected_full, cp_rank, cp_size) |
There was a problem hiding this comment.
We build the reference by rolling the full sequence without distributed communication, then extracting this rank’s CP shard.
The actual path shards the input first, then rolls it with cross-rank communication. Both paths should produce identical local values.
| rolled, rolled_sum = roll_tensor(local, shifts=-1, dims=-1, cp_group=cp_group) | ||
|
|
||
| assert torch.equal(rolled, expected_local) | ||
| assert rolled_sum == expected_local.sum() |
There was a problem hiding this comment.
These check the local output values—including exchanged boundary tokens and zero padding—and the returned local sum. They verify numerical behavior is preserved. They can also pass before the fix because the old implementation communicates with the correct peers through WORLD.
svcnvidia-nemo-ci
added
Approved
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/35040222266 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/35043683828 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/35044129580 |
What does this PR do?
Fixes MTP's CP-boundary roll (
roll_tensor/_roll_tensor_packed_seq) to communicate on the context-parallel process group instead of the default (world) process group.Details
Both roll variants compute their neighbor ranks from
cp_groupbut issueisend/irecvwithoutgroup=, so every CP boundary exchange actually runs on the default process group. When the world equals the CP group (typical unit tests) this is invisible. Whenever the CP group is a strict subgroup of the world (TP > 1 and/or multiple DP replicas), all CP rings in the job share one communicator:_roll_tensor_packed_seqposts the exchange per packed sequence per microbatch.The fix addresses the ops to
cp_groupand batches each rank's sends/receives withbatch_isend_irecv.Testing
TestRollTensorWithCPSubgroup) run both roll variants on a CP group that is a strict subgroup of the world (TP2 × CP2), asserting numeric equivalence with the single-rank semantics and that the exchange is addressed to the CP group (the latter fails on the pre-fix code).test_multi_token_prediction.pyrun before/after the fix on 4 GPUs produces identical results (failing cases on 4 GPUs are pre-existing world-size/fp8 environment mismatches, unchanged by this PR).tools/autoformat.shrun; black/isort/pylint/ruff clean (mypy findings are pre-existing, outside the changed lines).