Vol. XIV · Parallelism Plate 3 of 11collective communication, all-reduce, all-gather, reduce-scatter, all-to-all, ring algorithm, NCCL, bandwidthCE

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.

What you will learn
Workshop
S bytes per device in total (for all-gather, S is the size of the assembled result), link bandwidth 450 GB/s each way, 5 µs per step; ring algorithms for all-reduce, all-gather and reduce-scatter, pairwise exchange for all-to-all. Left: time against the number of devices. Right: bytes each device sends.

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.

Real libraries pick among ring, tree and other algorithms by message size and topology; trees cut latency for small messages, rings maximise bandwidth for large ones. The formulas here are the bandwidth-bound limit.

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.
Parallel training talks through four collectives; a ring all-reduce is a reduce-scatter plus an all-gather, and its cost per device barely grows with the group.

Sources

  1. Patarasuk and Yuan, “Bandwidth optimal all-reduce algorithms for clusters of workstations”, JPDC 69(2), 2009.
  2. NVIDIA Collective Communications Library (NCCL) documentation.