[Distributed] Add zero-SM all_gather and all-to-all (#79686)
* [Distributed] Add zero-SM all_gather and all-to-all Build the collectives on top of the symmetric-memory window registration: a group created with cta_policy=2 runs them on the Copy Engines and the RMA CPU proxy, so they consume no SM. - add paddle.distributed.zero_sm with empty(), all_gather() and alltoall_single(), mirroring the signatures of paddle.distributed.stream. The output is allocated from the registered pool and any input that is not already pooled is staged into it - take the single-call ncclAlltoAll when the splits are uniform and both buffers are registered: a group of ncclSend/ncclRecv always launches a point-to-point device kernel, while ncclAlltoAll is a symmetric collective and can reach the zero-SM path. Everything else keeps the Send/Recv implementation * [Distributed] Address review on the zero-SM collectives - gate the single-call ncclAlltoAll on the even split instead of on locally uniform split sizes. Explicit splits cannot be judged rank-locally: rank 0 may see [1, 1] while rank 1 legally sees [1, 2], which would send one rank into ncclAlltoAll and the others into Send/Recv and hang the group. With the even split a legal all-to-all forces one count group-wide, so the choice needs no agreement - let the list form of all_gather take a 0-D input, which stream.all_gather answers with one scalar per rank, rather than indexing an empty shape - settle the write-back in is_completed() as well as in wait(), so that a caller polling the task and then reading its list sees the gathered values - drop the per-call votes from _ZeroSMPool: every decision is local now and the caller owes the pool a rank-symmetric call sequence, which is what the docstrings spell out Verified on one node with two GPUs against NCCL 2.30.7: the even-split call still reaches the symmetric path and the uneven one keeps using Send/Recv. * address review opinion * address review opinion * address review opinion * retrigger AI review on current head
A
adam-xiaoyao committed
4ac7a7572b1697b538340e6a38944cd060d791ee
Parent: 80001ff
Committed by GitHub <noreply@github.com>
on 9/9/2026, 1:18:14 PM