Suppose your model fits on one GPU, but one GPU would take a month to train it. You have eight GPUs. What is the simplest way to make training eight times faster?

Give each GPU a full copy of the model and a different eighth of every batch. Each GPU runs the forward and backward pass on its own slice, the GPUs average their gradients, and every GPU applies the same averaged update to its own copy. The copies start identical and receive identical updates, so they stay identical forever. This is data parallelism, and it is how almost every large model is trained, usually combined with the other strategies of this chapter.

The whole cost of data parallelism is in one word of that description, average. Eight GPUs each hold a gradient as large as the model, and they must combine them into one, every step. How to do that fast, with a pattern of messages called all-reduce, is the subject of this section. Along the way we meet the handful of collective operations that every distributed training system is built from, and derive what each one costs.

Why averaging gradients is exactly large-batch training

Data parallelism is not an approximation. Take a global batch of $B$ examples with index set $\mathcal{I}$ and split it into $p$ shards $\mathcal{I}_1, \dots, \mathcal{I}_p$ of $B/p$ examples each, one per GPU. Write $\ell_i$ for the loss on example $i$. The loss is the mean over the batch, so its gradient is

In plain words, the gradient of the whole batch is the average of the GPUs' gradients, as long as every GPU averages over the same number of examples.

That condition is not a technicality. In language modeling the natural unit is the token, not the sequence, and different GPUs can see different numbers of non-padding tokens. If each GPU averages its loss over its own tokens and the GPUs then average those means, a GPU with 100 tokens gets the same say as one with 10,000. The fix is to sum the loss over tokens on each GPU, sum the token counts across GPUs, and divide once. The cell below checks both claims on a linear model.

Data parallelism changes the global batch size, not the algorithm. Eight GPUs with 32 sequences each train exactly like one GPU with 256 sequences. That is both its strength (no change to the math) and its limit. Beyond the critical batch size, larger batches stop reducing the number of steps needed (§8.7), so you cannot add GPUs forever just by growing the batch. We return to this limit at the end of the section.

The collective operations

When many GPUs work on one job, they talk in a small vocabulary of group operations called collectives, named in the MPI standard decades before deep learning and implemented for GPUs by libraries such as NVIDIA's NCCL. Every rank (one process, usually one GPU) calls the same collective at the same time, each with its own buffer. Picture $p$ ranks, each holding a buffer split into $p$ chunks.

The reductions usually sum, though maximum and others exist. Two identities are worth seeing right away. An all-reduce is a reduce-scatter followed by an all-gather (reduce each chunk on its owner, then hand the owners' chunks to everyone). And an all-to-all is a transpose of a $p\times p$ grid of chunks, every rank sending a different piece to every other rank. Here are all of them as plain functions on lists of arrays, one array per simulated rank.

Ring all-reduce, derived

Now the central question. Each of $p$ ranks holds a gradient of $M$ bytes. How long does it take to give every rank the sum?

The obvious algorithm sends everything to one rank, adds, and sends the result back. Rank 0 then receives $(p-1)M$ bytes and sends $(p-1)M$ bytes, so with links of bandwidth $\mathcal{B}$ (bytes per second) the time is at least $2(p-1)M/\mathcal{B}$. It grows linearly with the number of GPUs, and every other rank's link sits idle while rank 0 works. With 64 GPUs and a 2 GB gradient over 50 GB/s links, that is 5 seconds per step, far longer than the computation.

The ring all-reduce removes the bottleneck by making every rank do an equal share of the sending, adding, and receiving. Arrange the ranks in a ring, each sending only to its right neighbor $r+1$ (all indices modulo $p$), and split every buffer into $p$ chunks of $M/p$ bytes. The algorithm has two phases of $p-1$ steps each.

1. Reduce-scatter. At step $k = 0, 1, \dots, p-2$, rank $r$ sends its current copy of chunk $(r - k) \bmod p$ to rank $r+1$, which adds it into its own copy of that chunk. A chunk travels around the ring collecting one more contribution at every hop. After $p-1$ steps, it has visited every rank. Working through the indices, the chunk that arrives at rank $q$ in the last step is $(q+1) \bmod p$, so rank $q$ now holds the complete sum of chunk $(q+1) \bmod p$.
2. All-gather. At step $k = 0, \dots, p-2$, rank $r$ sends chunk $(r + 1 - k) \bmod p$ to rank $r+1$, which overwrites its copy. At step 0 each rank passes on the finished chunk it owns, and at every later step it passes on the finished chunk it just received. After $p-1$ steps every rank has every finished chunk.

The widget below runs it. Each square is one chunk on one rank, and the dots show which ranks' contributions it contains.

Now count. In each of the $2(p-1)$ steps, every rank sends one chunk of $M/p$ bytes and receives one. All $p$ links work at once, so a step takes $\alpha + (M/p)/\mathcal{B}$, where $\alpha$ is the fixed latency of sending any message at all (launching the transfer, synchronizing, crossing the switch). Writing $\beta = 1/\mathcal{B}$ for the time per byte (this latency-plus-bandwidth cost model, $\alpha + M\beta$ per message, is Hockney's model, the standard way to reason about collectives), the total is

In plain words, each rank sends and receives just under twice its buffer, however many ranks there are, but it pays the latency $2(p-1)$ times.

The bandwidth term is the remarkable part. As $p$ grows, $2(p-1)/p$ approaches 2 from below, so the data each GPU moves is essentially constant. Adding GPUs to a ring does not make the all-reduce slower, except through latency. And no algorithm can do better on bandwidth. Patarasuk and Yuan (2009) proved that some rank must send at least $2(p-1)M/p$ bytes in any all-reduce (assuming the data are not compressed). The argument is a count of messages. Summing one number held by $p$ ranks and delivering the sum to all of them takes at least $2(p-1)$ messages, each the size of one number. The elements of a buffer are independent, so an $M$-byte all-reduce needs at least $2(p-1)M$ bytes of messages in total, and spread over $p$ senders, at least one must send $2(p-1)M/p$. The ring meets that bound exactly. This is a theorem, and it is why rings are the default for large messages.

The two halves are themselves the collectives of the table. A reduce-scatter alone is phase 1, $(p-1)$ steps moving $(p-1)M/p$ bytes per rank. An all-gather alone is phase 2, with the same cost, where $M$ is the size of the full gathered result. So each of them costs half an all-reduce, which will matter when ZeRO splits an all-reduce into its halves (§10.3).

def collective_bytes(op, M, p):
    return 2 * M          # placeholder

import numpy as np

def ring_allreduce(bufs):
    total = sum(np.asarray(b, float) for b in bufs)
    return [total.copy() for _ in bufs], []          # right answer, but no ring and no log

Latency, trees, and choosing an algorithm

The ring's weakness is its $2(p-1)$ latency terms. For a large gradient bucket the bandwidth term dominates and the ring is ideal, but for a small message, say a few kilobytes of loss values or a tensor-parallel activation of one token, the latency term dominates, and it grows linearly with the number of ranks. Algorithms built on trees pay latency only logarithmically.

Recursive doubling is the simplest. In step $j = 0, 1, \dots, \log_2 p - 1$, rank $r$ exchanges its whole current buffer with rank $r \oplus 2^j$ (the rank whose index differs in bit $j$) and both add. After $\log_2 p$ steps everyone has the sum, so

It has the fewest steps possible, but it sends the whole buffer every step, $M\log_2 p$ bytes instead of the ring's $2M$. Recursive halving and doubling (Rabenseifner's algorithm) fixes that by exchanging half of the remaining data at each step, which gives $2\log_2 p$ latency terms and the optimal $2(p-1)M/p$ bandwidth when $p$ is a power of two. NCCL uses rings and double binary trees, two complementary trees that each carry half of the data so that every rank does useful work in one of them, which gives logarithmic latency at close to full bandwidth for large clusters.

Setting $t_{\text{ring}} = t_{\text{rd}}$ gives the message size where the ring starts to win,

With $\alpha = 10\ \mu$s and 50 GB/s links, $\alpha/\beta = 500$ KB. For eight ranks $M^\star \approx 4.4$ MB, and for 512 ranks about 72 MB. Small messages want trees, large ones want rings, and the crossover moves up as the cluster grows. Libraries pick the algorithm per call from the message size and the topology. Play with the trade-off below.

import numpy as np

def allreduce_time(M, p, alpha, beta, algo):
    return 2 * (p - 1) * alpha + 2 * M * beta        # placeholder

def ring_crossover(p, alpha, beta):
    raise NotImplementedError

The hardware the messages run on

The formulas need two numbers, latency and bandwidth, and the hardware gives very different ones depending on whether the two GPUs share a server.

The gap is the most important fact in this chapter. Inside a DGX H100, each GPU can send 450 GB/s to the others through NVSwitch. Between nodes, each GPU gets its own 400 Gb/s InfiniBand NIC, 50 GB/s, nine times less, and latency across the network is higher too. Almost every layout decision in §10.7 comes down to keeping the chattiest communication on NVLink. Measured numbers are lower than these nominal rates, so for real planning run NVIDIA's nccl-tests on your cluster.

A note on reading those benchmarks. nccl-tests reports an algorithm bandwidth, $M$ divided by the time, and a bus bandwidth, the algorithm bandwidth times $2(p-1)/p$ for an all-reduce. The bus bandwidth is the rate at which each link actually carried data in the ring analysis above, so it can be compared directly with the link's rated speed, whatever $p$ is.

A hierarchical all-reduce exploits the gap. Across $n$ nodes of $g$ GPUs, first reduce-scatter within each node over NVLink, so each GPU holds the node's sum of $1/g$ of the buffer. Then all-reduce those $1/g$ pieces across nodes, each GPU talking to its counterparts in other nodes through its own NIC. Then all-gather within each node. Each GPU's InfiniBand traffic is now $2(n-1)/n \cdot M/g$ instead of $2(ng-1)/(ng) \cdot M$.

Real libraries build rings and trees that use every NIC in parallel, in the spirit of this calculation, so you rarely write it yourself. What matters is the conclusion. Traffic that crosses nodes should be as small as possible, and every GPU's NIC should carry its share.

Overlapping communication with the backward pass

Now put the all-reduce into a training step. The naive version runs the full backward pass, then all-reduces every gradient, then steps the optimizer. The GPUs compute while the network idles, then the network works while the GPUs idle.

The fix is to notice that gradients do not appear all at once. The backward pass produces them from the last layer to the first, so the gradients of the last layer are final while the first layers are still being differentiated. Their all-reduce can start right away, on a separate CUDA stream, while computation continues. Only the communication for the first layers, which finish last, is left exposed at the end.

Launching one all-reduce per parameter tensor would pay the latency $\alpha$ hundreds of times per step, so gradients are grouped into buckets. Li et al. (2020) describe how PyTorch's DistributedDataParallel (DDP) does it. It assigns parameters to buckets of about 25 MB (bucket_cap_mb) in roughly the reverse of their order in the model, which approximates the order in which their gradients become ready. An autograd hook on each parameter marks its gradient ready, and when every gradient in a bucket is ready, DDP launches an asynchronous all-reduce of that bucket. Before the optimizer step it waits for all of them. The bucket size is a real trade-off. Tiny buckets pay latency many times, and huge buckets wait too long to start, so the last bucket's communication sticks out past the end of the backward pass.

DDP has a few more details worth knowing. At construction it broadcasts rank 0's parameters and buffers to every rank, so all replicas start identical. It divides the summed gradients by the number of ranks, giving the average. Inside with model.no_sync(): it skips the all-reduce, which is how gradient accumulation avoids communicating on every micro-batch. And it needs find_unused_parameters=True when some parameters receive no gradient in a step, because otherwise it would wait forever for a bucket that never fills.

def ddp_step_time(grads, bucket_cap, p, alpha, beta):
    total = sum(n for n, _ in grads)
    end_bwd = grads[-1][1]
    t = end_bwd + 2 * (p - 1) * alpha + 2 * (p - 1) / p * total * beta    # no overlap at all
    return t, t - end_bwd

When data parallelism stops scaling

Can the all-reduce always hide behind the backward pass? Count both sides for one step. Each GPU processes $n_{\text{tok}}$ tokens. The backward pass of a model with $N$ parameters costs about $4N$ FLOPs per token (two-thirds of the $6N$ training cost of §9.2), so at $F$ FLOP/s it takes $4Nn_{\text{tok}}/F$ seconds. The all-reduce moves about $2 \times 2N$ bytes of bfloat16 gradients per GPU, taking $4N/\mathcal{B}$ seconds. The communication fits inside the backward pass when

In plain words, each GPU must chew through at least $F/\mathcal{B}$ tokens per step, whatever the model size, to give the network time to deliver the gradients.

For an H100 at its peak of 989 TFLOP/s in bfloat16, over a 50 GB/s InfiniBand NIC, $F/\mathcal{B} \approx 19{,}800$ tokens per GPU per step. At a realistic 40% of peak it is about 7,900. Within a node, over NVLink, it is only about 2,200. These are generous estimates, since they assume perfect overlap and ignore latency.

Now combine this with the batch size limit. If each of $n_{\text{dp}}$ data-parallel GPUs needs at least 8,000 tokens per step, the global batch is at least $8{,}000\,n_{\text{dp}}$ tokens. Pretraining batches are typically a few million tokens, bounded by the critical batch size (§8.7). A 4M-token batch therefore allows only about 500 data-parallel GPUs before communication is exposed or each GPU gets too few tokens to be efficient. To use 16,000 GPUs, as Llama 3 405B did, you have to split the model itself, so that a group of GPUs works on the same tokens. That is the job of the next four sections.

In PyTorch

A DDP training script is launched with one process per GPU, for example torchrun --nproc_per_node=8 train.py, which sets the environment variables RANK, LOCAL_RANK, and WORLD_SIZE that the script reads. This code cannot run in a single-process notebook, so it is shown as reference.

The collectives are also available directly, and writing data parallelism by hand shows what DDP automates. A hook on each parameter fires once its gradient is fully accumulated and launches an asynchronous all-reduce, and the step waits for all of them.

This version launches one all-reduce per tensor and pays latency for each, which is exactly the problem DDP's buckets solve.

Read more

• Li et al. (2020). PyTorch Distributed: Experiences on Accelerating Data Parallel Training. How DDP buckets, overlaps, and synchronizes, with measurements.
• Patarasuk and Yuan (2009). Bandwidth optimal all-reduce algorithms for clusters of workstations. The lower bound on all-reduce traffic and the ring that meets it.
• Thakur, Rabenseifner, and Gropp (2005). Optimization of Collective Communication Operations in MPICH. Rings, recursive doubling, and recursive halving with their cost models.
• Sergeev and Del Balso (2018). Horovod: fast and easy distributed deep learning in TensorFlow. The paper that brought ring all-reduce from HPC into deep learning.
• Goyal et al. (2017). Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour. Data parallelism on 256 GPUs, and the learning-rate scaling that makes large batches work.
• NVIDIA. NCCL Tests, performance notes. Algorithm bandwidth versus bus bandwidth, explained.
• Douillard et al. (2023). DiLoCo: Distributed Low-Communication Training of Language Models. Data parallelism that synchronizes every few hundred steps, with an inner and an outer optimizer.
• Austin et al. (2025). How to Scale Your Model. The chapters on networking and on training derive when data parallelism is compute-bound, for TPUs and GPUs.

Watch next

▶

▶

▶