[wdykas]

…it deadlock

The NCCL copy service submitted every send/recv of a reshard plan as one batch_isend_irecv, i.e. one ncclGroup with the whole model in it (~15k ops per rank for a 67B MoE refit, sizes from a few KB to hundreds of MB, several peers). NCCL splits such a group into multiple kernel plans by an internal budget, and because the split point depends on each rank's own op list, matching sends and recvs can land in different kernel plans on the two ranks. The resulting device-side dependency cycle hangs the refit forever: pytorch/pytorch#174288

Observed as a NeMo-RL Megatron non-colocated refit (TP8/EP16 -> TP4/EP4, 32 ranks across two racks) never returning, while TP8/EP8 (half the ops per rank) completed. Reproduced standalone on 4 GPUs by replaying the exact plan op list: 24k ops per rank in one group hangs on NCCL 2.30.4 and 2.30.7, 8k ops passes, and the same 24k ops split into task-id windows passes.

Two layers, both deterministic across ranks so no rank can disagree on which submission a transfer belongs to:

• planner: execution batches are now on by default, capped at DEFAULT_MAX_PARAMS_PER_BATCH (32) logical parameters per lockstep batch in addition to the existing byte limit. All shards/replicas of a parameter stay in one batch and every destination rank is served in every batch, so the fleet stays busy. MEGATRON_REFIT_MAX_PARAMS_PER_BATCH overrides; 0 disables. Only copy services with supports_multiple_runs_per_plan (NCCL) consume the batches; gloo/NIXL/M2N behavior is unchanged.
• NCCLCopyService.run(): if a submission still exceeds MAX_OPS_PER_GROUP (256) ops, issue it as consecutive batch_isend_irecv calls over windows of TASK_WINDOW (256) global task ids. Sender and receiver share a transfer's task id, so both place it in the same window.

Validated end to end in NeMo-RL (TP2/EP2 -> TP1/EP1 refit): 241 batches of <=75 ops, transfer 1.8s, total sync time unchanged; windowed fallback 2.1s.

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[copy-pr-bot]

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.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[wdykas]

/ok to test 85d0da4

[cspades]

NCCL splits such a group into multiple kernel plans by an internal budget, and because the split point depends on each rank's own op list, matching sends and recvs can land in different kernel plans on the two ranks

So this is literally a designed race (or I guess... load?) condition in NCCL?

half the ops per rank

An example of an op would be per-parameter, going from a TP8 shard to a TP4 shard, or a EP16 shard to an EP4 shard? Asking cause I also tested various combos, like TP8 EP16 to TP8 EP4 to see if it was because we couldn't do two at a time, and I think it still hung the same way. (I'll test this PR but had to ask.)

[wdykas]

NCCL splits such a group into multiple kernel plans by an internal budget, and because the split point depends on each rank's own op list, matching sends and recvs can land in different kernel plans on the two ranks

So this is literally a designed race (or I guess... load?) condition in NCCL?

half the ops per rank

An example of an op would be per-parameter, going from a TP8 shard to a TP4 shard, or a EP16 shard to an EP4 shard? Asking cause I also tested various combos, like TP8 EP16 to TP8 EP4 to see if it was because we couldn't do two at a time, and I think it still hung the same way. (I'll test this PR but had to ask.)

Written with help from ai but the TLDR is that nccl dumbly cuts up send/recv into groups underneath and that can cause a send an recv to go to different groups. Hate this

NCCL's group launch has a fixed budget per kernel plan (work-struct bytes / ops per plan). When a group exceeds it, NCCL cuts the group into consecutive plans and launches them in stream order, so plan k+1 on a rank can't start until plan k finishes. The cut is made from each rank's own queue, and a sender's queue (say 16 sends to 4 peers, big expert tensors) looks nothing like its receiver's queue (recvs from 7 peers, different sizes). So rank A can put "send to B, op #900" in its plan 1 while B puts "recv from A, op #900" in its plan 2, and B's plan 1 is waiting on something in A's plan 2. Neither plan can finish, so nothing launches. Below the budget there is only one plan per rank and no cycle is possible.

One op is one contiguous slice transfer for one parameter, one (src rank, dst rank) pair. Counts per destination rank in the plan:

• TP-sharded dense param, TP8 into TP4: 2 ops (the dst shard is assembled from 2 source shards). TP8 into TP8: 1 op.
• Expert weight, EP16 into EP4: 1 op per expert per destination replica, but each destination now holds 128 experts instead of 64, from 4 owners instead of 2, and each source sends every expert to 4 replicas instead of 2. This is where most ops live (44 MoE layers times 128 experts times 2 weights, so roughly 11k of the ~15k ops per rank).
• Replicated params (norms, biases, router bias): 1 op each.

That's why your TP8/EP16 into TP8/EP4 test hung too: TP was unchanged, but the EP change alone doubled the expert ops and bytes per rank and raised the peer count from 2 to 3 up to 4 to 5, which is enough to cross the plan budget. It has nothing to do with doing two reshards at once; it's purely how many ops end up in one batch_isend_irecv. My offline replay of the real plan gives, per rank at 16 layers: TP8/EP8 1,477 ops, TP8/EP4 2,373, TP4/EP4 2,714 (multiply by about 5.5 for 89 layers), and the standalone test hangs at 24k ops but passes at 8k on the same NCCL.

[wdykas]

/ok to test 80687c3

[wdykas]

/ok to test b51ffdb

[cspades]

Tested on my script here: https://github.com/cspades/RL/blob/cye/nemotron-omni-super-base/scripts/submit_nemotron_35_super_67b_vision_single_controller_8n4g.sh -> https://wandb.ai/adlr/mllm-rl-dev/runs/4auzuugb?nw=nwusercye_nv

Except with TP4 and EP4 from TP8 and EP16 on training.

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/34874681967