Parallelism I: data

Data parallelism. Every worker (a GPU, or a whole machine) holds a full copy of the model, the Decoder site's model here, and gets its own share of each batch. Each runs the forward and backward passes of chapters 1 and 4 on its share; then the workers' gradients are averaged and every copy takes the same optimiser step, so the copies stay identical. The average of $W$ workers' mean gradients over $B/W$ windows each is exactly the mean gradient over the whole batch of $B$: data parallelism changes where the arithmetic happens, not what it computes. Checked on the model: chapter 13's four simulated workers with four windows each, averaging their gradients, end the default run at a validation loss of 2.299304; the default run, sixteen windows on one worker, at 2.299304 (the additions happen in a different order, so the last bits differ). More workers means a bigger batch per step, and chapter 6's critical batch size says how far that helps.

All-reduce, and the ring. Averaging the gradients is a collective operation called an all-reduce: every worker starts with its own vector and ends with the sum (or mean) of all of them. The naive way, everyone sending its whole gradient to one worker, makes that worker's link the bottleneck. The ring all-reduce in the animation does better: cut each vector into $p$ chunks; in $(p-1)$ reduce-scatter steps each worker passes one chunk to its neighbour, which adds it to its own, until each worker holds the finished sum of one chunk; in $(p-1)$ all-gather steps the finished chunks go round the ring until everyone has all of them. Each worker sends $2(p-1)/p$ of its vector, no more than twice the vector however many workers there are, which is the least any all-reduce can send (Patarasuk & Yuan 2009). What does grow with the ring is the number of steps, $2(p-1)$, each paying a link's latency $\alpha$. For LLaMA 7B's gradient (that is 13.5 GB in BF16), a ring all-reduce over the 8 GPUs of one H100 server's NVLink takes 0.052 s by this site's systems model, over 64 GPUs on the network 0.53 s; a binary tree, which climbs $\log_2 p$ levels up and down instead, has fewer latency hops but, done naively (each level waiting for the whole vector), takes 3.2 s on the same 64 GPUs. Real libraries pipeline both.

All-gather, and why its bytes grow with the workers. An all-gather is the second half of a ring all-reduce on its own: every worker contributes a piece and ends up with everyone's pieces. Inside an all-reduce the pieces are chunks of one sum, added on the way. That only works when every worker's vector has the same layout, as dense gradients do. If each worker sends something different, say the 1 value in 16 it judged most important, with its position (chapter 13), the pieces cannot be added as they pass: each must arrive whole. A ring all-gather then has each worker send its own piece and forward the other workers' pieces, $(p-1)$ pieces in all, so the bytes each worker sends grow in proportion to the number of workers, where an all-reduce's stay below twice the vector. That is why compressing a gradient 16-fold saves far less than 16-fold on the wire once there are many workers.

Payload and bytes on the wire. Two numbers describe a sync. The payload is what a worker originates: its dense gradient is 2 bytes per parameter in BF16 (chapter 9), so LLaMA 7B's is 13.5 GB. The wire bytes are what it actually sends in the collective, forwarding included: for a ring all-reduce $2(p-1)/p$ times the payload (that is 26.5 GB per GPU among 64); for an all-gather of compressed payloads $(p - 1)$ times the payload. Link time follows the wire bytes, not the payload; the Compute in Transit results site quotes both, and so does chapter 13.

ZeRO and FSDP: stop every worker holding everything. Chapter 10 counted 16 bytes per parameter for mixed-precision Adam: weights, gradients, two moments and an FP32 master copy, all replicated on every data-parallel worker. ZeRO (Rajbhandari et al. 2020) shards them instead: stage 1 gives each of $d$ workers $1/d$ of the optimiser states (each updates only its own slice of the weights), stage 2 also shards the gradients (a reduce-scatter instead of an all-reduce: each worker receives only its slice's sum), stage 3 also shards the weights themselves, all-gathering each layer's weights just before it is needed. PyTorch's FSDP (Zhao et al. 2023) is stage 3. For LLaMA 7B over 64 GPUs, the model states per GPU fall from 107.8 GB (no sharding) to 28.2 GB (stage 1), 15.0 GB (stage 2) and 1.7 GB (stage 3). Stages 1 and 2 send no more than plain data parallelism (a reduce-scatter and an all-gather are an all-reduce cut in two); stage 3 sends half as much again, because the weights are all-gathered once more in the backward pass (that is 39.8 GB per GPU per step instead of 26.5 ).

Loading the animation…

Step time = compute + communication, and overlap. A data-parallel step costs the compute of the forward and backward passes plus the time to sync the gradients. Nothing forces the sync to wait for the whole backward pass: the last layer's gradients are ready first, so frameworks send them in buckets while the backward pass is still working on earlier layers (PyTorch DDP, Li et al. 2020). When the communication fits inside the remaining backward pass, roughly two thirds of the compute, it costs nothing extra. One step of LLaMA 7B on 64 H100s with 1,048,576 tokens, by the systems model: compute 1.64 s, gradient all-reduce 0.53 s; 2.17 s per step if the sync waits, 1.64 s if it overlaps. The same arithmetic over a much slower link is chapter 13's subject: there the communication outgrows the compute and no overlap can hide it.

MFU: model FLOPs utilisation. How well a cluster is used is measured as model FLOPs utilisation (Chowdhery et al. 2022, PaLM, §4.1 and App. B): the floating-point operations the model needs per second (about $6N$ per token for $N$ parameters, plus attention's $12\,L\,H\,Q\,T$, counted from the tokens actually processed) divided by the hardware's peak. Recomputed activations (chapter 10) don't count, so MFU compares systems fairly; counting them gives hardware FLOPs utilisation. PaLM 540B reached 46.2% MFU on 6,144 TPU v4 chips; Llama 3 405B reached 38–43% on up to 16,384 H100s (Grattafiori et al. 2024, §3.3.2); Megatron-LM's 1-trillion-parameter run reached 52% of peak counting recomputation (Narayanan et al. 2021). In the LLaMA 7B example the step's MFU is 35.7% with the sync waiting and 47.3% with it overlapped: the second is the systems model's ceiling, its fitted matrix-multiply efficiency $\eta =$ 0.473 (chapter 12 shows where that number comes from and how well the model predicts published runs).

The hardware underneath. Racks, chapter 4 runs this ring all-reduce on a rack's NVLink, inside its switches and over the network, and chapter 9 shows where a job's data-parallel traffic lands on racks and the network (the Racks site, in the Hardware track).