A model too large for one accelerator must be split across many. The splitting is never free — and the cost, like everything else here, is a short calculation you can do from a timing diagram.
A model that will not fit on one accelerator must be cut into pieces and spread across many. Nothing about that is free. Cut it into sequential stages and the accelerators spend part of their time waiting on each other; spread its state across devices and you must first know how large that state actually is. Both costs are countable, and counting them is this chapter.
Assign consecutive layers to consecutive devices — pipeline parallelism — and no device can begin until the one before it has finished, so the fleet idles through part of every step. Divide the optimizer state instead and the arithmetic is trivial division, but the figure you must compute first is the undivided total, because that is what tells you whether the run is possible at all. Both penalties shape real systems, and both are a few lines of counting.
Split a network into \( P \) sequential stages, one per accelerator. A batch flows through stage 1, then stage 2, and so on. If we fed one whole batch through, only one accelerator would be busy at a time — catastrophic waste. So we cut the batch into \( M \) micro-batches and stream them, so that while micro-batch 2 is in stage 1, micro-batch 1 is already in stage 2. But the pipeline still has to fill at the start and drain at the end, and during fill-and-drain some accelerators sit idle. That idle fraction is the bubble.
The bubble fraction from a timing diagram
Let each micro-batch take one unit of time per stage. Draw the schedule: micro-batch \( m \) can enter stage \( s \) only after it has left stage \( s-1 \) and after micro-batch \( m-1 \) has left stage \( s \). Tracing the diagonal, the last micro-batch enters stage 1 at time \( M-1 \) and exits the final stage \( P \) at time \( M + P - 1 \). So the whole batch completes in \( M + P - 1 \) time units.
The useful work is \( MP \) stage-units (each of \( M \) micro-batches through each of \( P \) stages). The total accelerator-time spent is \( P(M + P - 1) \) (all \( P \) accelerators, held for the whole duration). The idle — bubble — fraction is
\[ 1 - \frac{MP}{P(M+P-1)} = \frac{P-1}{M+P-1}. \]Two boundary cases carry the meaning. With \( M = 1 \) — no micro-batching, naive model parallelism — the bubble is \( (P-1)/P \), which for \( P = 8 \) is \( 7/8 = 87.5\% \) idle. That catastrophe is the entire reason micro-batching exists. As \( M \) grows large, the bubble tends to zero: more micro-batches amortize the fill-and-drain.
How many micro-batches to hide the bubble
For an 8-stage pipeline, how many micro-batches keep the bubble at or below 5%? Require \( \frac{7}{M+7} \le 0.05 \), i.e. \( M + 7 \ge 140 \), so \( M \ge 133 \). You need well over a hundred micro-batches per optimizer step just to keep one-twentieth of your hardware from sitting idle — which is why pipeline parallelism is a tool of last resort, reached for only when a model will not fit any other way.
A 16-stage pipeline runs with \( M = 48 \) micro-batches. What fraction of accelerator-time is lost to the bubble?
\( (P-1)/(M+P-1) = 15/(48+15) = 15/63 = 0.238 \), about 24%.
The second penalty is memory, and its headline number is worth building rather than accepting. Training keeps far more per parameter than the parameter, and what it keeps is countable.
Sixteen bytes for every parameter
Modern training runs in mixed precision: the arithmetic is done in half precision because that is what the hardware is fast at, while a full-precision copy of each weight is kept so that many small updates are not lost to rounding. Count what has to be resident, per parameter, to take one step.
\( \quad \) Half-precision weight, read by the forward and backward passes — \( \textbf{2} \) bytes.
\( \quad \) Half-precision gradient, produced by the backward pass — \( \textbf{2} \) bytes.
\( \quad \) Full-precision master weight, which the update is actually applied to — \( \textbf{4} \) bytes.
\( \quad \) Full-precision first moment, the optimizer's running mean of gradients — \( \textbf{4} \) bytes.
\( \quad \) Full-precision second moment, its running mean of squared gradients — \( \textbf{4} \) bytes.
The two moments are what an adaptive optimizer keeps so that each coordinate's step can be scaled by its own gradient history — the cheap approximation to per-axis conditioning sketched in the challenge at the end of Chapter 3. They are why the figure is 16 and not 8: plain gradient descent would keep no moments at all, and the same run would need half the memory and converge far worse.
Note which entries a technique could hope to remove, because it decides what any memory saving can possibly be worth. Freezing a parameter removes its gradient and all three full-precision copies and leaves only the 2 bytes that must be read to compute anything at all — fourteen of the sixteen. That is the arithmetic underneath Chapter 13's low-rank adaptation, and it is why the saving there is two orders of magnitude rather than a factor of two.
Multiply by the parameter count and the total is enormous; the remedy is to shard that state across many accelerators, each holding a slice. The arithmetic is mere division — but you should always compute the unsharded figure first, because that is the wall the sharding exists to climb.
The wall, and the ladder over it
A 50-billion-parameter model at 16 bytes of optimizer state per parameter needs \( 16 \times 5\times 10^{10} = 8\times 10^{11} \) bytes — 800 gigabytes — of state. No single accelerator holds that; it is the memory wall. Shard it with no redundancy across 64 accelerators and each holds \( 800/64 = 12.5 \) gigabytes, which fits comfortably. The sharding did nothing clever — it divided by 64 — but computing the unsharded 800 GB first is what shows you why the division was necessary.
A 13-billion-parameter model at 16 bytes/parameter of state is sharded across 32 accelerators. Give the unsharded total and the per-accelerator share.
Unsharded: \( 16 \times 1.3\times 10^{10} = 2.08\times 10^{11} \) bytes \( = 208 \) GB. Per accelerator: \( 208/32 = 6.5 \) GB.
The bubble is a fraction of accelerator-time, not an amount of time, and the two get confused constantly. A 17.9% bubble on an eight-stage pipeline does not mean seventeen-point-nine of anything — it means that across the whole fleet, for the whole step, that proportion of available accelerator-seconds produced nothing. Doubling the micro-batch count halves the fraction while lengthening the wall-clock step, because more micro-batches means more work per step. Efficiency and latency move in opposite directions here, and a claim about one says nothing about the other.
The same caution applies to sharding. “12.5 GB per accelerator” is a stock; “800 GB of state” is the stock that matters for feasibility. Quoting only the first hides whether the run is possible at all.
Weights and optimizer state are not the whole bill. The backward pass needs the activations the forward pass produced — Section 2.3 — so every intermediate result must be held from the moment it is computed until the gradient reaches it. That store scales with quantities the previous section never mentioned: the batch, the sequence length, and the depth.
Guess the size before computing it. Optimizer state for a 7-billion-parameter model is 112 GB and does not depend on the batch at all; activations depend on nothing else so strongly. Most people expect activations to be the smaller term. At training batch sizes they are routinely the larger one, and the reason is that they are counted per token rather than per parameter.
What the forward pass leaves behind
A transformer layer of width \( d \) stores, per token, a handful of intermediate vectors: the block's input, the three projections, the attention output, and the two feed-forward activations at widths \( d \) and \( 4d \). Call it \( c\,d \) values per token per layer, with \( c \) around 10 once everything is counted.
For \( d = 4096 \), \( L = 32 \), half precision, a batch of 8 sequences of 4096 tokens — \( B s = 32{,}768 \) tokens in flight:
\[ 10 \times 4096 \times 32 \times 32{,}768 \times 2 \ \text{bytes} = 8.6\times10^{10} = 86 \ \text{GB}. \]Against 112 GB of optimizer state and 14 GB of weights, the activations are the second-largest term and within a factor of two of the largest — on a batch most practitioners would call modest. Double the sequence length and they pass everything else.
The remedy is the trade this book keeps meeting: spend the abundant resource to relieve the binding one. Store the activations at only a few points in the network — checkpoints — and recompute the rest during the backward pass from the nearest one below.
The square root of depth
Split \( L \) layers into \( g \) equal segments, storing only each segment's input. Two costs result. The stored activations are the \( g \) checkpoints, plus, while recomputing inside one segment, that segment's \( L/g \) layers:
\[ \text{stored} \propto g + \frac{L}{g}. \]Differentiate with respect to \( g \) — Chapter 1's constrained-minimum reflex, with no constraint to substitute this time. Setting \( 1 - L/g^2 = 0 \) gives \( g = \sqrt{L} \), and the stored total is \( 2\sqrt{L} \) rather than \( L \).
The price is one extra forward pass over each segment during the backward sweep: about \( 2N \) FLOPs per token on top of the \( 6N \) of Chapter 2, so roughly a third more arithmetic. At \( L = 64 \) the memory falls by a factor of \( 64/(2\times 8) = 4 \) for that third. Memory is bought with arithmetic at a rate the roofline of Chapter 11 will make it easy to judge.
(a) For \( L = 100 \) layers, give the optimal number of checkpoint segments and the factor by which stored activations fall. (b) A run holds 86 GB of activations and 112 GB of optimizer state. Checkpointing at \( \sqrt{L} \) with \( L = 32 \); what is the new activation figure, and which term now binds?
(a) \( g = \sqrt{100} = 10 \); stored goes from \( 100 \) to \( 2\sqrt{100} = 20 \), a factor of 5. (b) \( 2\sqrt{32} = 11.3 \) against 32, so activations fall to \( 86 \times 11.3/32 = 30 \) GB. Optimizer state now dominates, and the next thing to attack is the 16 bytes of Derivation 8.2 rather than the activations.
One cost remains, and it is the one that decides which axis of parallelism you reach for. Splitting work across devices means the devices must talk, and their link is slower than their memory by roughly the same margin that memory is slower than arithmetic.
Data parallelism — every device holds the whole model and processes a different slice of the batch — requires that the gradients be summed across devices before the step. That sum moves every gradient once: \( 2N \) bytes per device per step in half precision, whatever the batch size. Tensor parallelism, which splits individual matrices across devices, instead exchanges activations at every layer, so its traffic scales with tokens in flight rather than with parameters. Pipeline parallelism moves only the activations at stage boundaries, which is the least traffic of the three — and buys that with the bubble of Derivation 8.1.
A 7-billion-parameter model is trained data-parallel in half precision over a link carrying 50 GB/s. Give the bytes summed per device per step, and the time that takes. If a step's arithmetic occupies 400 ms, what fraction of the step is communication?
\( 2N = 1.4\times10^{10} \) bytes \( = 14 \) GB, taking \( 14/50 = 0.28 \) s. Against a 400 ms step that is 280 ms of talking to 400 ms of working — 41% of the step, and the reason gradient summation is overlapped with the backward pass rather than performed after it.
None of the penalties of scale is mysterious; all submit to counting. The pipeline bubble, \( (P-1)/(M+P-1) \), comes straight from a timing diagram and explains why micro-batching is mandatory and why pipelines are a last resort. Memory sharding is division, but the unsharded figure is the wall that motivates it. Activations are counted per token rather than per parameter, which is why they overtake the optimizer state on any serious batch, and checkpointing buys them back at \( 2\sqrt{L} \) for about a third more arithmetic. Communication is the fourth cost, and the axis you choose is a choice about which of parameters or tokens you would rather move. These are the arithmetic facts that shape how the largest training runs are laid out across hardware — and they return, transformed, in Chapter 11, where the same roofline thinking governs not training but the economics of serving a trained model one token at a time. Before that, Chapter 9 confronts a subtler problem: once models are this large and this capable, how do we know what they can actually do — and when a benchmark is lying to us?
A · Drills
B · Problems
C · Challenge
Reproduce Derivation 8.1 from a hand-drawn timing diagram — the completion time and the bubble fraction — on blank paper, and evaluate both boundary cases. Attempt B-1 and B-2 closed book. You pass when you can draw the staircase, read \( M+P-1 \) off it, and produce the \( 87.5\% \) idle figure for a naive 8-stage pipeline without hesitation.