Training Infrastructure and Systems

A 70-billion-parameter model does not fit on one GPU. Not the weights, not the gradients, and definitely not the optimizer. And even if it did fit, one GPU grinding through 15 trillion tokens would take centuries. So the whole game of large-scale pretraining is a systems game: how do you spread one model over thousands of chips and keep every chip busy?

This chapter is the "plumbing" chapter. There are no new learning ideas in it. Everything is about memory, bandwidth, and keeping the matrix multiplies fed. But it is the plumbing that decides whether your run costs one week or four, so it is worth understanding properly.

We will build up in the natural order. First, count bytes: where does the memory go? Second, shrink the bytes: mixed precision. Third, split the work: data, tensor, pipeline, sequence, and expert parallelism. Fourth, move the bytes faster: communication and FlashAttention. Fifth, survive: checkpointing and hardware failures. Sixth, measure: tokens per second and MFU.

The problem: one GPU is too small and too slow

Start with the two numbers that matter. An NVIDIA H100 has 80 GB of high-bandwidth memory and a peak of roughly 989 trillion bf16 floating-point operations per second (dense, no sparsity tricks). Those are the walls we keep hitting.

Memory first. Training a model needs far more than the weights. You need the gradients, the optimizer's running statistics, and the activations saved during the forward pass so that backprop can use them. We will do the arithmetic in a moment, but the punchline is that with Adam and mixed precision, training needs about 16 bytes per parameter before you even count activations. A 70B model needs about 1.1 terabytes. That is fourteen H100s just to hold the state, doing nothing yet.

Now time. Recall from the scaling-laws chapter that training costs roughly $6ND$ floating-point operations for $N$ parameters and $D$ tokens. For $N = 7\times 10^{10}$ and $D = 1.5\times 10^{13}$ that is $6.3\times 10^{24}$ FLOPs. Divide by one H100 running perfectly at $10^{15}$ FLOP/s and you get $6.3\times 10^{9}$ seconds, which is about 200 years. Nobody waits 200 years.

Memory accounting: where do the bytes go?

Before we can split anything, we have to know what we are splitting. Let us count the bytes used to train a model with $N$ parameters using Adam in the standard mixed-precision recipe.

The 16 bytes per parameter

Each parameter carries five pieces of state during training. Here is the list, with the sizes used in the usual bf16 mixed-precision setup:

So the state alone is $16N$ bytes. This is the number that the ZeRO paper (Rajbhandari et al. 2020) built its whole argument on: 2 + 2 for the half-precision weights and gradients, plus 12 for what they call the "optimizer states" (master weights and both Adam moments).

Worked example: 7B and 70B

Activations, and activation checkpointing

The state is the fixed cost. The variable cost is activations: every intermediate tensor the forward pass produces that backprop will need later. For a transformer they scale with sequence length times batch size times hidden width, per layer.

Korthikanti et al. (2022) worked out the per-layer budget for a standard transformer block. With sequence length $s$, micro-batch $b$, hidden width $h$, and $a$ attention heads, storing everything in bf16 costs roughly:

The 34 covers the inputs to every linear, the norms, the GeLU/SwiGLU input, and the dropout masks. The second term is the attention score matrix, one $s \times s$ block per head, and it is the term that explodes with sequence length. FlashAttention (later in this chapter) makes that second term disappear because the scores are never stored, leaving roughly $34\,s\,b\,h$ bytes per layer.

The fix is activation checkpointing (Chen et al. 2016), also called gradient checkpointing or recomputation. Instead of saving every intermediate, save only the input to each block. During backward, recompute the block's forward pass on the fly to regenerate what you need. Memory drops to about $2\,s\,b\,h$ per layer plus one full layer's worth of scratch; compute rises by roughly one extra forward pass, which is about a third of a training step.

Mixed precision: shrinking the bytes

Why not just train in fp32 and be done with it? Because fp32 is slow and fat. Tensor cores run bf16 matmuls several times faster than fp32, and every byte you do not store is a byte you do not have to move. So the question is: how low can the precision go before training breaks?

fp32, fp16, bf16: what the bits buy you

A floating-point number is a sign, an exponent, and a mantissa. The exponent sets the range (how big and how small you can go); the mantissa sets the precision (how many significant digits). The three formats split 32 or 16 bits differently:

Here is the catch with fp16. Its largest value is 65,504 and its smallest normal value is about $6\times 10^{-5}$. Gradients in a deep network routinely fall below that, and they silently become zero. Activations in attention occasionally exceed it, and they become infinity. Either way the run dies.

Loss scaling, and why bf16 won

Micikevicius et al. (2017) fixed the underflow problem for fp16 with a trick called loss scaling: multiply the loss by a large constant $S$ (say 1024) before backward, so every gradient is $S$ times bigger and stays in range, then divide by $S$ before the optimizer step. Dynamic loss scaling grows $S$ until an overflow appears, then backs off.

It works, but it is a fiddly extra moving part. bf16 sidesteps it entirely: it keeps fp32's 8 exponent bits, so its range matches fp32 and nothing under- or overflows in practice. The price is a 7-bit mantissa, so each individual number is only good to about two decimal digits. Basically bf16 trades precision for range, and it turns out that training tolerates low precision far better than it tolerates a wrong range.

That is why every large run since roughly 2021 uses bf16 for the working copies and keeps an fp32 master weight for the update. The tiny per-step update $-\eta\,\hat m/\sqrt{\hat v}$ would often vanish if it were added to a bf16 weight (2 digits of precision cannot represent "1.0000 plus 0.00001"), so the accumulation happens in fp32 and the result is rounded back to bf16 for the next forward pass.

fp8 in a paragraph

The next step down is 8 bits. Micikevicius et al. (2022) proposed two fp8 formats: E4M3 (4 exponent bits, 3 mantissa bits, more precision) for the forward pass and E5M2 (more range) for gradients. H100 tensor cores run fp8 matmuls at about twice bf16 speed. The trouble is the tiny range, so fp8 needs per-tensor or per-block scaling factors, and most recipes keep the accumulation, the master weights and the optimizer in higher precision. DeepSeek-V3 (2024) reported training a 671B-parameter MoE with fp8 matmuls and fine-grained block scaling, which made fp8 pretraining credible at scale; it is still less routine than bf16.

Data parallelism: the same model, different data

Now to the splitting. The simplest idea: put a full copy of the model on each of $G$ GPUs, give each a different slice of the batch, and average the gradients. Every GPU then applies the same averaged update and stays in sync. This is data parallelism (DP), and it is how every training run scales first.

The one communication step is the gradient average. Each GPU has its own gradient vector $g_i$, and every GPU needs $\bar g = \frac{1}{G}\sum_i g_i$. The collective that does "everyone ends with the sum of everyone's vector" is called all-reduce.

Ring all-reduce

The naive way, sending every gradient to one GPU that sums and broadcasts back, makes one link carry $G$ times the data. The trick that fixes it is the ring. Arrange the $G$ GPUs in a ring and chop each gradient vector into $G$ chunks. Then run two phases:

Gradient accumulation

Data parallelism has a cousin that costs nothing: gradient accumulation. If the batch you want does not fit in memory, split it into micro-batches, run forward and backward on each, and let the gradients add up in the .grad buffers before stepping the optimizer once. Mathematically it is identical to one big batch (the loss is a mean, so scale each micro-batch loss by $1/\text{accum steps}$). The all-reduce only needs to run once, at the end, which is exactly what code/lumen/train.py does.

ZeRO and FSDP: stop replicating the state

Plain data parallelism has a glaring waste. Every GPU holds a full copy of all 16 bytes per parameter, even though the optimizer step for parameter $j$ only needs to happen once. With 64 GPUs, you store the optimizer state 64 times. ZeRO ("Zero Redundancy Optimizer", Rajbhandari et al. 2020) asks: what if each GPU owned only $1/G$ of the state?

Stages 1, 2 and 3: what gets sharded

ZeRO comes in three stages, each sharding one more thing:

Stage 1 is nearly free: the ring all-reduce already has a reduce-scatter phase in which GPU $k$ ends up with the summed chunk $k$. Just let GPU $k$ own the optimizer state for that chunk, step it there, and all-gather the updated weights instead of the gradients. Same bytes, three quarters of the memory gone.

Stage 3 is the big one: nobody holds the full model. Each layer's weights are gathered from all GPUs just before that layer runs, used, and thrown away. PyTorch's implementation of this idea is FSDP, Fully Sharded Data Parallel (Zhao et al. 2023). Llama 3 reports using FSDP for its data-parallel dimension, with weights sharded but not re-sharded after the forward pass, to avoid a second gather during backward.

Worked example: memory per GPU

Tensor parallelism: split the matrix multiply

ZeRO-3 gathers every layer's full weights onto every GPU just before use, so each GPU still needs to hold at least one layer at a time, and each GPU still does all the FLOPs for its micro-batch. What if a single layer's matmul is itself too big, or you want several GPUs to share the FLOPs of one forward pass? Then you split the matrix.

Tensor parallelism (TP), as done in Megatron-LM (Shoeybi et al. 2019), cuts a linear layer's weight matrix into pieces that live on different GPUs. The trick is to pick the cut so that the pieces can be used without talking, and then to alternate two kinds of cut so that only one communication is needed per block.

Column split, then row split

Take the MLP: $Y = \text{GeLU}(XA)\,B$. Split $A$ by columns into $[A_1 \mid A_2]$ across two GPUs. Then $XA = [XA_1 \mid XA_2]$: each GPU computes its half of the hidden activation with the full $X$ and no communication, and GeLU is element-wise so it applies to each half independently.

Now split $B$ by rows into $\begin{bmatrix} B_1 \\ B_2 \end{bmatrix}$. Then $\text{GeLU}(XA_1)\,B_1 + \text{GeLU}(XA_2)\,B_2 = Y$: each GPU produces a partial $Y$ from its own half, and one all-reduce sums the partials. One communication for the whole MLP. Attention gets the same treatment by giving each GPU a subset of the heads (a column split of $W_Q, W_K, W_V$) followed by a row split of the output projection.

The cost is that the all-reduce sits inside every layer, twice per step, on the critical path. TP therefore needs very fast links, which is why it is almost always confined to the 8 GPUs of a single node connected by NVLink.

Pipeline parallelism: split the layers

TP is limited to a node. To spread one model across nodes, cut it the other way: layers 1 to 20 on GPU 0, 21 to 40 on GPU 1, and so on. Each GPU (a "stage") runs its layers and hands the activations to the next. Only a small activation tensor crosses the link, once per stage per micro-batch, which slow inter-node links can afford. This is pipeline parallelism (PP).

The bubble

The obvious problem: while stage 0 processes the batch, stages 1 to $p-1$ sit idle, and while stage $p-1$ works, the others wait. With one batch and $p$ stages, only one GPU is ever busy. GPipe (Huang et al. 2018) fixes most of this by splitting the batch into $m$ micro-batches and streaming them through: stage 0 starts micro-batch 2 while stage 1 works on micro-batch 1.

Some idleness remains at the start (filling the pipe) and the end (draining it). Count time in units of one micro-batch forward-or-backward on one stage. The forwards take $m + p - 1$ units, of which $p - 1$ are bubble; backward is the mirror image. So the fraction of wasted time is:

The lesson: use many more micro-batches than stages. With $p = 8$ and $m = 8$, 47% of time is bubble; with $m = 64$, it is 10%. But more micro-batches means each is smaller, which hurts matmul efficiency, and in GPipe's schedule all $m$ micro-batches' activations must be held until backward starts.

1F1B: one forward, one backward

The 1F1B schedule (PipeDream, Narayanan et al. 2019; used in Megatron-LM, Narayanan et al. 2021) does not shrink the bubble, but it fixes the memory. After a short warm-up, each stage alternates: one forward on the next micro-batch, then one backward on the oldest in-flight micro-batch. A stage never holds more than $p$ micro-batches of activations, however large $m$ is. Interleaved schedules (each GPU holds several non-contiguous chunks of layers) shrink the bubble further at the cost of more communication.

Sequence and context parallelism: split the tokens

There is one more thing to split: the sequence itself. Activation memory grows linearly with $s$, and if you want to train at 128k context, no amount of weight sharding helps because the per-layer activations of a single sequence exceed a GPU on their own.

Sequence parallelism in the Megatron sense (Korthikanti et al. 2022) is a companion to tensor parallelism. The parts of a block that TP does not shard (LayerNorm, dropout, the residual stream) are instead sharded along the sequence axis across the same TP group, and the all-reduce between them is replaced with an all-gather plus a reduce-scatter, which move the same bytes. Result: the un-sharded activations go away.

Context parallelism (CP) goes further: each GPU holds a contiguous chunk of the sequence for the whole forward pass, and attention, the one operation where every token needs every other token, is handled by passing key/value blocks around a ring (Ring Attention, Liu et al. 2023). Llama 3 reports CP as its tool for the 128k-context stage, with an all-gather of K and V rather than a ring, because with GQA the K and V tensors are small relative to Q, and because the gather was easier to combine with document masking.

Expert parallelism, in a paragraph

Mixture-of-experts models add a fourth axis. Each MoE layer contains many expert MLPs, and each token is routed to only a few. Expert parallelism (EP) puts different experts on different GPUs and uses an all-to-all collective to ship each token to its experts and back. The communication is irregular (it depends on which tokens go where), which is exactly why MoE training is harder to make efficient. The MoE chapter covers routing, load balancing, and the systems trade-offs.

3-D and 4-D parallelism: how Llama 3 combined them

None of these is enough alone. TP is fast but confined to a node. PP crosses nodes cheaply but bubbles. DP scales best but replicates. So real runs stack them. "3-D parallelism" (Narayanan et al. 2021) means TP inside the node, PP across a few nodes, and DP across the rest. The product of the three degrees equals the GPU count.

Grattafiori et al. (2024) describe Llama 3's 4-D parallelism, tensor, context, pipeline, and data, arranged in that order from innermost (most communication, fastest links) to outermost. For the 405B model on 16,384 H100s at 8k context, they report TP 8, CP 1, PP 16, DP 128 (8 × 1 × 16 × 128 = 16,384), with CP rising to 16 and DP dropping to 8 for the 128k-context stage. The reported BF16 model FLOPs utilization was 38 to 43% across these configurations. The Llama 3 case study goes through the whole recipe.

Communication: the bandwidth hierarchy

Every parallelism choice above came down to "how fast is the link?" So let us put numbers on the links.

NVLink versus InfiniBand

Inside an H100 node, the 8 GPUs are connected by NVLink through NVSwitch at about 900 GB/s per GPU (bidirectional, reported peak). Between nodes, the standard is InfiniBand or RoCE Ethernet at 400 Gb/s per port, which is 50 GB/s. So the in-node link is roughly 18 times faster than the out-of-node link, and inter-node latency is also higher.

Compute-communication overlap

Bandwidth is not the only lever. The other is overlap: start sending as soon as you can, while the GPU keeps computing. Data-parallel gradient all-reduce is the classic case. Backprop produces gradients layer by layer from the top, so the all-reduce for layer 40's gradients can run while the GPU is still computing layer 39's. Frameworks bucket gradients into 25 MB chunks and launch each bucket's all-reduce the moment it is full. If the backward pass is longer than the all-reduce, the communication is hidden completely.

FSDP does the same with its weight all-gathers (prefetch layer $\ell+1$'s weights while computing layer $\ell$), and pipeline schedules overlap activation sends with the next micro-batch. A well-tuned run has the network busy nearly all the time and the GPUs never waiting for it.

FlashAttention: never write the score matrix

So far we have moved bytes between GPUs. Now look inside one GPU, because the same "bandwidth is the bottleneck" story plays out there, and attention is the worst offender.

The memory hierarchy inside a GPU

A GPU has two kinds of memory that matter. The big one is HBM (high-bandwidth memory), 40 to 80 GB, at about 1.5 to 3 TB/s. The small one is SRAM, on-chip, about 20 MB in total across the streaming multiprocessors, at around 19 TB/s. Matmuls are compute-bound: tensor cores chew through data faster than it is fetched only when the arithmetic per byte is high. Softmax, masking, dropout, and other element-wise ops are memory-bound: they do one operation per byte and spend all their time waiting on HBM.

Standard attention writes the $L\times L$ score matrix $S = QK^\top$ to HBM, reads it back to compute softmax, writes $P$ to HBM, reads it back to multiply by $V$. For $L = 8192$ and 32 heads, that is 32 × 8192 × 8192 × 2 bytes = 4.3 GB of scores per layer, round-tripped several times through the slow memory. The FLOPs are fine; the traffic is the problem.

Online softmax: the enabling trick

FlashAttention's plan is to process $K$ and $V$ in blocks small enough to fit in SRAM, and never materialize the full row of scores. But softmax needs the max and the sum over the whole row. How can you compute a softmax when you only see a block at a time?

The answer is online softmax (Milakov and Gimelshein 2018). Keep a running max $m$ and a running sum $\ell$ of exponentials relative to that max. When a new block arrives with a bigger max, rescale what you already have by $e^{m_{\text{old}} - m_{\text{new}}}$ and carry on. Concretely, for a new block of scores $s_{\text{blk}}$:

The same rescaling factor applies to the running output $o$ (the partial $P V$ product): $o_{\text{new}} = o_{\text{old}}\,e^{\,m_{\text{old}} - m_{\text{new}}} + \sum_j e^{\,s_j - m_{\text{new}}}\, v_j$. At the end, divide $o$ by $\ell$ and you have exactly the softmax-weighted sum, computed one block at a time.

Tiling: the whole algorithm

With online softmax in hand, the algorithm is simple to state. Load a block of $Q$ rows into SRAM. Loop over blocks of $K$ and $V$: load them into SRAM, compute the block of scores, update the running max, sum, and output using the rescaling rule, and move on. Write only the final normalized output $O$ to HBM (plus the per-row $m$ and $\ell$ for backward). The score matrix never exists in HBM.

Two consequences. Memory for attention drops from $O(L^2)$ to $O(L)$, which is what removed the $5as/h$ term from the activation formula earlier. And wall-clock speed goes up two to four times on typical shapes, not because there are fewer FLOPs (there are slightly more, because backward recomputes the scores) but because HBM traffic drops by an order of magnitude. FlashAttention-2 (Dao 2023) reorganised the work across GPU threads for another roughly 2× on the same idea.

Kernel fusion and torch.compile

FlashAttention is one instance of a general idea called kernel fusion. Each separate GPU operation ("kernel") reads its inputs from HBM and writes its outputs back. A chain like x = x * scale; x = x + bias; x = gelu(x) is three kernels and six trips through HBM for data that could have stayed in registers. Fusing them into one kernel makes the memory-bound chain nearly free. torch.compile (PyTorch 2.0, 2023) does this automatically: it traces your model into a graph, fuses element-wise chains, and generates Triton kernels. Typical training speedups are in the 1.2 to 2× range depending on how much of the model was memory-bound to begin with. The companion code does not require it, but it is a one-line addition worth trying on any model you train.

Checkpointing and fault tolerance

Now zoom back out to the cluster, and to a different kind of problem. With 16,000 GPUs running for months, something is always broken. A GPU fails, a network switch drops, a host reboots, a power fluctuation trips a rack. If a single failure killed the run, no large run would ever finish.

The reality of failures at scale

Grattafiori et al. (2024) report the numbers for Llama 3 405B, and they are worth reading slowly. During a 54-day snapshot of pretraining they observed 466 job interruptions, of which 419 were unexpected. About 78% of the unexpected ones were attributed to confirmed or suspected hardware issues; faulty GPUs alone were the largest single category at 58.7% of unexpected interruptions. That is roughly eight unexpected interruptions per day. Despite that, they report over 90% effective training time, and that only three interruptions required significant manual intervention.

They also report subtler effects: silent data corruption (a GPU computing a wrong answer with no error), and a 1 to 2% diurnal swing in throughput as midday temperatures changed GPU clock behaviour. At scale, the hardware is not a fixed thing; it is weather.

What makes a run survivable

Three ingredients. First, frequent model checkpoints: write the full sharded state (weights, optimizer, data-loader position, RNG state) to a parallel filesystem often enough that a crash costs minutes of work, not hours. Llama 3 reports its storage fabric was sized so a checkpoint of the 405B model (trillions of bytes of state) could be written without stalling the GPUs. Second, automatic detection and restart: a failed job is detected by health checks, the bad host is cordoned off, and the job resumes from the last checkpoint on healthy hardware, all without a human. Third, determinism: the data order and RNG streams are reproducible from the checkpoint, so a restart continues the same run rather than a subtly different one.

Throughput metrics: tokens per second and MFU

Finally, measurement. You have a run going. Is it fast? "Fast" needs a yardstick, and there are two.

Tokens per second

The raw number is tokens processed per second across the whole cluster. It tells you when the run will finish: $D / (\text{tokens/s})$ seconds. It does not tell you whether you are using the hardware well, because a bigger model processes fewer tokens per second even at perfect efficiency.

Model FLOPs Utilization

MFU (Chowdhery et al. 2022, the PaLM paper) fixes that by comparing achieved useful FLOPs to the hardware peak. The useful FLOPs per token are the $6N$ from the scaling-laws chapter (2 for forward, 4 for backward, per parameter). So:

MFU counts only the FLOPs a perfect implementation would need. Recomputation from activation checkpointing does not count, even though the GPU really does it; a related metric, hardware FLOPs utilization (HFU), includes it. The gap between the two tells you how much you are paying for memory savings.

Reported MFU for well-tuned large runs is typically 35 to 45% on H100s; Llama 3 reported 38 to 43%. Above 50% is rare because attention, norms, and communication are not matmuls and cannot use the tensor cores at full rate. The important use of MFU is as a diagnostic: a sudden drop means something went wrong, and a number far below 30% means the run is leaving money on the floor.

In code: mixed precision and gradient accumulation

Everything in this chapter beyond one GPU needs a cluster, but the two ideas you will use every single day, autocast and gradient accumulation, fit in twenty lines. This is the core of the loop in code/lumen/train.py.

import torch

scaler = torch.cuda.amp.GradScaler(enabled=(dtype == torch.float16))  # only fp16 needs loss scaling
accum = 8                                # micro-batches per optimizer step

for step, batches in enumerate(loader):  # loader yields lists of `accum` micro-batches
    for micro in batches:
        x, y = micro
        with torch.autocast(device_type="cuda", dtype=dtype):   # bf16 or fp16 matmuls
            logits = model(x)
            loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))
        scaler.scale(loss / accum).backward()   # gradients add up across micro-batches
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)                  # fp32 master weights live inside the optimizer
    scaler.update()
    optimizer.zero_grad(set_to_none=True)

Three things to notice. The loss is divided by accum so the accumulated gradient is the mean over the full batch, not the sum. The clip happens after unscaling and after all micro-batches, on the full accumulated gradient. And with bf16 the GradScaler is disabled, because bf16 has fp32's range and does not need it; the code keeps the object so the same loop works for fp16 on older GPUs.

Practice

Further reading

• Rajbhandari et al. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. The 16-bytes-per-parameter accounting and the three sharding stages.
• Shoeybi et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. Column/row tensor parallelism.
• Narayanan et al. (2021). Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. 3-D parallelism, 1F1B and interleaved pipeline schedules, the bubble analysis.
• Huang et al. (2018). GPipe. Micro-batch pipelining.
• Korthikanti et al. (2022). Reducing Activation Recomputation in Large Transformer Models. The activation-memory formula and sequence parallelism.
• Zhao et al. (2023). PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel.
• Liu et al. (2023). Ring Attention with Blockwise Transformers for Near-Infinite Context. Context parallelism.
• Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness; Dao (2023). FlashAttention-2.
• Milakov and Gimelshein (2018). Online normalizer calculation for softmax. The rescaling trick.
• Micikevicius et al. (2017). Mixed Precision Training; Micikevicius et al. (2022). FP8 Formats for Deep Learning.
• Chen et al. (2016). Training Deep Nets with Sublinear Memory Cost. Activation checkpointing.
• Chowdhery et al. (2022). PaLM. Defines MFU (Appendix B).
• Grattafiori et al. (2024). The Llama 3 Herd of Models. Section 3.3 on infrastructure, 4-D parallelism, MFU, and the failure statistics quoted here.