Add CTRAN dispatch and arg validation to ncclAlltoAll - #3561
Open
YulunW wants to merge 1 commit into
Open
Conversation
Summary: NCCLX ships two all-to-all entry points with identical signatures but different implementations, sitting ~230 lines apart in the same file: - `ncclAlltoAll` (lowercase `t`, `collectives.cc:117`) — the upstream NVIDIA API, added in NCCL 2.28. Dispatches straight to `ncclEnqueueCheck(ncclFuncAlltoAll)`. No CTRAN, no argument validation. - `ncclAllToAll` (capital `T`, `collectives.cc:346`) — Meta-created before upstream had an all-to-all. Has the CTRAN switch plus argument checks, and falls back to a hand-rolled grouped `baseSend`/`baseRecv` loop. This is the first diff in a stack that de-dups the two down to `ncclAlltoAll`. It brings `ncclAlltoAll` up to parity so it is a drop-in replacement, without touching `ncclAllToAll` yet — both entry points remain fully functional after this diff, so nothing is broken at any point in the stack. Grafted onto `ncclAlltoAll`, mirroring how `ncclAllGather` (`collectives.cc:93`) already layers CTRAN on top of the upstream baseline: - `count == 0` early-out - `SetCudaDevRAII` - `CudaPtrCheck` on `sendbuff` and `recvbuff` - in-place (`sendbuff == recvbuff`) rejection with `ncclInvalidArgument` - the `NCCL_ALLTOALL_ALGO` / `ctranAllToAllSupport` / `ctranAllToAll` dispatch The non-CTRAN fallback stays on the upstream `ncclEnqueueCheck(ncclFuncAlltoAll)` path rather than adopting the Meta send/recv loop. These are closer than they look: `taskAppend` (`enqueue.cc:3044-3051`) already decomposes `ncclFuncAlltoAll` into the same `2 * nRanks` P2P send/recv tasks the Meta loop builds by hand — it never becomes a `ncclTaskColl` and never reaches a device kernel. Staying on the upstream path also propagates `collAPI = ncclFuncAlltoAll` to the net plugin (already whitelisted at `net.cc:218`), emits one `AlltoAll` NVTX range instead of `2N` `Send`/`Recv` ranges, and keeps divergence from upstream small for future rebases. Applied identically to `v2_29` and `v2_30`; the two files are byte-identical here. Differential Revision: D115510404
Contributor
|
@YulunW has exported this pull request. If you are a Meta employee, you can view the originating Diff in D115510404. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
NCCLX ships two all-to-all entry points with identical signatures but different
implementations, sitting ~230 lines apart in the same file:
ncclAlltoAll(lowercaset,collectives.cc:117) — the upstream NVIDIA API,added in NCCL 2.28. Dispatches straight to
ncclEnqueueCheck(ncclFuncAlltoAll).No CTRAN, no argument validation.
ncclAllToAll(capitalT,collectives.cc:346) — Meta-created before upstreamhad an all-to-all. Has the CTRAN switch plus argument checks, and falls back to a
hand-rolled grouped
baseSend/baseRecvloop.This is the first diff in a stack that de-dups the two down to
ncclAlltoAll. Itbrings
ncclAlltoAllup to parity so it is a drop-in replacement, without touchingncclAllToAllyet — both entry points remain fully functional after this diff, sonothing is broken at any point in the stack.
Grafted onto
ncclAlltoAll, mirroring howncclAllGather(collectives.cc:93)already layers CTRAN on top of the upstream baseline:
count == 0early-outSetCudaDevRAIICudaPtrCheckonsendbuffandrecvbuffsendbuff == recvbuff) rejection withncclInvalidArgumentNCCL_ALLTOALL_ALGO/ctranAllToAllSupport/ctranAllToAlldispatchThe non-CTRAN fallback stays on the upstream
ncclEnqueueCheck(ncclFuncAlltoAll)path rather than adopting the Meta send/recv loop. These are closer than they look:
taskAppend(enqueue.cc:3044-3051) already decomposesncclFuncAlltoAllinto thesame
2 * nRanksP2P send/recv tasks the Meta loop builds by hand — it never becomesa
ncclTaskColland never reaches a device kernel. Staying on the upstream path alsopropagates
collAPI = ncclFuncAlltoAllto the net plugin (already whitelisted atnet.cc:218), emits oneAlltoAllNVTX range instead of2NSend/Recvranges,and keeps divergence from upstream small for future rebases.
Applied identically to
v2_29andv2_30; the two files are byte-identical here.Differential Revision: D115510404