Collective Communication.
Four ways for a group of devices to exchange data — and an all-reduce is just a reduce-scatter followed by an all-gather.
Every parallel scheme in this volume communicates through a small vocabulary of collectives, implemented in libraries such as NVIDIA’s NCCL. In an all-reduce every device ends with the sum of everyone’s data — how data parallelism sums gradients. In an all-gather each device contributes a shard and every device ends with all of them — how sharded weights are assembled for use. A reduce-scatter sums and leaves each device one shard of the sum. An all-to-all sends a different piece to every other device — how a mixture of experts routes tokens.
With a ring or similar bandwidth-optimal algorithm over n devices with link bandwidth B, an all-gather or reduce-scatter of S bytes takes about S(n − 1)/(nB), and an all-reduce, being one of each, twice that. The cost per device barely grows with n, which is what makes thousand-device data parallelism possible. The workshop compares the four collectives for a chosen message size.
Below: left, time of each collective for the chosen message size over NVLink-speed links, against the number of devices. Right, bytes sent per device by each.
- Define all-reduce, all-gather, reduce-scatter and all-to-all.
- Compute their costs with ring algorithms.
- Show that all-reduce = reduce-scatter + all-gather.
- Match each collective to the parallelism that uses it.
Four collectives, one cost model
Across 8 devices, an all-reduce of 1 GB takes 3.96 ms, the same as a reduce-scatter followed by an all-gather (3.96 ms) — because that is how ring all-reduce works: each device first ends with one summed shard, then the shards are passed round until everyone has them all. From 2 devices to 128, the all-reduce time grows only 2.54 times, since each device sends about twice the message whatever the group size. What does grow is latency: a tiny 1 KB all-reduce over 128 devices still takes 1.27 ms in 254 sequential steps, which is why small messages are bucketed together.
Each parallel scheme is a choice of collective. Data parallelism all-reduces gradients once per step (plate XIV.4); sharded data parallelism replaces it with a reduce-scatter of gradients and all-gathers of weights (plate XIV.5); tensor parallelism all-reduces activations inside every layer (plate XIV.6); expert parallelism all-to-alls tokens to experts and back (plate XIV.9); pipeline parallelism sends activations point to point between stages (plate XIV.7). Knowing the bytes each moves, and over which links, is most of the work of choosing a layout.
Terms
- All-reduce
- Every device ends with the sum of all devices’ data.
- All-gather
- Every device ends with every device’s shard.
- Reduce-scatter
- The sum is computed and each device keeps one shard of it.
- All-to-all
- Each device sends a different piece to every other device.
Sources
- Patarasuk and Yuan, “Bandwidth optimal all-reduce algorithms for clusters of workstations”, JPDC 69(2), 2009.
- NVIDIA Collective Communications Library (NCCL) documentation.