All-Reduce: Synchronizing Gradients in Data-Parallel SGD

"Everyone brought a number. I made sure everyone left with the same number. That is the entire job, and somehow it is the hardest one in the building."

An All-Reduce That Has Seen Some Gradients

Chapter 4 opened by arguing that communication, not raw compute, is what bounds distributed training, and Section 4.2 separated point-to-point messages from collective operations that the whole group participates in at once. All-reduce is the most important collective in this book, and the reason is specific rather than general: it is the exact shape of the synchronization step in data-parallel stochastic gradient descent. We could introduce all-reduce as an abstract reduction over a process group, the way a parallel-computing textbook would, but that hides why a deep-learning engineer cares. So we introduce it through its defining use. Every worker has just finished a backward pass and is holding a gradient. To take a correct optimizer step, they all need the average of those gradients. The primitive that delivers exactly that, to every worker at once, is all-reduce.

1. The Operation, Defined Precisely Beginner

Fix a group of $K$ workers, indexed $k = 0, 1, \ldots, K-1$, and suppose each worker $k$ holds a vector $g_k \in \mathbb{R}^{P}$ of the same length $P$. The vectors differ across workers; in training, $g_k$ is the gradient computed on worker $k$'s shard of the batch. An all-reduce with reduction operator $\oplus$ produces a single vector $r$ and leaves a copy of $r$ on every worker:

The operator $\oplus$ is applied elementwise and must be associative and commutative so the result does not depend on the order in which the workers' contributions are combined; sum, mean, max, min, and product all qualify. The two facts that make this a collective rather than a plain reduction are worth stating separately. First, the reduction is over all workers: no single worker can compute $r$ from its own data alone. Second, the result is shared with all workers: when the call returns, every worker holds the same $r$, byte for byte. A reduction that left the answer on only one worker would be a reduce; broadcasting that answer back to everyone is what promotes it to an all-reduce, and the name simply concatenates the two halves.

For gradient synchronization the operator we want is sum, and the mean is one scalar division away. If worker $k$ contributes the sum of gradients over its own $n_k$ examples, then an all-reduce with $\oplus = +$ leaves every worker holding the total gradient sum over the whole batch; dividing by the batch size $N = \sum_k n_k$ gives the averaged gradient that the optimizer needs. This is the same combine-then-divide pattern you performed by hand in Section 1.1, now recognized as a single named primitive.

2. Averaging Gradients Is an All-Reduce Intermediate

The link back to Section 1.1 is not an analogy; it is an identity. Recall the loss is an average over $N$ examples, so its gradient is an average of per-example gradients, and an average decomposes across any partition of the examples. Split the batch into $K$ shards, give shard $k$ to worker $k$, and let worker $k$ compute the local sum $s_k = \sum_{i \in \text{shard } k} \nabla \ell(w; x_i, y_i)$. The averaged gradient the optimizer wants is

The parenthesized sum is precisely an all-reduce with $\oplus = +$ over the vectors $s_k$. After the all-reduce every worker holds the same total $\sum_k s_k$, divides by $N$, and now holds the identical averaged gradient. Because the workers started with identical model parameters and now apply the identical gradient with the identical optimizer, they remain in lockstep: the model on every worker stays bit-for-bit the same at every step, which is what makes data-parallel training a single coherent training run rather than $K$ diverging ones. The gradient identity of Section 1.1 is the correctness proof; all-reduce is the mechanism that realizes it across machines without ever gathering all the data in one place.

3. Building All-Reduce From Scratch Intermediate

To make the semantics concrete and impossible to misread, we simulate $K$ ranks in a single process, each holding its own private vector, and implement all-reduce so that the only way information crosses between ranks is through the explicit combine step. We then check the contract directly: after the call, do all $K$ ranks hold one and the same reduced vector, and does that vector equal the independently computed ground-truth sum? The demonstration below uses the sum reduction and then a mean, mirroring gradient synchronization.

import numpy as np

def all_reduce_sum(rank_buffers):
    """Simulate all-reduce(SUM) over K ranks, each holding its own vector.
    Input: a list of K arrays (one private buffer per rank).
    Output: a list of K arrays, every entry the elementwise sum, so that
    the SAME result is now resident on every rank (the 'all' in all-reduce)."""
    total = np.zeros_like(rank_buffers[0])
    for buf in rank_buffers:          # the only cross-rank data movement
        total = total + buf           # associative, order-independent
    return [total.copy() for _ in rank_buffers]   # broadcast a copy to each rank

rng = np.random.default_rng(0)
K, P = 6, 4                                       # 6 simulated ranks, length-4 vectors
rank_buffers = [rng.standard_normal(P) for _ in range(K)]   # each rank's private vector

reduced = all_reduce_sum(rank_buffers)            # the collective
ground_truth = np.sum(rank_buffers, axis=0)       # what every rank SHOULD now hold

# Contract 1: all ranks hold the identical vector after the call.
all_identical = all(np.array_equal(reduced[0], reduced[k]) for k in range(K))
# Contract 2: that shared vector equals the true elementwise sum.
matches_truth = np.allclose(reduced[0], ground_truth)

print("ranks K                 :", K)
print("all ranks identical     :", all_identical)
print("equals ground-truth sum :", matches_truth)
print("shared reduced vector   :", np.array2string(reduced[0], precision=4))

# Gradient synchronization is the same call followed by one division.
mean_grad = reduced[0] / K
print("averaged (mean) vector  :", np.array2string(mean_grad, precision=4))

Running this with the project Python interpreter produces the output below.

ranks K                 : 6
all ranks identical     : True
equals ground-truth sum : True
shared reduced vector   : [-4.1115 -0.2046 -0.1783  1.7551]
averaged (mean) vector  : [-0.6852 -0.0341 -0.0297  0.2925]

Two booleans both reading True is the entire semantics of all-reduce on display: consensus (every rank holds the same vector) and correctness (it is the right reduction). Our implementation gathered the buffers into one running total, which is fine for a teaching simulation but is exactly the asymmetric, leader-centric pattern that real systems avoid. A production all-reduce never funnels every contribution through one place; it arranges the communication so each worker moves only a small share of the data and the bandwidth load is balanced. How that is done, and why the ring all-reduce in particular reshaped deep-learning systems, is the whole of Section 4.4.

4. The Library Shortcut and Automatic Synchronization Intermediate

In a real multi-machine job you never gather buffers into one Python list, because the buffers live on different machines. PyTorch exposes all-reduce as a single collective over a process group, and each worker calls it on its own local tensor; when the call returns, that tensor has been overwritten in place with the reduced result, identical on every worker. The manual loop of Code 4.3.1 collapses to one line.

# Launch with: torchrun --nproc_per_node=6 thisfile.py
import torch, torch.distributed as dist

dist.init_process_group("nccl")                 # join the group of K workers
local_grad = compute_shard_gradient()           # this worker's gradient tensor (length P)

dist.all_reduce(local_grad, op=dist.ReduceOp.SUM)   # every worker now holds the SUM
local_grad /= dist.get_world_size()                 # divide by K to get the mean gradient
# local_grad is now the averaged gradient, identical on every worker

So the same primitive appears at three altitudes: the explicit combine you can read in Code 4.3.1, the one-line collective in Code 4.3.2, and the invisible hook that fires automatically inside DDP's backward pass. They compute the identical averaged gradient; they differ only in how much machinery the framework hides. What none of them changes is the cost, which is set by how the bytes actually move across the interconnect, and that is what we turn to next.

5. Why the Algorithm Still Matters Advanced

The naive picture of all-reduce, every worker ships its full gradient to one place that sums and rebroadcasts, has a fatal scaling property: the central node must send and receive roughly $K$ times the gradient, so synchronization time grows with the number of workers even though the useful answer is the same size regardless of $K$. For a billion-parameter model and hundreds of workers, that central bottleneck would dominate the step and erase the speedup that distribution was supposed to buy, exactly the communication tax that Chapter 4 opened with and that the alpha-beta cost model of Chapter 3 lets you quantify. The fix is to choose an all-reduce algorithm whose cost does not grow with $K$, by spreading the data movement evenly so no single worker is the funnel. That is why a primitive with a one-line interface still gets a whole section of algorithmic attention.

We have defined all-reduce, tied it by identity to the gradient average of Section 1.1, built it from scratch with its contract verified, and reduced it to one library call that DistributedDataParallel fires automatically. What remains is the part the practical example just foreshadowed: the algorithms that make all-reduce fast enough to be worth using, and why the ring all-reduce in particular changed how deep learning is trained. That is Section 4.4.