Transformer Training — Presentation 10

Data Parallelism, All-Reduce and ZeRO

Give every device the whole model and its own share of the batch, then make the copies agree. The ring all-reduce stepped through, the α–β link model, all-gather and why its bytes grow, payload against wire bytes, ZeRO and FSDP, overlap and MFU.

Data parallelism Ring all-reduce Reduce-scatter All-gather α–β model ZeRO / FSDP MFU
Shard batch → Backward → All-reduce → Same step
00

Topics We’ll Cover

Live companion and sources

This deck follows chapter 11, parallelism I: data, where four workers hold real gradients of the trained model and a ring sums them. Numbers: results.md §11b and §13 (the site’s training-systems model: arithmetic on an α–β link model, validated against published runs in deck 11).

01

Data Parallelism

Every worker (a GPU, or a whole machine) holds a full copy of the model and gets its own share of each batch. Each runs the forward and backward passes on its share; then the workers’ gradients are averaged and every copy takes the same optimiser step, so the copies stay identical.

worker 1: windows 1–4 → g1
worker 2: windows 5–8 → g2
…
→
all-reduce: ¼∑ gi
→
same step everywhere

The average of \(W\) workers’ mean gradients over \(B/W\) windows each is exactly the mean gradient over the whole batch: data parallelism changes where the arithmetic happens, not what it computes. Checked on the model: four simulated workers with four windows each end the default run at 2.299304; one worker with sixteen windows at 2.299304 (the additions happen in a different order, so the last bits differ). More workers means a bigger batch per step, and the critical batch (deck 05) says how far that helps.

02

The Ring All-Reduce, Step by Step

Averaging the gradients is a collective called an all-reduce: every worker starts with its own vector and ends with the sum. The ring cuts 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; then each worker holds the finished sum of one chunk. In \(p-1\) all-gather steps the finished chunks go round until everyone has all of them.

Orange: a chunk just received. Green: a finished sum of all four. Each worker sends \(2(p-1)/p\) of its vector, never more than twice it however many workers there are, which is the least any all-reduce can send (Patarasuk & Yuan 2009). What grows with the ring is the number of steps, \(2(p-1)\), each paying a link’s latency.

03

The α–β Link Model

Sending \(n\) bytes over a link costs \(\alpha + n/B\): a fixed latency and the bytes over the bandwidth. For \(p\) ranks and a vector of \(M\) bytes:

$$\begin{gathered} t_{\text{ring all-reduce}} = 2(p-1)\,\alpha + 2\,\tfrac{p-1}{p}\,\tfrac{M}{B} \\ t_{\text{reduce-scatter}} = t_{\text{all-gather}} = (p-1)\,\alpha + \tfrac{p-1}{p}\,\tfrac{M}{B} \\ t_{\text{tree, not pipelined}} = 2\,\lceil \log_2 p \rceil\,(\alpha + M/B) \end{gathered}$$

H100 figures as the site’s systems model uses them: NVLink 900 GB/s total (450 each way) and 989 TFLOP/s dense BF16 from the H100 SXM datasheet; 400 Gb/s per GPU on the network (Llama 3, §3.3.1). The WAN option is a 100 GbE link (12.5 GB/s) with the network’s latency, an illustration for deck 12.

04

Collective Times for LLaMA 7B

LLaMA 7B’s gradient (6,738,149,376 parameters, 13.48 GB in BF16) on H100s, by the systems model (results.md §11b):

WhereCollectiveTimeBytes each GPU sends
NVLink, 8 GPUsring all-reduce52.5 ms23.58 GB
NVLink, 8 GPUsbinary-tree all-reduce (not pipelined)179.7 ms26.95 GB
NVLink, 8 GPUsreduce-scatter26.2 ms11.79 GB
NVLink, 8 GPUsall-gather26.2 ms11.79 GB
network, 64 GPUsring all-reduce531.9 ms26.53 GB
network, 64 GPUsbinary-tree all-reduce (not pipelined)3.23 s26.95 GB
network, 64 GPUsreduce-scatter265.9 ms13.27 GB
network, 64 GPUsall-gather265.9 ms13.27 GB

A binary tree climbs \(\log_2 p\) levels up and down, so it has fewer latency hops, but done naively (each level waiting for the whole vector) it is several times slower at this size. Real libraries pipeline both.

05

All-Gather, and Why Its Bytes Grow

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.

Same layout: all-reduce

Bytes each worker sends: \(2\,\tfrac{p-1}{p}M\), below twice the vector however large \(p\) grows.

Different layouts: all-gather

If each worker sends something different (say the 1 value in 16 it judged most important, with its position), the pieces cannot be added as they pass: each must arrive whole. Each worker sends its own piece and forwards the others: \((p-1)\) pieces, growing with the number of workers.

That is why compressing a gradient 16-fold saves far less than 16-fold on the wire once there are many workers: the central point of deck 12.

06

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 (deck 08), so LLaMA 7B’s is 13.5 GB. The wire bytes are what it actually sends in the collective, forwarding included:

$$\begin{gathered}\text{all-reduce: } 2\,\tfrac{p-1}{p}\times\text{payload} \\ \text{all-gather (compressed): } (p-1)\times\text{payload}\end{gathered}$$

Among 64 GPUs a ring all-reduce has each GPU send 26.5 GB per step. Link time follows the wire bytes, not the payload; the Compute in Transit results site quotes both, and so does deck 12.

07

ZeRO and FSDP: Stop Every Worker Holding Everything

Plain data parallelism replicates all 16 bytes per parameter on every worker (deck 09). ZeRO (Rajbhandari et al. 2020) shards them across the \(d\) data-parallel workers: stage 1 the optimiser states, stage 2 the gradients too (a reduce-scatter instead of an all-reduce), stage 3 the weights themselves, all-gathered layer by layer just before use. PyTorch’s FSDP (Zhao et al. 2023) is stage 3.

$$\begin{aligned}\text{stage 0: } & (2 + 2 + K)\Psi \\ \text{1: } & 4\Psi + K\Psi/d \\ \text{2: } & 2\Psi + (2 + K)\Psi/d \\ \text{3: } & (2 + 2 + K)\Psi/d\end{aligned}$$
GPUs dStageWeightsGradientsOptimiserTotalSent per step
8013.48 GB13.48 GB80.86 GB107.81 GB23.58 GB
8113.48 GB13.48 GB10.11 GB37.06 GB23.58 GB
8213.48 GB1.68 GB10.11 GB25.27 GB23.58 GB
831.68 GB1.68 GB10.11 GB13.48 GB35.38 GB
64013.48 GB13.48 GB80.86 GB107.81 GB26.53 GB
64113.48 GB13.48 GB1.26 GB28.22 GB26.53 GB
64213.48 GB0.21 GB1.26 GB14.95 GB26.53 GB
6430.21 GB0.21 GB1.26 GB1.68 GB39.80 GB

LLaMA 7B, \(K = 12\) (results.md §11b). Stages 1 and 2 send no more than plain data parallelism (a reduce-scatter plus an all-gather is 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.

08

Step Time = Compute + Communication, and Overlap

A data-parallel step costs the forward and backward compute plus the gradient sync. 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 works on earlier layers (PyTorch DDP, Li et al. 2020). When the communication fits inside the remaining backward pass it costs nothing extra.

Nothing overlapped

Compute 1.64 s + all-reduce 0.53 s = 2.17 s per step; MFU 35.7%.

All-reduce hidden

1.64 s per step; MFU 47.3%, the model’s ceiling (its fitted matrix-multiply efficiency η = 0.473).

One data-parallel step of LLaMA 7B on 64 H100s with 1,048,576 tokens, by the systems model (results.md §11b). Over a much slower link the communication outgrows the compute and no overlap can hide it: deck 12.

09

MFU: Model FLOPs Utilisation

$$\text{MFU} = \frac{D\,(6N + 12\,L\,H\,Q\,T)}{t \cdot n \cdot P}$$

The floating-point operations the model needs per second (about \(6N\) per token for \(N\) parameters, plus attention’s \(12LHQT\), counted from the tokens \(D\) actually processed in time \(t\)) divided by the peak \(P\) of \(n\) devices (Chowdhery et al. 2022, PaLM, §4.1 and App. B). Recomputed activations don’t count, so MFU compares systems fairly; counting them gives hardware FLOPs utilisation.

RunUtilisationSource
PaLM 540B, 6,144 TPU v4 chips46.2% MFUChowdhery et al. 2022
Llama 3 405B, up to 16,384 H100s38–43% MFU (BF16)Grattafiori et al. 2024, §3.3.2
Megatron-LM 1T GPT, 3,072 A100s52% of peak, counting recomputationNarayanan et al. 2021
LLaMA 7B on 64 H100s (systems model)35.7% waiting / 47.3% overlappedresults.md §11b
10

What to Take Away