How one training step is cut into pieces that fit, run at once, and agree at the end.
Chapter 5 ended with two numbers that do not fit on a device: 6.5 TB of optimiser state for Beacon and 3.8 × 10²⁵ operations for its run. Chapter 10 gave you the machine: a GPU with 80 GB and a few hundred usable teraflops, joined to its neighbours by fast links inside a server and slower links between servers. This chapter is the bridge between the two. It takes the training step of Chapter 7, the same forward, bill, backward, nudge, and shows the four ways it is split across sixteen thousand GPUs so that it fits, runs in a few seconds, and produces exactly the update a single impossibly large GPU would have produced. Chapter 12 then asks which architectural choices make that splitting cheaper; Chapter 13 runs the run.
The question this chapter answers: a training step needs terabytes of state and quintillions of operations; how is it divided across thousands of GPUs without changing the mathematics, and what does the division cost?
Sixteen thousand translators, one encyclopedia, a deadline. There are only three ways to split the work, and a good project uses all three at once.
Give each translator different pages. Everyone works independently on their own volume. But the encyclopedia must read as one book, so every evening the translators meet and reconcile their glossaries: each one's discoveries about how to render a term get averaged into a shared glossary that everyone uses tomorrow. The meeting is the cost. It grows with the size of the glossary, not with the number of pages.
Split each page among several translators. One does the nouns, one the verbs, one the idioms, and they pass fragments back and forth constantly. This only works if they sit at the same table, because they talk after every sentence.
Run an assembly line. Team A translates, passes the page to team B who edits, who passes to team C who typesets. Every team is always busy, once the line has filled. But while A is translating the first page, B and C stand idle, and while C typesets the last page, A and B do. The idle time is the bubble, and the way to shrink it is to pass many small batches of pages down the line rather than one big one.
Data parallelism is the first, with the all-reduce as the evening meeting. Tensor parallelism is the second, kept inside one server where the links are fast. Pipeline parallelism is the third, with micro-batches to fill the line. And the sharing of the glossary's storage across the translators, so no one has to carry the whole thing, is ZeRO.
| In the picture | In the machine | The word we will use |
|---|---|---|
| Different pages per translator, same glossary | Each GPU holds a full model copy and a different slice of the batch; gradients are averaged | data parallelism (DP) |
| The evening reconciliation | Every GPU ends with the same averaged gradient | all-reduce |
| Not everyone carries the whole glossary | Optimiser state, gradients, and weights sharded across the DP group | ZeRO stages 1–3, FSDP |
| One page split between translators at the same table | One matrix multiplication split across GPUs in a server | tensor parallelism (TP) |
| The assembly line | Consecutive blocks on different GPUs | pipeline parallelism (PP), stages |
| Idle teams while the line fills and drains | Fraction of time stages wait | pipeline bubble |
| Small batches of pages down the line | The batch split into pieces that flow through stages | micro-batches, 1F1B schedule |
| A page too long for one table, split by paragraphs | One sequence's tokens split across GPUs | context (sequence) parallelism (CP) |
| Fraction of the day actually spent translating | Achieved FLOP/s over peak | MFU |
Start from Chapter 5's numbers. Beacon has 405 billion parameters. To train it, each parameter needs about 16 bytes: 2 for the working bf16 weight, 2 for its gradient, 4 for a 32-bit master copy, and 8 for Adam's two moments. That is 6.5 TB before a single activation is stored. A GPU has 80 GB (H100) or 192 GB (B200) public. The state alone needs at least 34 to 81 GPUs' worth of memory, and activations, the intermediate values kept for the backward pass, multiply that.
Then time. The run is 3.8 × 10²⁵ operations public. A single H100 sustaining 4 × 10¹⁴ useful operations per second (about 40% of its bf16 peak, Chapter 10) would need 9.5 × 10¹⁰ seconds: three thousand years. Sixteen thousand of them, perfectly parallel, would need 68 days. Llama 3 405B was trained on 16,384 H100s public, and the run took on the order of two months, which tells you the parallelism was far from perfect and yet good enough.
So the split must do two things at once: divide the state so that each GPU's share fits, and divide the work so that all GPUs are busy nearly all the time. The constraint that shapes every choice is that the split must not change the mathematics. At the end of every step, the model must have taken exactly the update of Chapter 5, as if one enormous GPU had done it. Everything below is a way of arranging computation and communication so that this remains true.
The simplest split, and the one every run uses at the outermost level. Make DP copies of the model. Give each copy a different slice of the batch: with a 16M-token batch and DP = 128, each copy sees 125k tokens. Each copy does its own forward and backward on its slice and ends up with its own gradient, which is the average over its tokens. The gradient the mathematics wants is the average over all tokens, which is the average of the copies' gradients. So the copies exchange gradients and average them, and then each applies the identical update. After the step, all copies are still identical. That exchange is the all-reduce.
The naive all-reduce, "everyone sends everything to everyone", costs DP times the gradient size per GPU and would swamp any network. The ring algorithm, which is what NCCL implements, costs each GPU 2·(DP − 1)/DP times the gradient size, less than twice the gradient, regardless of how many GPUs are in the ring public.
One GPU in Beacon's run holds a 1/(TP·PP) = 1/128 shard of the model: 3.16 B parameters, whose bf16 gradient is 6.3 GB. Across a DP ring of 128 over 400 Gb/s InfiniBand (about 50 GB/s):
bytes each GPU must send 2 × (127/128) × 6.3 GB ≈ 12.5 GB time at 50 GB/s ≈ 0.25 s compute per step per GPU 6 × 405e9 × 16e6 / 16,384 / 4e14 ≈ 5.9 s
$ python code/ch11/bubbles_and_rings.pyPP= 16 micro-batches= 16 bubble 48.4% interleaved v=2: 31.9% PP= 16 micro-batches= 64 bubble 19.0% interleaved v=2: 10.5% PP= 16 micro-batches= 256 bubble 5.5% interleaved v=2: 2.8% all-reduce 6.3 GB across DP=128 over NVLink (≈ 450 GB/s per direction): 0.03 s all-reduce 6.3 GB across DP=128 over InfiniBand 400 Gb/s (≈ 50 GB/s): 0.25 s compute per step per GPU: 6 × 405e9 × 16e6 / 16384 GPUs / 4e14 FLOP/s ≈ 5.9 s
A quarter of a second of communication against six seconds of compute: 4%, if it were not overlapped at all. In practice it is overlapped. The backward pass produces gradients layer by layer from the top, so the all-reduce of block 126's gradient can begin while blocks 125 and below are still computing. Bucket the gradients into chunks, launch each chunk's all-reduce as soon as it is ready, and the communication hides almost entirely behind the backward pass. Data parallelism is cheap. Its problem is not time but memory: every one of the 128 copies holds the full state of its shard, 16 bytes per parameter, and nothing about data parallelism reduces that.
Look at what the 128 data-parallel copies hold that is identical: the same weights, the same optimiser state, and after the all-reduce, the same gradients. The insight of ZeRO is that this redundancy is unnecessary. Each copy only needs its part of the state at the moment it is used.
Stage 1 shards the optimiser state. GPU k in the DP group keeps Adam's m, v, and the fp32 master weights for a 1/DP slice of the parameters and updates only that slice; after the update, the updated bf16 weights for each slice are all-gathered so every copy has the full working weights again. Since the optimiser state is 12 of the 16 bytes per parameter, this alone cuts memory by nearly 4×. Stage 2 also shards the gradients: instead of all-reducing the full gradient to everyone, reduce-scatter it so each GPU receives only the averaged gradient of its own slice, the slice it will update. Same communication volume, since reduce-scatter is half of the ring all-reduce. Stage 3 shards the weights themselves: no GPU holds the full model; each block's weights are all-gathered just before that block is used, in forward and again in backward, and discarded after. Per-GPU state falls to 16·N/DP bytes at a cost of about 1.5× the communication of plain data parallelism public (the ZeRO paper's accounting). Stage 3 is what PyTorch calls fully sharded data parallelism, FSDP, and it is what Llama 3 used within its data-parallel groups public.
per-GPU bytes per parameter of the shard DP = 128 plain data parallel 2 + 2 + 12 = 16 → 16 ZeRO-1 (optimizer) 2 + 2 + 12/DP → 4.1 ZeRO-2 (+ gradients) 2 + 2/DP + 12/DP → 2.1 ZeRO-3 (+ weights) (2 + 2 + 12)/DP → 0.125 (plus one block's gathered weights at a time)
ZeRO alone can make a 70B model trainable on ordinary data parallelism across a few dozen GPUs. It does not, however, help with the other two problems: a single block's matrices may still be too large for one GPU's memory bandwidth to make efficient, and the activations of a long sequence still have to live somewhere. That is where the model-parallel axes come in.
A matrix multiply splits cleanly. Take the MLP of Chapter 2: y = act(x·W_in)·W_out. Cut W_in into TP column blocks and give one to each GPU. Every GPU receives the full input x (it is only [T, C], small compared to the weights) and computes its own 1/TP slice of the hidden activations, act(x·W_in[:, slice]). The activation is elementwise, so no communication is needed yet. Now cut W_out into the matching TP row blocks; each GPU multiplies its hidden slice by its row block and gets a partial sum of y. The full y is the sum of the partials: one all-reduce. Attention splits the same way, by heads: each GPU owns H/TP heads, computes them entirely locally, and the partial outputs through its rows of W_O are all-reduced.
The cost is what makes tensor parallelism a within-server technique. Each block needs two all-reduces in forward and two in backward, of activations of shape [T, C], and they sit on the critical path: nothing downstream can proceed until the sum is known. Over NVLink at hundreds of GB/s this is tolerable; over InfiniBand it is not. Every published configuration keeps TP at or below the number of GPUs in one server, 8 public (Megatron-LM's recommendation; Llama 3's TP of 8). What tensor parallelism buys is the ability to hold one block across 8 GPUs, so that each GPU's slice of a 16k-wide matrix is small enough to be efficient and the activations of one block fit.
The third axis puts different blocks on different GPUs. With PP = 16 and 126 blocks, stage 1 holds blocks 1 to 8, stage 2 blocks 9 to 16, and so on. A micro-batch of tokens flows through stage 1, its output activations cross to stage 2, and so on to the end; the backward pass flows the other way. Communication is a single [T, C] tensor per stage boundary, per direction, which is tiny. The problem is idleness.
If the whole batch went through as one piece, stage 16 would sit idle while stages 1 to 15 processed it, then stage 1 would idle during the backward pass. Only one stage would ever be busy. The fix is micro-batching: cut the batch into m pieces and push them down the line one after another, so that once the line is full, all stages work on different micro-batches at once. The bubble does not vanish; it is the fill and drain at each end, and its fraction of the step is
bubble = (p − 1) / (m + p − 1)
for p stages and m micro-batches. Sixteen stages and sixteen micro-batches waste half the step; 256 micro-batches waste 5%. The 1F1B schedule (one forward, one backward, alternating once the line is full) keeps this bubble but caps the activations each stage must hold at p micro-batches' worth rather than m public. Interleaving, where each GPU holds several non-adjacent groups of blocks, divides the bubble again by the number of groups at the cost of more boundary traffic.
Two refinements are worth knowing by name. Micro-batches are the same thing as gradient accumulation: the optimiser step waits until all m micro-batches' gradients have been summed, so the effective batch is m × micro-batch × DP. And DeepSeek-V3's DualPipe schedule runs forward and backward of different micro-batches simultaneously on the same stage, overlapping the communication of one with the compute of the other, which is how they trained a 671B MoE on 2,048 H800s with limited inter-node bandwidth public.
Two more axes appear in current runs. Context parallelism splits a single very long sequence across GPUs by tokens. Each GPU holds the activations for its slice of positions, which is what makes 128k-token sequences fit at all: the activations of one 128k sequence at C = 16384 are tens of gigabytes per layer. Attention then needs every query to see every earlier key and value, so the key/value blocks are passed round the context-parallel ring, one slice at a time, while each GPU accumulates its queries' outputs, the FlashAttention tiling of Chapter 10 stretched across machines public (Llama 3 used exactly this in its long-context stage). Expert parallelism belongs to mixture-of-experts models, Chapter 12: the experts of an MLP layer are placed on different GPUs, and each token's activations are sent to the GPU holding the expert it was routed to, an all-to-all exchange instead of an all-reduce.
Llama 3's 405B run used what its report calls 4D parallelism: TP 8 within each server, PP 16 across servers in a pod, DP 128 across the cluster, and CP for the long-context phase, with FSDP-style sharding of optimiser state, gradients, and weights inside each DP group public. The report gives 38 to 43% MFU for the main phase, around 400 TFLOP/s per GPU public. Beacon uses the same layout. Here is the memory on one of its GPUs, from the shapes.
16,384 GPUs = TP 8 × PP 16 × DP 128 model-parallel shard factor 8 × 16 = 128 → 405 B / 128 = 3.16 B parameters per GPU layers per pipeline stage 126 / 16 ≈ 8 tokens per DP replica per step 16 M / 128 = 125,000 (≈ 15 sequences of 8k) static state per GPU (bytes per parameter of the shard) weights bf16 3.16 B × 2 = 6.3 GB gradients bf16 3.16 B × 2 = 6.3 GB Adam state + master fp32 3.16 B × 12 = 38.0 GB → sharded across DP 128 (ZeRO-1): 0.3 GB activations (estimate: ≈ 34·C bytes per token per layer, no T² term with FlashAttention) 125k tokens × 8 layers × 34 × 16,384 / 8 (TP) ≈ 68.5 GB if a whole replica-batch were resident with 8 micro-batches ≈ 8.6 GB with activation checkpointing (keep block inputs only) ≈ 1.6 GB
$ python code/ch11/shard_memory.pyNo ZeRO, no checkpointing weights 6.3 grads 6.3 opt 38.0 acts 68.5 total 119.2 GB/GPU ZeRO-1 (optimizer sharded) weights 6.3 grads 6.3 opt 0.3 acts 68.5 total 81.5 GB/GPU ZeRO-1 + activation checkpointing weights 6.3 grads 6.3 opt 0.3 acts 12.7 total 25.7 GB/GPU ZeRO-3/FSDP + checkpointing weights 0.0 grads 0.0 opt 0.3 acts 12.7 total 13.1 GB/GPU
Read the first line: with the naive arrangement, one GPU would need 119 GB and has 80. Shard the optimiser and it is at the edge. Add micro-batching and checkpointing and it is comfortably inside, with room for the CUDA context, communication buffers, and the temporary gathered weights of the block being computed. This is the reason the layout has four dimensions: no single one of them would fit.
Three techniques run underneath all of the above and are worth naming precisely because they appear in every training log.
Mixed precision. The matrix multiplies run in bf16, 2 bytes per number with an 8-bit exponent (same range as fp32, less precision), because tensor cores are several times faster in 16-bit and memory traffic halves. But an Adam update of 10⁻⁵ on a weight of magnitude 1 is below bf16's resolution: added in bf16, it would round to nothing. So a 32-bit master copy of each weight receives the updates, and a bf16 copy is derived from it for the next forward pass. That master copy is the "4" in 16 bytes per parameter. The older fp16 format had a smaller exponent and needed loss scaling, multiplying the loss by a large constant so small gradients did not underflow, then dividing the gradients back; bf16 made that unnecessary and is now universal public. DeepSeek-V3 pushed further, running most matmuls in 8-bit floating point with fine-grained scaling and keeping accumulation in higher precision, and reported it matched bf16 quality public.
Gradient accumulation. When the batch the mathematics wants is larger than the batch the memory allows, run several micro-batches, sum their gradients, and step once. The result is identical to one large batch (the gradient of an average is the average of gradients); only the wall-clock changes. Pipeline micro-batches are gradient accumulation with the stages overlapped.
Activation checkpointing. The backward pass needs the input of every operation (Chapter 5: the gradient on W is xᵀ·δ, so x must be there). Keeping every sub-layer's input for 8k tokens across 126 layers is the dominant memory cost. Checkpointing keeps only each block's input, throws away everything inside, and recomputes the block's forward pass during backward. Cost: one extra forward, about a third more compute per step. Benefit: activation memory falls from "every intermediate" to "one vector per block per token". Every frontier run uses some form of it, often selectively, recomputing only the cheap-but-large intermediates public.
At this scale, hardware fails constantly, and the training loop is synchronous: every step, all 16,384 GPUs must finish before the update. One GPU that stalls stalls the run. Meta reported that during a 54-day snapshot of the 405B run there were 466 job interruptions, 419 of them unexpected, about 78% attributed to confirmed or suspected hardware problems, with GPU faults the single largest category, and yet over 90% effective training time public. The tools that make this survivable are unglamorous. Checkpoints of the full state (6.5 TB) are written to distributed storage frequently enough that a failure costs minutes of progress, not hours; the storage must absorb terabytes in seconds so the GPUs do not wait. Automated detection identifies the failed node, the job is restarted from the last checkpoint on a replacement, and the scheduler keeps spare nodes warm. Silent data corruption, where a GPU returns wrong numbers without an error, is caught by watching the loss for spikes that do not reproduce on replay and by periodic numerical self-tests public.
MFU, model FLOP utilisation, is the summary statistic: the fraction of peak arithmetic that went into the model's 6N operations per token. Everything above eats into it: pipeline bubbles, exposed communication, recomputation for checkpointing (which counts against MFU, since it is not "model" FLOPs), stragglers, restarts. Published frontier runs sit between 35 and 45% public. A run at 40% MFU on 16k H100s for two months is roughly 3.8 × 10²⁵ operations, which closes the loop with Chapter 5.
| Quantity | Llama 3.1 405B | DeepSeek-V3 (671B MoE) | Evidence |
|---|---|---|---|
| GPUs | 16,384 H100 | 2,048 H800 | public |
| Parallelism | TP 8 × PP 16 × DP 128 (+ CP for long context), FSDP sharding | PP 16 (DualPipe), EP 64, ZeRO-1 DP; no TP | public |
| Precision | bf16 with fp32 master | FP8 matmuls, fine-grained scaling | public |
| MFU | 38–43% | not reported as MFU; 2.79 M GPU-hours total | public |
| Interruptions | 466 in 54 days, 419 unexpected, ≈78% hardware | not reported | public / unknown |
| Closed frontier models | Cluster sizes of tens of thousands of accelerators are stated or inferred from disclosed compute; layouts not disclosed | inferred | |
What does a checkpoint cost? The full training state is the thing that must be saved: fp32 master weights plus Adam's two moments, 12 bytes per parameter, plus bf16 weights if saved separately.
state to save 12 × 405 B ≈ 4.9 TB (plus 0.8 TB bf16 weights) storage write bandwidth 2 TB/s (a large parallel file system, public figures for such clusters) time, if the GPUs waited ≈ 2.5 s per checkpoint → checkpoint every 30 min costs 0.14% of the run progress lost per failure ≤ 30 min of 16,384 GPUs ≈ 8,200 GPU-hours ≈ the cost of the restart, not the checkpoint
Checkpointing is cheap; failures are not. With 419 unexpected interruptions in 54 days, one every three hours, and a restart that takes minutes plus the lost half-hour, the arithmetic says roughly 5 to 10% of the run is spent recovering, which is consistent with the reported "over 90% effective time". Reducing the interval between checkpoints, and writing them asynchronously so the GPUs never wait, is worth more than almost any other engineering at this scale.
Use data parallelism only. Every GPU must hold the whole model's state, 6.5 TB. It does not fit on any GPU, or on any server; the run cannot start. Even with ZeRO-3 sharding the state across all 16k GPUs (0.4 GB each), every block's full weights must be all-gathered before use, 16 GB of traffic per block per pass, and the activations of a single 8k sequence through a 16k-wide model exceed one GPU's memory. Model parallelism is not an optimisation; it is a requirement above roughly 20B parameters.
Run tensor parallelism across servers. Two all-reduces per block per direction, 500 per step, each a synchronous wait on a 50 GB/s link instead of a 450 GB/s one. The critical path grows by tens of seconds per step; MFU collapses toward 10%. This is why every published layout keeps TP inside the NVLink domain and uses pipeline parallelism, with its once-per-stage traffic, to cross servers.
Use one micro-batch with sixteen pipeline stages. Bubble fraction 15/16: each stage works 6% of the time. Sixteen thousand GPUs deliver the throughput of a thousand. The remedy costs nothing but a loop: split the batch into 64 or 256 micro-batches.
Skip the fp32 master weights. Adam's typical update, η ≈ 10⁻⁴ times a unit-scale normalised step, is smaller than the spacing between adjacent bf16 values near 1 (about 8 × 10⁻³). Most updates round to zero; the model learns only where gradients are unusually large; the loss curve flattens early. Four bytes per parameter of master copy is what keeps a million small nudges from being lost.
Never checkpoint. One unexpected interruption every three hours, and each one restarts the run from the beginning. The run never completes. At frontier scale, checkpointing is as much a correctness requirement as the causal mask.
Let one straggling GPU be 20% slower. The step is synchronous: all 16,384 wait for the slowest. The whole cluster runs at 80% of its speed because of one device. Detecting and evicting stragglers, and choosing layouts where a slow node hurts only its pipeline stage rather than the global all-reduce, is a real part of the engineering.
Say it back. A frontier training step needs 16 bytes per parameter of state and 6N operations per token, neither of which fits on one device, so the step is cut along independent axes whose product is the GPU count. Data parallelism copies the model and splits the batch; the copies agree by an all-reduce whose cost per GPU is under twice the gradient size regardless of copy count, and which hides behind the backward pass because top-layer gradients finish first. ZeRO removes the redundancy between copies by sharding optimiser state, then gradients, then weights across the data-parallel group, at about 1.5× the communication. Tensor parallelism splits each matrix across the GPUs of one server, column-wise then row-wise, with an all-reduce of the activations after every sub-layer, which is why it stays inside the fast links. Pipeline parallelism puts consecutive blocks on consecutive servers, moving only one activation tensor per boundary, and pays a bubble of (p − 1)/(m + p − 1) that micro-batching shrinks. Context parallelism splits long sequences by tokens and passes key/value blocks round a ring. Underneath, matmuls run in bf16 against an fp32 master copy, micro-batches accumulate into the mathematically identical large-batch gradient, and activation checkpointing trades a third more compute for most of the activation memory. Sixteen thousand devices fail every few hours, so the state is checkpointed often and the job restarts automatically. The result is 35 to 45% of peak arithmetic going into the model, and a run that is exactly the update Chapter 5 described, taken sixteen thousand ways at once.
C = 8192, 80 layers) is trained on 512 GPUs with TP 8, PP 8, DP 8, ZeRO-1, and a 4M-token batch. Compute per-GPU bytes for bf16 weights, bf16 gradients, and optimiser state; tokens per DP replica; layers per stage; and the pipeline bubble for 32 micro-batches. Then repeat with ZeRO-3 and say which line changed.shard_memory.py with a search: for a given model and GPU count, enumerate all (TP, PP, DP) with TP ≤ 8 whose product is the GPU count, compute per-GPU memory under ZeRO-1 with checkpointing and 8 micro-batches, and print the configurations that fit in 80 GB sorted by pipeline bubble. Which one would you choose for Beacon, and why does it match, or not match, the published layout?