Gradient Synchronization and All-Reduce

"I compute my gradient in milliseconds, then spend half a second telling everyone else about it. We need to talk about my work-life balance."

An All-Reduce That Has Seen Some Gradients

In Section 15.3 we established the data-parallel recipe: replicate the model across $K$ workers, give each worker a different slice of the mini-batch, and let each one compute a gradient locally. The replicas start identical and must stay identical, step after step, so that they are genuinely training one model rather than $K$ slowly diverging ones. The mechanism that keeps them locked together is the subject of this section. After every backward pass, before any optimizer update, the $K$ local gradients are combined into one shared average gradient that every worker then applies. Because every worker applies the same update to the same starting weights, the replicas remain bit-for-bit identical, and the cluster behaves as a single model trained on the full batch. The combine is a collective operation, and which collective it is, and what it costs, is the whole story of data-parallel performance.

1. The Gradient Is an Average, So the Combine Is an All-Reduce Beginner

Recall the exact identity that Section 1.1 proved with a single calculation: the gradient of an average loss is the average of the per-example gradients, and an average decomposes over any partition of the data. Split a mini-batch of $B$ examples into $K$ shards, let worker $k$ compute the sum of gradients over its shard, and the full-batch gradient is

Every worker needs the same combined gradient $g$, not just its own piece, because every worker must apply the identical update to keep the replicas in sync. That requirement, every participant ends up holding the sum of all participants' contributions, is the exact definition of the all-reduce collective introduced in Section 4.3. A reduce alone would deposit the sum on one worker; a broadcast alone would copy one worker's value to all; all-reduce is the fused operation that does both at once and leaves every worker holding the reduced result. Data-parallel training does not merely use all-reduce occasionally; it is built around one all-reduce per step, and the parameter $K$ in the cost models below is the number of training replicas.

The figure below shows the per-step picture concretely. Four replicas each finish a backward pass holding a different local gradient vector; the all-reduce sums them component by component and hands the identical averaged gradient back to all four, which then take the same optimizer step.

2. Why This Is the Scaling Bottleneck Intermediate

The communication-bounds-training thesis of Section 4.1 stated the general principle: in a parallel system, the work that grows with the machine count or the model size, and that cannot be made faster by faster local compute, is the work that eventually limits you. Gradient synchronization is the textbook instance. Local gradient computation gets cheaper per worker as you add workers, because a fixed global batch is split more ways and each worker handles fewer examples. The all-reduce does not get cheaper. Its volume is the model size, which is constant, and the best ring algorithms make its time nearly independent of the worker count for large $K$. So as you scale out, the numerator of "communication time" holds steady while the denominator "compute time per worker" shrinks, and the fraction of each step spent communicating climbs toward one. That is the precise sense in which all-reduce, not the GPUs, sets the scaling ceiling.

The cost of a ring all-reduce was derived in Section 4.4; we use its bandwidth term here. A ring all-reduce of a payload of $S$ bytes across $K$ workers runs a reduce-scatter followed by an all-gather, and each worker sends $2(K-1)/K$ copies of its $S/K$-sized chunks over its slowest link. With per-link bandwidth $B$ and a model of $P$ parameters at $b$ bytes each (so $S = Pb$), the time is

Two features of this expression do the explaining. First, it is linear in the parameter count $P$: double the model and you double the synchronization time, every step, forever. Second, the $K$-dependence saturates fast; going from 32 to 512 workers barely changes the factor $2(K-1)/K$, which is already $1.94$ at $K=32$. So the all-reduce time is set almost entirely by the model, and adding workers buys more compute throughput without making the synchronization any cheaper. For the largest models this is brutal: a 70-billion-parameter gradient in 32-bit floats is 280 gigabytes of payload, and even at 100 gigabytes per second of link bandwidth the all-reduce term alone runs into seconds per step. The demonstration below makes these numbers concrete and shows the fraction of the step the all-reduce consumes.

# Ring all-reduce cost model for the per-step gradient synchronization.
# A gradient has P parameters, 4 bytes each (fp32). Ring all-reduce sends
# 2*(K-1)/K * (P*bytes) over the slowest link per worker; for large K the
# bandwidth term approaches 2*(P*bytes)/B and stops shrinking with K.
BYTES = 4
B = 100e9            # interconnect bandwidth per link: 100 GB/s (NVLink-class)
GPU_FLOPS = 300e12   # sustained 300 TFLOP/s per accelerator (mixed precision)
GLOBAL_TOKENS = 2_000_000   # fixed global batch (tokens) split across K workers

def allreduce_seconds(P, K):
    # Reduce-scatter then all-gather: 2*(K-1) chunk transfers of size P/K.
    return 2.0 * (K - 1) / K * (P * BYTES) / B

def compute_seconds(P, K):
    # Forward+backward ~ 6 FLOPs/param/token; each worker sees GLOBAL/K tokens.
    return 6.0 * P * (GLOBAL_TOKENS / K) / GPU_FLOPS

models = [
    ("ResNet-50",   25_000_000),
    ("BERT-large",  340_000_000),
    ("GPT-2 1.5B",  1_500_000_000),
    ("LLaMA 7B",    7_000_000_000),
    ("70B model",   70_000_000_000),
]
K = 64

print(f"Per step at K={K} workers, fixed global batch={GLOBAL_TOKENS:,} tokens, "
      f"B={B/1e9:.0f} GB/s/link\n")
print(f"{'model':<13}{'params':>15}{'compute (ms)':>14}"
      f"{'allreduce (ms)':>16}{'comm fraction':>15}")
for name, P in models:
    comp = compute_seconds(P, K)
    comm = allreduce_seconds(P, K)
    frac = comm / (comp + comm)
    print(f"{name:<13}{P:>15,}{comp*1e3:>14.1f}{comm*1e3:>16.1f}{frac:>14.0%}")

print("\nLLaMA 7B (P=7.0e9): comm fraction grows as workers are added")
print(f"{'workers K':>11}{'compute (ms)':>14}{'allreduce (ms)':>16}{'comm fraction':>15}")
P = 7_000_000_000
for k in (2, 8, 32, 128, 512):
    comp = compute_seconds(P, k)
    comm = allreduce_seconds(P, k)
    frac = comm / (comp + comm)
    print(f"{k:>11}{comp*1e3:>14.1f}{comm*1e3:>16.1f}{frac:>14.0%}")

Per step at K=64 workers, fixed global batch=2,000,000 tokens, B=100 GB/s/link

model                 params  compute (ms)  allreduce (ms)  comm fraction
ResNet-50         25,000,000          15.6             2.0           11%
BERT-large       340,000,000         212.5            26.8           11%
GPT-2 1.5B     1,500,000,000         937.5           118.1           11%
LLaMA 7B       7,000,000,000        4375.0           551.2           11%
70B model     70,000,000,000       43750.0          5512.5           11%

LLaMA 7B (P=7.0e9): comm fraction grows as workers are added
  workers K  compute (ms)  allreduce (ms)  comm fraction
          2      140000.0           280.0            0%
          8       35000.0           490.0            1%
         32        8750.0           542.5            6%
        128        2187.5           555.6           20%
        512         546.9           558.9           51%

Two trends in Output 15.4.1 carry the whole argument. The all-reduce column of the first table scales linearly with the parameter count, exactly as $T_{\text{all-reduce}} \propto P$ predicts, so the absolute synchronization cost is dictated by the model you train. The second table holds the model fixed and shows the more dangerous trend: as workers are added under a fixed global batch, per-worker compute falls but the all-reduce does not, and the communication fraction marches from under one percent to over fifty. A practitioner reads these two tables together as a single warning: a bigger model makes every all-reduce longer, and more workers make that all-reduce a larger share of the step, so the largest models on the largest clusters are exactly where synchronization dominates.

3. NCCL: The All-Reduce That Actually Runs on GPUs Intermediate

The cost model above assumes a competent ring algorithm running on a topology that supports it. On NVIDIA GPUs, the library that supplies both is NCCL, the NVIDIA Collective Communications Library introduced as the GPU-native collectives layer in Section 4.8. NCCL is what `torch.distributed` calls when you select the `nccl` backend, and it is the reason a single line of framework code turns into an efficient multi-GPU all-reduce. It does three things that a naive implementation would not. It discovers the interconnect topology (NVLink and NVSwitch inside a node, InfiniBand or RoCE between nodes) and builds rings and trees that match it, so the slowest link in the formula above is as fast as the hardware allows. It moves gradient buffers GPU-to-GPU directly, bypassing the host CPU and host memory through GPUDirect, which keeps the bandwidth term close to the hardware peak. And it overlaps the reduce-scatter and all-gather phases with the kernels still running on the device, which is the hook that Section 15.5 exploits to hide synchronization behind the backward pass.

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

dist.init_process_group("nccl")              # NCCL discovers NVLink/IB topology
local_grad = compute_shard_gradient()        # this replica's gradient tensor (on GPU)

# One collective: NCCL runs the ring reduce-scatter + all-gather of Section 4.4.
dist.all_reduce(local_grad, op=dist.ReduceOp.SUM)
local_grad /= dist.get_world_size()          # SUM -> mean keeps replicas identical

4. What This Sets Up Intermediate

The cost model leaves two honest exits, and the rest of the chapter walks through both. The first is to move fewer bytes: gradient compression, low-precision all-reduce, and the local-update schemes that communicate less often all attack the $Pb$ in the numerator directly, and we meet them as the communication-avoiding frontier below and again with full machinery in Chapter 10. The second exit, and the one the very next section takes, is to stop paying the all-reduce as a separate phase at all. If the gradient for the last layer is ready the instant its backward pass finishes, its all-reduce can run while the earlier layers are still computing their gradients, so the synchronization disappears behind compute that was happening anyway. Making that overlap practical requires grouping gradients into buckets and launching their collectives as soon as each bucket fills, which is precisely the gradient bucketing and communication-computation overlap of Section 15.5. Everything in this section is the cost that those mechanisms exist to hide.

We now have the defining cost of data parallelism stated exactly: one all-reduce per step, sized by the model, linear in the parameter count, and an ever-larger share of the step as the cluster grows. That single quantity explains why data parallelism alone cannot train the largest models, why every serious training stack obsesses over communication, and why the next section is about hiding rather than merely accepting this cost. The mechanics of that hiding begin in Section 15.5.