ReduceScatter and AllGather decomposition

For p ranks, split an N-element input into p equal blocks under a common layout. ReduceScatter reduces corresponding input elements across all ranks and leaves one reduced block with each rank. AllGather can then collect those reduced blocks in rank order so every rank holds the complete reduced tensor. Under compatible reduction, partitioning, and layout rules, that composition has the logical effect of AllReduce. It is not a promise of identical floating-point rounding or performance to a separately selected AllReduce implementation.

NCCL 2.32.3's ReduceScatter receive count denotes the per-rank reduced block, while each input buffer contains p such blocks. Its AllGather send count denotes one rank's contribution and the receive buffer holds the p-rank concatenation. Do not reuse one numeric count without checking which side of the operation it describes. MPI provides both equal-block and variable-count interfaces with their own rules. Zero or uneven partitions, padding, and alignment must be handled explicitly; a library cannot infer a missing element from a buffer that happens to be large enough.

A sharded optimizer may update its assigned parameter state using its reduced gradient shard, retaining less replicated gradient or optimizer memory during that phase. Which tensors need gathering afterwards depends on the training algorithm, parameter partition, and next computation. Do not assume that splitting a collective automatically shards every parameter, activation, or optimizer state. A complete memory plan records which buffers remain full, which are partitioned, and when temporary gathered views overlap with the retained shards.

Rank ownership is an explicit indexing contract. A change in communicator rank order can change which tensor slice a rank receives or where its AllGather contribution appears. Checkpoint metadata and optimizer ownership must agree with that mapping after restart or resharding. A buffer with the right number of elements but the wrong slice can silently corrupt training without causing a count mismatch. Validate values and offsets using a small identifiable tensor before benchmarking an anonymous array of random numbers.

Decomposition can support overlap or memory savings when useful shard-local work occurs between phases. It can also add launch overhead, force a less efficient algorithm choice, or create extra buffers if applied without a dependency plan. The correct question is which consumer needs a full tensor and at what time. Preserve stream dependencies from reduction completion through shard consumption to gathering, and do not overwrite a shard while an AllGather still reads it. A memory estimate taken between phases must not be advertised as the peak for the whole training step.

The modular lab shows an ownership-order contrast only. Reset, set both m=8 and a=1, predict residue one after one Step and period eight, then Step. Separately Reset, set both m=8 and a=3, predict residue three after one Step and the same period eight, then Step. Both additive orbits visit every residue but in different orders. They perform no reduction or gathering. Explain why a real AllGather must place values according to its rank-layout contract rather than simply append chunks in whichever physical arrival order happens first.