Distributed Communication Collectives
The setting
A communicator (a process group) is an ordered set of $P$ ranks, each holding a tensor. A collective is an operation that every rank in the group calls, with matching arguments, that transforms the group's data as a whole. Two facts shape everything:
Collectives are synchronizing
Every rank must call the same collectives in the same order. A rank that skips one, or calls them in a different order, deadlocks the group or silently pairs up mismatched buffers.
Ranks can belong to many groups
The same device can be in a data-parallel group along one mesh axis and a tensor-parallel group along another. Groups are how a device mesh is expressed in code.
The six primitives
Let each rank start with a tensor of $N$ bytes, and write $x_p$ for rank $p$'s data. Let $\oplus$ be an elementwise reduction, usually sum.
| Collective | Before | After | Bytes each rank moves |
|---|---|---|---|
| Broadcast | root has $x$ | everyone has $x$ | $\approx N$ |
| Reduce | rank $p$ has $x_p$ | root has $\bigoplus_p x_p$ | $\approx N$ |
| All-reduce | rank $p$ has $x_p$ | everyone has $\bigoplus_p x_p$ | $2N\frac{P-1}{P}$ |
| All-gather | rank $p$ has shard $x_p$ of size $N/P$ | everyone has $[x_0;\dots;x_{P-1}]$ | $N\frac{P-1}{P}$ |
| Reduce-scatter | rank $p$ has full-size $x_p$ | rank $p$ has shard $p$ of $\bigoplus_q x_q$ | $N\frac{P-1}{P}$ |
| All-to-all | rank $p$ has $P$ shards, one per destination | rank $p$ has shard $p$ from every rank | $N\frac{P-1}{P}$ |
All-to-all is the odd one out: it is a distributed transpose, moves no reduction, and is the primitive that expert parallelism is built on.
The cost model
Communication time is modelled with two constants: a per-message latency $\alpha$ and an inverse bandwidth $\beta$ measured in seconds per byte. A single message of $n$ bytes costs
An algorithm that takes $s$ sequential communication steps and moves $n$ bytes per rank therefore costs $s\alpha+\beta n$. Two regimes follow immediately, and they call for different algorithms:
Small messages are latency-bound
The $s\alpha$ term dominates. You want an algorithm with few steps, so a tree with $O(\log P)$ depth wins even though it moves more data.
Large messages are bandwidth-bound
The $\beta n$ term dominates. You want the algorithm that moves the least data per rank, which is the ring, even though it takes $O(P)$ steps.
Gradient buffers in training are megabytes to gigabytes, so training is firmly in the bandwidth-bound regime. That is why the ring algorithm is the one worth understanding in detail.
Ring all-reduce, derived
Arrange the ranks in a logical ring where rank $p$ only ever sends to rank $p+1 \bmod P$. Split each rank's tensor into $P$ chunks of $N/P$ bytes. The algorithm runs in two phases of $P-1$ steps each.
The only difference between the phases is add versus overwrite. Counting the traffic:
When the ring is the wrong choice
The ring is bandwidth-optimal but its latency grows as $2(P-1)\alpha$. For a small tensor across many nodes that term dominates, and a tree does much better:
| Algorithm | Latency term | Bandwidth term | Best for |
|---|---|---|---|
| Ring | $2(P-1)\alpha$ | $2\frac{P-1}{P}N\beta$ | large tensors, the training case |
| Recursive halving/doubling | $2\log_2 P\cdot\alpha$ | $2\frac{P-1}{P}N\beta$ | power-of-two $P$, medium tensors |
| Binary tree (reduce + broadcast) | $2\log_2 P\cdot\alpha$ | $2N\beta$ at the root links | small tensors, latency-bound |
Production libraries choose per call, based on message size and topology, and typically run hierarchically: reduce within a node over the fast intra-node fabric, then all-reduce across nodes over the slower network with only one representative per node, then broadcast back down inside each node. Since intra-node links are often an order of magnitude faster than inter-node ones, this turns a $P$-way problem into a much smaller one.
Making communication free
The best collective is one you never wait for. Gradients become available progressively during the backward pass, starting with the last layer, so there is a long window in which to send them while the remaining layers are still computing.
Bucket. Issuing one collective per parameter tensor pays $\alpha$ hundreds of times. Group gradients into buckets of a few tens of megabytes and issue one collective per bucket. This is a pure latency optimization.
Launch asynchronously. Fire the collective as soon as a bucket is full and keep computing. The work is genuinely overlapped only if it runs on a separate stream from the compute kernels.
Wait as late as possible. Block on the handles only just before the optimizer step needs the values.
When overlap works, communication time vanishes from the critical path entirely and MFU stays high. When it does not — because buckets are too small, or the collective shares a stream with compute — you see it immediately as a gap in a profiler trace and as a disappointing utilization number.
All-reduce is reduce-scatter followed by all-gather, costs $2N\frac{P-1}{P}$ bytes per rank, and is therefore almost independent of cluster size. Everything else is choosing the right algorithm for the message size, putting the chattiest collective on the fastest link, and hiding it behind compute.
Check yourself
- Show that all-reduce moves exactly twice the bytes of all-gather, using the decomposition identity. derivation
- For $N=1$ GB, $P=64$, $\beta=1/(25\ \text{GB/s})$ and $\alpha=5\ \mu$s, compute the ring all-reduce time and say which term dominates. calculation
- Explain why the bandwidth term of a ring all-reduce approaches a constant as $P$ grows. reasoning
- Give the state of each chunk after phase 1 completes, and prove the diagonal claim. derivation
- Implement bucketed, overlapped gradient all-reduce and measure the utilization difference against a naive per-tensor version. implementation