Training at scale: memory, parallelism, rooflines and infrastructure design

Before this lesson: lesson 07 (the one-GPU training step, 16 bytes per parameter, gradient accumulation), lesson 08 (the transformer layer that tensor parallelism splits) and lesson 11 ($C \approx 6ND$). The full theory is in the ML Systems track: training memory, collectives, data parallelism and ZeRO, tensor and pipeline parallelism and the roofline model. Time: about 105 minutes, plus the notebook.

---

Why this matters #

"Will a 13B model train on 8 GPUs with 80 GB?" and "How would you train this model on 1,000 GPUs?" are standard questions in research engineer and ML engineer interviews at LLM labs, and in ML-for-science groups that train foundation models. Some rounds ask for code: a data-parallel step with all_reduce, a simulated ring all-reduce, a column-parallel linear layer, or a pipeline schedule. Design rounds ask for a full training plan: GPU-hours, a parallel layout, a check that the network and the data loader keep up, what happens when a node dies, and how you know the model is getting better. The answers are a few formulas and a few hardware numbers. This lesson gives you both, with proofs, so that you can derive them on a whiteboard in minutes.

1. Predict first #

2. The picture #

One GPU first. A training step keeps four kinds of tensors in GPU memory:

• parameters $W$, used by the forward and the backward pass;
• gradients $\nabla W$, made by the backward pass and used by the optimizer;
• optimizer states: for AdamW the first and second moments $m$ and $v$, and in mixed precision also an fp32 master copy of the weights;
• activations: the values made by the forward pass and kept until the backward pass has used them.

The first three are the model states. Their size depends only on the number of parameters $\Psi$. The activations depend on the number of tokens processed at once, on the width and on the depth.

Interactive figure: loading…

Parallelism is a choice of what to split.

Method	What is split	What is sent	How often
data parallelism (DDP)	the batch	all-reduce of the gradients	once per step, during the backward pass
ZeRO-1, ZeRO-2	the batch + optimizer states (+ gradients)	reduce-scatter of gradients, all-gather of parameters	once per step
ZeRO-3 / FSDP	the batch + all model states	all-gather of parameters (forward and backward), reduce-scatter of gradients	layer by layer
tensor parallelism (TP)	every weight matrix	all-reduce of activations	2 per layer forward, 2 backward
pipeline parallelism (PP)	the layers	activations and their gradients, to the next stage	per micro-batch, at each stage boundary

Three collectives do almost all the communication. A collective is an operation in which all GPUs of a group take part. Take $p$ GPUs, each with a vector of $n$ bytes:

• all-reduce: every GPU ends with the sum of all $p$ vectors;
• reduce-scatter: GPU $i$ ends with chunk $i$ (of size $n/p$) of the sum;
• all-gather: every GPU starts with one chunk and ends with all $p$ chunks.

An all-reduce is a reduce-scatter followed by an all-gather. The ring algorithm arranges the GPUs in a ring. In each step every GPU sends one chunk of $n/p$ bytes to its neighbour. Each phase takes $p - 1$ steps, so an all-reduce sends

almost independent of the number of GPUs. This is why data parallelism scales: more GPUs do not mean more bytes per GPU.

Interactive figure: loading…

3. Definitions and main results #

Notation (as in the Megatron and ZeRO papers). $\Psi$ number of parameters; $L$ layers; $h$ hidden size; $a$ attention heads; $s$ sequence length; $b$ micro-batch size (sequences per GPU per forward pass); $N_d$ data-parallel GPUs; $t$ tensor-parallel degree; $p$ pipeline stages; $m$ micro-batches per step; $N$ GPUs in total; $D$ training tokens; $F$ FLOP/s per GPU; $\text{BW}$ bandwidth in bytes/s. Hardware (dense, bf16): A100-80GB SXM 312 TFLOP/s and 2.04 TB/s HBM; H100 SXM 989 TFLOP/s and 3.35 TB/s HBM, NVLink 450 GB/s per direction inside a node, about 50 GB/s per GPU between nodes (one 400 Gb/s network card per GPU). 1 GB $= 10^9$ bytes.

3.1 The memory of one training step #

Model states. Mixed-precision AdamW keeps, per parameter (ZeRO paper, §3.1):

bf16 weight	bf16 gradient	fp32 master weight	fp32 Adam $m$	fp32 Adam $v$	total
2	2	4	4	4	16 bytes

The last three are the optimizer states: $K = 12$ bytes per parameter. PyTorch autocast keeps fp32 weights instead (4 for the weight, 4 for the gradient, 8 for Adam): also 16 bytes, split differently. A 7B model needs $16 \cdot 7 \cdot 10^9 = 112$ GB before any activation.

Parameter count. A GPT-style layer (MLP width $4h$) has $12h^2 + 13h$ parameters: attention $4h^2 + 4h$, MLP $8h^2 + 5h$, two LayerNorms $4h$. So $\Psi \approx 12Lh^2$, plus the embedding.

Activations. For $y = xW$ the weight gradient is $\nabla W = x^\top \nabla y$, so the backward pass needs the input $x$. Every matmul input, every nonlinearity input and every dropout mask is stored. Korthikanti et al. (2022, §4.1) count, per layer, with 16-bit activations and 1-byte dropout masks:

• $34\,sbh$ covers the matmul inputs, the GeLU input, the LayerNorm inputs and the dropout masks.
• $5as^2b$ is the attention core: the $s \times s$ softmax output, its dropout mask and the dropout output, for each of the $a$ heads. It grows with $s^2$.
• FlashAttention never stores the $s \times s$ matrices: about $34\,sbh$ per layer.
• Full activation checkpointing stores only the input of each layer ($2sbh$ bytes) and recomputes the rest in the backward pass. This costs one extra forward pass: about $8\Psi$ instead of $6\Psi$ FLOPs per token, 33% more compute.
• Tensor parallelism with sequence parallelism of degree $t$ divides all of it by $t$.

Example: $s = 4096$, $b = 1$, $h = 4096$, $a = 32$ (a 7B model). $sbh \approx 1.68 \cdot 10^7$ and $5as/h = 160$. Without FlashAttention: $194\,sbh \approx 3.25$ GB per layer, 104 GB for 32 layers. With FlashAttention: $34\,sbh \approx 0.57$ GB per layer, 18.3 GB in total. With full checkpointing: $2sbh \cdot 32 \approx 1.07$ GB, plus one layer during recomputation.

3.2 Data parallelism #

Theorem (data parallelism = one large batch). Let the loss be the mean over a batch of $B$ samples, where each sample's loss depends only on that sample and $W$ (no BatchNorm statistics across the batch). Split the batch into $N_d$ shards of equal size. Then the mean of the $N_d$ shard gradients equals the full-batch gradient.

Proof.

1. $\mathcal L(W) = \frac{1}{B}\sum_{i=1}^{B} \ell_i(W)$, where $\ell_i$ is the loss of sample $i$.
2. Let $S_k$ be shard $k$ and $\mathcal L_k = \frac{N_d}{B}\sum_{i \in S_k} \ell_i$ its mean loss. Each sample is in exactly one shard, so $\mathcal L = \frac{1}{N_d}\sum_{k} \mathcal L_k$.
3. The gradient is linear: $\nabla \mathcal L = \frac{1}{N_d}\sum_{k} \nabla \mathcal L_k$.
4. GPU $k$ computes $\nabla \mathcal L_k$. An all-reduce (sum), then a division by $N_d$, gives $\nabla \mathcal L$ on every GPU.
5. All GPUs start from the same $W$ and apply the same update to the same gradient, so the copies stay identical. $\blacksquare$

With unequal shards (a ragged last batch, or a different number of tokens per GPU), weight each $\nabla \mathcal L_k$ by its count; section 4(a) shows the error of the plain mean. Gradient accumulation is the same proof in time instead of space: run $m$ micro-batches one after another, add up their gradients, then take one optimizer step. The global batch is $b \cdot m \cdot N_d$ sequences.

How DDP hides the all-reduce. PyTorch DDP groups the gradients into buckets (25 MiB by default). When a bucket is complete, its all-reduce starts while the backward pass continues with earlier layers. Inside no_sync() (the accumulation steps) DDP skips the all-reduce.

When is it hidden? With $T$ tokens per GPU per step, the backward pass costs about $4\Psi T$ FLOPs (4 of the $6\Psi$ FLOPs per token), and the bf16 ring all-reduce sends about $2 \cdot 2\Psi = 4\Psi$ bytes per GPU. Communication hides behind the backward pass when

The model size cancels. With $F = 400$ TFLOP/s achieved and $\text{BW} = 50$ GB/s between nodes, each GPU needs at least 8,000 tokens per step. With gradient accumulation, only the last micro-batch's backward pass overlaps with the all-reduce, so count the tokens of that micro-batch.

3.3 ZeRO and FSDP #

ZeRO (Zero Redundancy Optimizer) is data parallelism without copies of the model states: each GPU keeps a $1/N_d$ slice and collects the rest only when it needs it.

Theorem (ZeRO memory and communication). With $\Psi$ parameters, $N_d$ data-parallel GPUs and $K = 12$, the model states per GPU and the values each GPU sends per step are (times $(N_d-1)/N_d$ for ring collectives; 2 bytes per value in bf16):

Stage	What is sharded	Model states per GPU (bytes)	Sent per GPU per step
DDP (stage 0)	nothing	$(2 + 2 + K)\Psi = 16\Psi$	$2\Psi$
ZeRO-1	optimizer states	$2\Psi + 2\Psi + K\Psi/N_d$	$2\Psi$
ZeRO-2	+ gradients	$2\Psi + (2 + K)\Psi/N_d$	$2\Psi$
ZeRO-3 / FSDP	+ parameters	$(2 + 2 + K)\Psi/N_d$	$3\Psi$

Proof.

1. Memory. Each stage keeps the unsharded parts in full and the sharded parts as a $1/N_d$ slice. Add the bytes per parameter: 2 (weights), 2 (gradients), $K$ (optimizer).
2. DDP. A ring all-reduce of $\Psi$ gradient values is a reduce-scatter ($\Psi$) plus an all-gather ($\Psi$): $2\Psi$.
3. ZeRO-1 and ZeRO-2. GPU $k$ updates only its $1/N_d$ of the parameters, so it needs only its $1/N_d$ of the summed gradient: a reduce-scatter ($\Psi$). After the update, every GPU needs all the new parameters: an all-gather ($\Psi$). This is the DDP all-reduce cut in two, with the optimizer step in the middle: $2\Psi$, no extra traffic.
4. ZeRO-3. Parameters are not kept between uses. Before each layer's forward pass, all-gather its weights, then free them: $\Psi$ over the whole model. The backward pass needs them again: a second all-gather ($\Psi$). Then reduce-scatter the gradients ($\Psi$). No all-gather after the update, because the next forward pass does it. Total $3\Psi$, which is $1.5\times$ DDP. $\blacksquare$

FSDP is PyTorch's ZeRO-3. It wraps the model in units (usually one per transformer block), prefetches the all-gather of the next unit while the current unit computes, and frees the gathered weights after use. HSDP (hybrid sharding) shards inside a node and replicates across nodes, so that the all-gathers stay on NVLink.

What ZeRO does not shard: activations. Each GPU still holds the activations of its own micro-batch. For those you need checkpointing, a smaller micro-batch, or TP with sequence parallelism.

3.4 Tensor parallelism (Megatron) #

Simplest case: one matmul $Y = XA$ on two GPUs.

• Split $A$ by columns, $A = [A_1, A_2]$. Then $Y = [XA_1, XA_2]$: each GPU computes its own columns. No communication, if $X$ is on both GPUs.
• Split $A$ by rows, and split $X = [X_1, X_2]$ to match. Then $Y = X_1A_1 + X_2A_2$: each GPU has a partial sum, and an all-reduce adds them.

The MLP block $Z = \operatorname{GeLU}(XA)\,B$, with $A$ of size $h \times 4h$ and $B$ of size $4h \times h$ (Shoeybi et al. 2019, §3):

1. Split $A$ by columns. GPU $i$ computes $Y_i = \operatorname{GeLU}(XA_i)$. GeLU acts on each element, so no communication is needed.
2. Split $B$ by rows. GPU $i$ computes $Z_i = Y_iB_i$, a partial sum. One all-reduce gives $Z = \sum_i Z_i$.
3. Why not split $A$ by rows? Then GPU $i$ would hold $X_iA_i$, and $\operatorname{GeLU}(X_1A_1 + X_2A_2) \ne \operatorname{GeLU}(X_1A_1) + \operatorname{GeLU}(X_2A_2)$: a second all-reduce before the GeLU.
4. In the backward pass, every GPU used the same $X$, so $\nabla X$ is a sum over GPUs: one more all-reduce.

Attention is split by heads: the Q, K, V projections are column-parallel (each GPU gets $a/t$ heads), attention runs locally, and the output projection is row-parallel. So each layer needs 2 all-reduces in the forward pass and 2 in the backward pass, each on $b \cdot s \cdot h$ values.

Cost. TP divides weights, gradients, optimizer states and (with sequence parallelism) the activations by $t$. But its communication sits on the critical path of every layer. In the forward pass a GPU computes $24bsh^2/t$ FLOPs per layer and sends $2 \cdot 2\frac{t-1}{t} \cdot 2bsh$ bytes, so

For $h = 8192$, $t = 8$ and $F = 989$ TFLOP/s this is about 0.63 on NVLink (450 GB/s), but about 5.6 between nodes (50 GB/s). So TP stays inside a node.

3.5 Pipeline parallelism #

Split the $L$ layers into $p$ stages, one per GPU. Between stages only point-to-point messages travel: the activations of a micro-batch forward, their gradients backward ($b \cdot s \cdot h$ values each). This is little data, so PP works well across nodes. The naive pipeline sends the whole batch from stage to stage, so only one GPU works at a time. GPipe cuts the batch into $m$ micro-batches, so stage 2 works on micro-batch 1 while stage 1 works on micro-batch 2.

Theorem (pipeline bubble). With $p$ stages, $m$ micro-batches, and times $t_f$, $t_b$ for the forward and backward pass of one micro-batch on one stage (equal on all stages, communication ignored), a GPipe step takes $(m + p - 1)(t_f + t_b)$. The idle ("bubble") fraction of the step is

Proof.

1. Micro-batch $j$ ($j = 0, \dots, m-1$) can run on stage $k$ ($k = 0, \dots, p-1$) only after it left stage $k-1$, and after micro-batch $j-1$ left stage $k$. By induction on $j + k$, its forward pass on stage $k$ ends at time $(j + k + 1)\,t_f$.
2. So the last forward pass ($j = m-1$, $k = p-1$) ends at $(m + p - 1)\,t_f$.
3. The backward phase is the mirror image and takes $(m + p - 1)\,t_b$.
4. Each stage does useful work for $m(t_f + t_b)$, so it is idle for $(p-1)(t_f + t_b)$: a fraction $(p-1)/(m+p-1)$ of the step. $\blacksquare$

Relative to the ideal time $m(t_f + t_b)$, the bubble is $(p-1)/m$; the Ultra-Scale Playbook and Narayanan et al. use this form. Say which one you use.

Interactive figure: loading…

• More micro-batches shrink the bubble. With $m \ge 4p$ the bubble is below 20% of the step, because $(p-1)/(5p-1) < 1/5$.
• But GPipe stores the activations of all $m$ micro-batches, because every forward pass runs before the first backward pass.
• 1F1B (one forward, one backward) starts each backward pass as early as possible. Stage $k$ runs $p - k - 1$ warm-up forward passes, then alternates. The bubble is the same, but stage $k$ holds at most $p - k$ micro-batches of activations (Narayanan et al. 2021, §2.2).
• Interleaved 1F1B gives each GPU $v$ non-adjacent chunks of layers. The bubble drops to $(p-1)/(vm + p - 1)$, at the cost of $v$ times more point-to-point messages.

3.6 Putting them together #

Large runs combine the methods: total GPUs $=$ DP $\times$ TP $\times$ PP (and more, see below). The order of decisions (Ultra-Scale Playbook, "Finding the Best Training Configuration"):

1. Fit one step in memory. Below about 10B parameters, ZeRO with FlashAttention (and, if needed, activation checkpointing) is usually enough. Larger models add TP inside a node (at most its GPU count, 8 on an H100 node) and PP or ZeRO-3 across nodes.
2. Reach the target global batch with DP and gradient accumulation.
3. Map to the hardware. Put the method with the most traffic (TP) on the fastest links (NVLink), and PP and DP across nodes.

Context parallelism (CP) splits the $s$ tokens of one long sequence over GPUs. Ring attention passes key and value blocks around a ring of GPUs; DeepSpeed-Ulysses uses an all-to-all to switch from a split over tokens to a split over heads. CP divides the activation memory by its degree; use it when one sample's activations do not fit (ML Systems 10). Expert parallelism puts the experts of a mixture-of-experts layer on different GPUs, and tokens travel with an all-to-all (Deep Learning 15).

Example: Llama 3 405B trained on 16,384 H100s with TP 8, PP 16 and FSDP data parallelism 128 (DP is outermost, because it tolerates the longest network latency); the long-context stage added CP 16 (Llama 3 paper, §3.3.2 and Table 4).

3.7 Rooflines for compute, network and storage #

Arithmetic intensity $I$ = FLOPs / bytes moved. For any resource with rate $R$ (bytes/s), the roofline bound is

The ridge point is $I^* = F/\text{BW}_{\text{HBM}}$: an operation with $I > I^*$ is compute-bound, one with $I < I^*$ is memory-bound. On an H100, $I^* = 989/3.35 \approx 295$ FLOP/byte; on an A100, $312/2.04 \approx 153$.

Result 1 (when a matmul is compute-bound). A bf16 matmul $(B \times D)\cdot(D \times F')$ that reads its inputs and writes its output once has $I = \frac{BDF'}{BD + DF' + BF'}$, which tends to $B$ when $B \ll D, F'$.

1. FLOPs: $BF'$ outputs, each a length-$D$ dot product: $2BDF'$.
2. Bytes: read $BD$ and $DF'$ values, write $BF'$ values, 2 bytes each: $2(BD + DF' + BF')$.
3. Divide. For $B \ll D, F'$ the weight term $DF'$ dominates the denominator, so $I \approx B$. $\blacksquare$

Interactive figure: loading…

What this means for training. In training, $B$ is the number of tokens in a micro-batch, $b \cdot s$, usually thousands. So every weight matmul is compute-bound. The memory-bound parts are elementwise operations (LayerNorm, GeLU, residual adds) and an unfused attention that writes the $s \times s$ scores to HBM. The fixes are kernel fusion (torch.compile) and FlashAttention. Small matmuls are a serving problem (lesson 13).

Network and storage rooflines use the same max rule. Network: compare the bytes each GPU sends per step with its link bandwidth, and the result with the compute time (section 3.2 did this for DP). Storage: the data loader must deliver (bytes per sample) $\times$ (samples per step) / (step time). Compare it with what the shared file system sustains for your job, not with its peak. If the work on different resources overlaps, the step takes about the maximum of the times; if nothing overlaps, the sum.

3.8 Compute budget and MFU #

Result 2 (training FLOPs per token). With $\Psi \approx 12Lh^2$ weight parameters, training costs per token

1. A weight matrix of shape $m \times n$ applied to one token costs $2mn$ FLOPs. Over all weight matrices: $2\Psi$ in the forward pass.
2. The backward pass has two matmuls per forward matmul (gradient with respect to the input and to the weight), each of the same size: $4\Psi$. Total $6\Psi$.
3. The attention scores have no weights. For one query token in one layer, $QK^\top$ costs $2sh$ FLOPs (a dot product with $s$ keys over the total head width $h$), and $PV$ another $2sh$. Forward $4sh$, with backward $12sh$ per layer: $12Lsh$.
4. With a causal mask a token sees on average $s/2$ keys: $6Lsh$. $\blacksquare$

The ratio attention / weights is $12Lsh/(6 \cdot 12Lh^2) = s/(6h)$ for full attention. It is small for an LLM ($s = 4096$, $h = 4096$, causal: about 8%) but large for long sequences.

MFU (model FLOPs utilization) $=$ (model FLOPs per token $\times$ tokens per second) / (peak FLOP/s of all GPUs). It counts only the FLOPs the model needs; recomputation does not count. Llama 3 405B reached 38–43% BF16 MFU on 8k–16k H100s (Table 4). Plan with 35–45% for a dense model, and say so. Then

Goodput is the fraction of the elapsed time spent on useful training. Llama 3 reports more than 90% on 16k GPUs. MFU measures how well one step uses the GPUs; goodput measures how much of the calendar is made of useful steps.

3.9 Checkpointing and failures #

Job failure rate. A synchronous job stops when any of its $N$ GPUs (or their hosts, links or software) fails. If each GPU fails independently at rate $1/M_{\text{gpu}}$ (exponential lifetimes), the time to the first failure satisfies $P(\min_i T_i > t) = \prod_i e^{-t/M_{\text{gpu}}} = e^{-Nt/M_{\text{gpu}}}$. So the job's mean time between failures (MTBF) is

Llama 3 405B saw 419 unexpected interruptions in 54 days on 16,384 GPUs, about 78% from hardware (§3.3.4, Table 5). That is $M \approx 3.1$ hours for the job, and $M_{\text{gpu}} = 54 \times 16{,}384 / 419 \approx 2{,}111$ days $\approx 5.8$ years per GPU.

Result 3 (Young's checkpoint interval). Save a checkpoint after every $\tau$ of useful work. Each checkpoint blocks training for $\delta$; failures come at random with mean time $M$; a restart costs $R$. If $\delta \ll \tau \ll M$, the extra time per unit of useful time is

Why this is true, in one sentence: frequent checkpoints waste time writing, rare checkpoints lose more work at each failure, and the best interval balances the two.

1. Writing: each $\tau$ of work pays $\delta$, so $\delta/\tau$ per unit of work.
2. A failure hits a random point of a segment, so on average half a segment, $\tau/2$, of work is lost, plus the restart $R$.
3. Failures come at rate $1/M$, so this costs $(\tau/2 + R)/M$ per unit of time.
4. Add: overhead $\approx \delta/\tau + \tau/(2M) + R/M$ (first order in $\delta/\tau$, $\tau/M$, $R/M$).
5. Set the derivative to zero: $-\delta/\tau^2 + 1/(2M) = 0$, so $\tau^* = \sqrt{2\delta M}$.
6. At $\tau^*$ the first two terms are equal, $\sqrt{\delta/(2M)}$ each; their sum is $\sqrt{2\delta/M}$. $\blacksquare$

Interactive figure: loading…

The square root is the key point: 100 times more GPUs make $M$ 100 times shorter and the overhead 10 times larger. This is why large runs cut $\delta$ with asynchronous checkpoints: copy the state to CPU memory in seconds, then write it to storage in the background. Young (1974) gave this first-order formula; Daly (2006) gave corrections for $\tau$ close to $M$.

What a checkpoint must contain, so that a resumed run continues the same run:

• the model weights (the fp32 master copy) and the optimizer states (AdamW $m$, $v$ and its step counter): 12–16 bytes per parameter;
• the learning-rate scheduler state, the global step and the tokens seen;
• the random number generator (RNG) states of every rank (dropout and data sampling use them);
• the data position. Simplest: make the data order a pure function of (seed, epoch, step), so the step counter is enough;
• the gradient-scaler state (fp16 only), the config and the code version (git commit).

How to write it. Sharded: every rank writes its own shard in parallel (PyTorch Distributed Checkpoint, DCP), which can also reload into a different number of GPUs. Atomic: write to a new directory and mark it complete only at the end, so that a crash during the write leaves the previous checkpoint valid. Retention: keep the last 2–3, plus a permanent one every few thousand steps for analysis and for rollback after a loss spike (PaLM restarted about 100 steps before each spike and skipped 200–500 batches).

The restart path. Detect failures (collective timeouts, heartbeats, a jump in step time: one slow GPU slows every rank); restart without a human (for example Slurm --requeue); exclude bad nodes and keep spare ones; load the latest complete checkpoint. And test it: kill the job on purpose and check that the resumed loss curve matches the uninterrupted one.

3.10 The eval pipeline #

Two loops.

• Inside the training loop (cheap, every few hundred steps): training loss, loss on a small fixed validation set, gradient norm, learning rate, tokens/s and MFU, memory, and the time the loader makes the GPUs wait.
• Outside the loop (asynchronous): an eval runner watches the checkpoint directory and starts eval jobs on other GPUs. Training never waits for them. Each result is stored with its key: run, step, eval-set version, eval-code version and settings (for example decoding parameters).

Eval data. Freeze the eval sets and store their hashes. For time-dependent data (weather, finance, logs), split by time, never at random; choose checkpoints on the validation period and open the test period once, for the final report. For LLMs, check contamination (n-gram overlap between the training data and the eval sets).

Eval cost. Count it like training: about $2\Psi$ forward FLOPs per token (plus attention), times the tokens of the eval suite. Keep it to a few percent of the training compute, and check that one eval finishes before the next checkpoint arrives.

Answers to section 1. (1) DDP: $16 \cdot 7.5 \cdot 10^9 = 120$ GB per GPU, more than one 80 GB GPU. ZeRO-3: $120/64 = 1.875$ GB (ZeRO paper, Figure 1; ZeRO-1 gives 31.4 GB and ZeRO-2 16.6 GB). (2) Only $1.5\times$. Each all-gather moves every parameter once, like each half of the DDP all-reduce. Two all-gathers and one reduce-scatter make $3\Psi$, against $2\Psi$ for DDP.

4. Worked example #

(a) Data parallelism by hand #

A model $\hat y = wx$ with one weight $w = 1$, squared loss, mean over 4 samples: $x = (1, 2, 3, 4)$, $y = (2, 4, 5, 9)$. The gradient of $\mathcal L = \frac{1}{B}\sum_i (wx_i - y_i)^2$ is $\frac{2}{B}\sum_i (wx_i - y_i)\,x_i$.

• Full batch. Residuals $wx_i - y_i = (-1, -2, -2, -5)$; times $x_i$: $(-1, -4, -6, -20)$. So $\nabla \mathcal L = \frac{2}{4}(-31) = -15.5$.
• Two GPUs. GPU 0 gets samples 1–2: $\frac{2}{2}(-1 - 4) = -5$. GPU 1 gets samples 3–4: $\frac{2}{2}(-6 - 20) = -26$.
• All-reduce. The sum is $-31$; divided by $N_d = 2$ it is $-15.5$ ✓. Without the division the gradient is twice too large (for SGD, a doubled learning rate).
• Uneven shards. GPU 0 gets sample 1, GPU 1 gets samples 2–4: gradients $-2$ and $-20$. Their plain mean is $-11$: wrong. Weighted by the counts: $(1 \cdot (-2) + 3 \cdot (-20))/4 = -15.5$ ✓.

(b) The design round: pretrain a 7B model on 1T tokens #

Prompt. "Pretrain a 7B decoder ($L = 32$, $h = 4096$, $a = 32$, $s = 4096$, causal) on 1 trillion tokens within two weeks. You have H100 nodes with 8 GPUs (80 GB, NVLink), 50 GB/s per GPU between nodes, and a shared file system. Global batch about 4M tokens."

Step 1 — compute budget.

• FLOPs per token: $6\Psi + 6Lsh = 4.2 \cdot 10^{10} + 3.2 \cdot 10^{9} \approx 4.52 \cdot 10^{10}$ (attention about 7%).
• Total: $D \times 4.52 \cdot 10^{10} = 4.52 \cdot 10^{22}$ FLOPs.
• At 40% MFU, $F_{\text{eff}} = 0.4 \times 989 \approx 396$ TFLOP/s per GPU: $4.52 \cdot 10^{22} / 3.96 \cdot 10^{14} \approx 1.14 \cdot 10^{8}$ GPU-seconds $\approx$ 31,750 GPU-hours.
• Two weeks with 90% goodput: $N \ge 31{,}750 / (336 \times 0.9) \approx 105$. Take 128 GPUs (16 nodes): 10.3 days of compute, about 11.5 days of wall-clock time.

Step 2 — memory per GPU.

• Global batch: $4{,}194{,}304$ tokens $= 1{,}024$ sequences, so each of 128 GPUs processes 8 sequences per step: micro-batch $b = 1$, $m = 8$ accumulation steps.
• Model states with DDP: $16\Psi = 112$ GB. Does not fit.
• ZeRO-1 over $N_d = 128$: $4\Psi + 12\Psi/128 = 28 + 0.66 \approx 28.7$ GB.
• Activations ($b = 1$, FlashAttention): 18.3 GB. Without FlashAttention: 104 GB, impossible.
• Total about 47 GB of 80 GB, with room for the logits and the communication buffers.
• Layout: ZeRO-1 data parallelism over 128 GPUs. No TP, no PP, no activation checkpointing. Why ZeRO-1 and not ZeRO-3? ZeRO-1 communicates once per optimizer step, even with accumulation. ZeRO-3 would all-gather the weights for each of the 8 micro-batches: about 8 times the traffic.

Step 3 — roofline checks.

• HBM: the MLP matmul $(4096 \times 4096)\cdot(4096 \times 16384)$ has $I = \frac{4096 \cdot 16384}{4096 + 16384 + 16384} \approx 1{,}820$ FLOP/byte, far above 295. Compute-bound. Fuse the elementwise operations with torch.compile.
• Network: one optimizer step is 8 micro-batches of $4.52 \cdot 10^{10} \times 4096 / 3.96 \cdot 10^{14} \approx 0.47$ s, so 3.75 s. The gradient reduce-scatter + all-gather (bf16, $2\Psi = 14$ GB) sends $2 \cdot \frac{127}{128} \cdot 14 \approx 27.8$ GB per GPU. Worst case, every byte crosses the 50 GB/s link and nothing overlaps: 0.56 s, 15% of the step. A hierarchical all-reduce (reduce inside the node over NVLink, all-reduce of each GPU's $1/8$ across nodes, gather inside the node) sends only $2 \cdot \frac{15}{16} \cdot \frac{14}{8} \approx 3.3$ GB across nodes: 66 ms, under 2%. Note the overlap rule: the last micro-batch has 4,096 tokens, less than $F_{\text{eff}}/\text{BW} \approx 7{,}900$, so a flat ring would not hide fully.
• Storage: 4.2M tokens per step at 4 bytes per token ID is 17 MB per 3.75 s: about 5 MB/s. Not a risk.

Step 4 — checkpointing and failures.

• Content: fp32 weights + $m$ + $v$ $= 12\Psi = 84$ GB, plus RNG states, step, scheduler, data position and config. Sharded: 0.66 GB per GPU. Assume a blocking save of $\delta \approx 30$ s.
• $M = 2{,}111/128 \approx 16.5$ days $\approx 1.43 \cdot 10^6$ s. Young: $\tau^* = \sqrt{2 \cdot 30 \cdot 1.43 \cdot 10^6} \approx 9{,}250$ s $\approx 2.6$ h; overhead $\sqrt{2\delta/M} \approx 0.65\%$. Expect less than one failure in the whole run.
• Choose every 2,000 steps (2.1 h). Keep the last 3 and every 20,000th. Do a kill test first.
• Contrast with 16,384 GPUs: $M \approx 3.1$ h, $\tau^* \approx 14$ min and 7.3% overhead with $\delta = 30$ s; with asynchronous saves ($\delta = 5$ s), $\tau^* \approx 5.6$ min and 3%.

Step 5 — eval pipeline.

• In-loop: validation loss every 500 steps on a fixed held-out set.
• Asynchronous, at every checkpoint (about 120 over the run): a benchmark suite of about 50M tokens. Forward cost $\approx (2\Psi + 2Lsh) \times 5 \cdot 10^7 \approx 7.6 \cdot 10^{17}$ FLOPs, 0.5 GPU-hours per checkpoint, about 63 GPU-hours in total: 0.2% of training.

Step 6 — monitoring and trade-offs. Dashboards: loss, gradient norm, tokens/s, MFU per rank, loader wait, exposed communication time. Alerts: NaN or a loss spike, MFU 10% below normal (a slow node), checkpoint older than 4 hours. At 10 times the model (70B): model states 1.12 TB, so TP 8 inside the node plus PP or ZeRO-3 across nodes (check yourself, question 6). At 10 times the GPUs: failures every 1–2 days, so asynchronous checkpoints and automatic node exclusion.

(c) A different case: a long-sequence ViT #

A vision transformer for global weather fields: $L = 24$, $h = 2048$, $s = 16{,}200$ patch tokens per sample, full (non-causal) attention.

• Attention share: $s/(6h) = 16{,}200/12{,}288 \approx 1.32$, so attention is $1.32/2.32 \approx 57\%$ of the FLOPs. The rule $6\Psi D$ would underestimate the compute 2.3 times.
• Activations with FlashAttention: $34\,sbh \approx 1.13$ GB per layer, 27 GB per sample: the activations, not the model states ($16\Psi \approx 19$ GB), fill the GPU.
• Storage: one sample (input and target state, 70 channels on a $721 \times 1440$ grid, 16-bit) is about 290 MB. At 64 samples per 0.7-second step this is about 27 GB/s from the file system: here the data path, not the network, is the main risk. The same three rooflines, a different answer.

5. In code #

NumPy checks the numbers above: the data-parallel gradients of 4(a), the ZeRO table, the Megatron split, the bubble and Young's interval. The notebook has these cells and the exercises.

import numpy as np

# (a) data parallelism = one large batch: y_hat = w x, squared loss, mean over the batch
x, y, w = np.array([1., 2., 3., 4.]), np.array([2., 4., 5., 9.]), 1.0
grad = lambda xs, ys: 2 * np.mean((w * xs - ys) * xs)
full = grad(x, y)
g0, g1 = grad(x[:2], y[:2]), grad(x[2:], y[2:])            # two GPUs, two samples each
print(full, g0, g1, (g0 + g1) / 2)                          # -15.5 -5.0 -26.0 -15.5
u0, u1 = grad(x[:1], y[:1]), grad(x[1:], y[1:])            # uneven shards: 1 and 3 samples
print((u0 + u1) / 2, (1 * u0 + 3 * u1) / 4)                 # -11.0 (wrong) -15.5 (count-weighted)

# (b) model-state memory per GPU (GB) for mixed-precision AdamW, K = 12
def zero_gb(psi, n):
    return [16 * psi, 4 * psi + 12 * psi / n, 2 * psi + 14 * psi / n, 16 * psi / n]
print([round(v / 1e9, 2) for v in zero_gb(7.5e9, 64)])     # [120.0, 31.41, 16.64, 1.88]

# (c) Megatron: split the first MLP weight by columns, the second by rows
gelu = lambda z: 0.5 * z * (1 + np.tanh(np.sqrt(2 / np.pi) * (z + 0.044715 * z**3)))
rng = np.random.default_rng(0)
X, A, B = rng.normal(size=(4, 8)), rng.normal(size=(8, 32)), rng.normal(size=(32, 8))
Z = gelu(X @ A) @ B                                          # the unsplit MLP
Z1 = gelu(X @ A[:, :16]) @ B[:16]                            # GPU 0: its 16 columns of A, 16 rows of B
Z2 = gelu(X @ A[:, 16:]) @ B[16:]                            # GPU 1
print(np.allclose(Z, Z1 + Z2))                               # True: one all-reduce (a sum) at the end
P1, P2 = X[:, :4] @ A[:4], X[:, 4:] @ A[4:]                  # row split of A instead
print(np.allclose(gelu(X @ A), gelu(P1) + gelu(P2)))         # False: GeLU of a sum is not a sum of GeLUs

# (d) pipeline bubble and (e) Young's interval
p, m = 4, 8
print(round((p - 1) / (m + p - 1), 4))                       # 0.2727
M_gpu = 54 * 16384 / 419                                     # days per GPU between failures (Llama 3)
for n in (128, 16384):
    M = M_gpu * 86400 / n                                    # job MTBF in seconds
    for delta in (30, 5):
        tau = np.sqrt(2 * delta * M)
        print(n, delta, round(M / 3600, 1), round(tau / 60, 1), round(100 * np.sqrt(2 * delta / M), 2))
# 128 30 395.9 154.1 0.65      (GPUs, delta in s, M in h, tau* in min, overhead in %)
# 128 5 395.9 62.9 0.26
# 16384 30 3.1 13.6 7.34
# 16384 5 3.1 5.6 3.0

A Monte Carlo check of Young's formula, with random failures (time in hours):

import numpy as np

M, delta, R = 24.0, 0.05, 0.1         # hours: job MTBF, checkpoint cost, restart cost
work = 2000.0                         # hours of useful training


def simulate(tau, rng):
    """Elapsed hours to finish `work`, with a checkpoint after every tau hours of work."""
    t, done, next_fail = 0.0, 0.0, rng.exponential(M)
    while done < work:
        seg = min(tau, work - done) + delta            # work, then write a checkpoint
        if t + seg <= next_fail:                       # the segment finishes: progress is saved
            t, done = t + seg, done + seg - delta
        else:                                          # failure: lose the segment, pay the restart
            t = next_fail + R
            next_fail = t + rng.exponential(M)
    return t


rng = np.random.default_rng(0)
taus = np.array([0.25, 0.5, 1.0, 1.55, 2.5, 4.0, 8.0])
sim = np.array([np.mean([simulate(tau, rng) / work - 1 for _ in range(40)]) for tau in taus])
formula = delta / taus + taus / (2 * M) + R / M
print(np.round(100 * sim, 1))       # [21.2 11.8  7.8  7.2  8.2 10.9 19.8]
print(np.round(100 * formula, 1))   # [20.9 11.5  7.5  6.9  7.6 10.  17.7]
print(round(np.sqrt(2 * delta * M), 2), taus[np.argmin(sim)])   # 1.55 1.55

The minimum is at $\tau^* = 1.55$ h in both. The simulation is a little above the first-order formula, and more so for large $\tau$, where $\tau \ll M$ holds less well.

6. Where it shows up #

• LLM pretraining. Every large run is 3-D or 4-D parallel, with an MFU around 40% and failures every few hours at the largest scale (Llama 3, §3.3). The data-parallel size sets the global batch, and the global batch sets the learning rate.
• Fine-tuning. FSDP with LoRA or QLoRA: the frozen base model is sharded, and only the small adapters have gradients and Adam states (lesson 11).
• ML for science. Weather, climate and other scientific foundation models often have moderate parameter counts but very long token sequences and large samples. Attention FLOPs, activations and the input pipeline then matter more than ZeRO (section 4(c)). HPC clusters add node-hour allocations, job time limits and a shared parallel file system.
• RL post-training. An inference engine runs inside the training loop to generate rollouts, so the design mixes the training numbers of this lesson with the serving numbers of lesson 13.

7. Common mistakes #

• "2 bytes per parameter" for training. That is bf16 inference. Training with AdamW needs 16 bytes per parameter, plus the activations.
• "ZeRO-3 shards activations." ZeRO-1/2/3 shard model states only. Activations need checkpointing, a smaller micro-batch, TP with sequence parallelism, or context parallelism.
• "ZeRO-3 costs $N_d$ times the communication." It is $3\Psi$ against $2\Psi$.
• Gradient accumulation with ZeRO-2/3. ZeRO-2 reduce-scatters after every micro-batch, and ZeRO-3 also all-gathers the weights for every micro-batch: about $m$ times the traffic of ZeRO-1.
• Sum instead of mean. Without the division by $N_d$ the gradient is $N_d$ times too large. With uneven shards, weight by the counts.
• TP across nodes. Its per-layer all-reduces need NVLink-class bandwidth.
• Too few micro-batches in a pipeline. With $m < p$ most stages wait; aim for $m \ge 4p$ and use 1F1B so that the activations do not grow with $m$.
• $6\Psi D$ at long sequence lengths. It drops the attention term; for $s \gg 6h$ it can be off by a factor of 2 or more.
• Planning at peak FLOP/s. Use 35–45% MFU for a dense model, and say so.
• Weights-only checkpoints, written in place, never tested. Without optimizer state, RNG states and data position a "resume" is a different run; a crash during an in-place write destroys the only copy; an untested restart fails when you need it.

8. Check yourself #

1. A 70B model, mixed-precision AdamW, 64 GPUs with 80 GB each. Model-state memory per GPU with ZeRO-1, ZeRO-2 and ZeRO-3? Which ones fit?
   Answer

   ZeRO-1: $4\Psi + 12\Psi/64 = 280 + 13.1 \approx 293$ GB. ZeRO-2: $2\Psi + 14\Psi/64 = 140 + 15.3 \approx 155$ GB. ZeRO-3: $16\Psi/64 = 17.5$ GB. Only ZeRO-3 fits (or add TP/PP); the activations come on top.
2. Why does ZeRO-2 send no more data than DDP?
   Answer

   A ring all-reduce already is a reduce-scatter plus an all-gather. ZeRO-2 keeps the two halves and puts the optimizer step on the shard between them: reduce-scatter the gradients ($\Psi$), update the own shard, all-gather the parameters ($\Psi$). Total $2\Psi$, the same as DDP.
3. $p = 8$ stages, $m = 32$ micro-batches. What is the bubble with GPipe and with 1F1B, and with interleaving ($v = 2$)? How many micro-batches of activations does stage 0 hold?
   Answer

   GPipe and 1F1B: $7/39 \approx 17.9\%$ of the step. Interleaved: $7/71 \approx 9.9\%$. Stage 0 holds 32 micro-batches with GPipe and at most 8 with 1F1B.
4. DDP across nodes with $F = 400$ TFLOP/s achieved per GPU and $\text{BW} = 50$ GB/s per GPU. How many tokens per GPU per step hide the all-reduce? Does the answer change for a model 10 times larger?
   Answer

   $T \ge F/\text{BW} = 8{,}000$ tokens. No: both the backward FLOPs ($4\Psi T$) and the all-reduce bytes ($4\Psi$) are proportional to $\Psi$, so $\Psi$ cancels.
5. A job on 1,024 GPUs ($M_{\text{gpu}} \approx 2{,}111$ days), blocking checkpoint $\delta = 60$ s, restart $R = 10$ min. Give $M$, $\tau^*$ and the overhead.
   Answer

   $M = 2{,}111 \cdot 86{,}400/1{,}024 \approx 1.78 \cdot 10^5$ s $\approx 49.5$ h. $\tau^* = \sqrt{2 \cdot 60 \cdot 1.78 \cdot 10^5} \approx 4{,}620$ s $\approx 77$ min. Overhead $\sqrt{2\delta/M} \approx 2.6\%$ plus $R/M \approx 0.34\%$: about 2.9%.
6. Pretrain a 70B decoder ($L = 80$, $h = 8192$, $s = 8192$) on 1,024 H100s (128 nodes of 8). Why can it not use ZeRO-1 alone, and what layout would you propose?
   Answer

   ZeRO-1 keeps full bf16 weights and gradients: $4\Psi = 280$ GB per GPU. Propose TP 8 inside each node (states and activations divided by 8), PP 4 across nodes and DP 32 with ZeRO-1. Each GPU holds $\Psi' = \Psi/(t\,p) \approx 2.2 \cdot 10^9$ parameters, so its model states are $4\Psi' + 12\Psi'/32 \approx 8.75 + 0.82 \approx 9.6$ GB. The TP communication/compute ratio is about 0.63 on NVLink; the pipeline bubble with $m = 32$ is $3/35 \approx 8.6\%$.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. How many bytes per parameter does mixed-precision AdamW keep, and how are they split?

   Show answer

2. How many bytes does each GPU send in a ring all-reduce of $n$ bytes over $p$ GPUs?

   Show answer

3. Why does data parallelism with equal shards give the full-batch gradient?

   Show answer

4. When does the DDP all-reduce hide behind the backward pass?

   Show answer

5. What are the model states per GPU with ZeRO-3 on $N_d$ GPUs?

   Show answer

6. Why does ZeRO-3 send $1.5\times$ the data of DDP?

   Show answer

7. How many all-reduces does a Megatron TP layer (attention + MLP) need per training step?

   Show answer

8. What is the pipeline bubble fraction with $p$ stages and $m$ micro-batches?

   Show answer

9. What is the mean time between failures of a job on $N$ GPUs?

   Show answer

10. Give Young's checkpoint interval and its overhead.

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

Recognize the task from the wording.

The prompt says	Start with
"will it fit?", "out of memory"	the memory table: $16\Psi$, activations $34\,sbh$ per layer, then ZeRO stage, TP, checkpointing
"how would you train this on N GPUs?"	DP (with ZeRO) first; TP inside a node; PP across nodes only if needed
"the step is slow", "low GPU utilization"	rooflines per resource: HBM, network (exposed communication), data loader
"communication is the bottleneck"	$T \ge F/\text{BW}$; hierarchical collectives; ZeRO-1 instead of ZeRO-3 with accumulation
"how long will training take?", "how many GPUs?"	FLOPs per token $\times D$, MFU 35–45%, goodput 90%
"nodes fail", "a three-month run"	$M = M_{\text{gpu}}/N$, Young's $\tau^*$, asynchronous sharded checkpoints, restart path
"very long sequences", "images or fields as tokens"	attention FLOPs $12Lsh$, FlashAttention, context parallelism, the data path

The routine to say aloud (about 30 seconds): "Let me first get the numbers: parameters, layers, width, sequence length, tokens, GPUs per node, memory, link bandwidths and the deadline. Then I compute the FLOPs and GPU-hours, then the memory per GPU, then I choose the smallest layout that fits, and I check compute, network and storage with rooflines. Last, failures and evals."

A plan for a 45-minute design round.

Minutes	Step	What you write down
0–5	Clarify	model shape, data, token budget, deadline, hardware, what "done" means
5–10	Compute budget	FLOPs per token, total FLOPs, GPU-hours at an assumed MFU, number of GPUs
10–15	Memory	model states, activations, what must be sharded
15–22	Layout	DP / ZeRO / TP / PP / CP degrees and their mapping to nodes
22–27	Roofline checks	HBM, network (is communication hidden?), storage (is the loader fast enough?)
27–33	Reliability	checkpoint content and size, interval, restart path
33–39	Eval pipeline	in-loop signals, asynchronous checkpoint evals, eval sets, cost
39–45	Monitoring, trade-offs	dashboards and alerts; what changes at 10× the model, data or cluster

Typical questions and short model answers.

1. "Will a 13B model train on one node of 8 × 80 GB?" — "The model states are $16 \times 13 = 208$ GB, so DDP does not fit. ZeRO-3 over 8 GPUs gives 26 GB per GPU; ZeRO-2 gives $2 \cdot 13 + 14 \cdot 13/8 \approx 48.8$ GB. Activations with FlashAttention at $s = 4096$, $h = 5120$, 40 layers: $34\,sbh \approx 0.71$ GB per layer, 28.5 GB per sequence. So ZeRO-3 with one sequence per GPU fits at about 55 GB; with ZeRO-2 I would add activation checkpointing. NVLink makes the ZeRO-3 all-gathers cheap inside the node."
2. "Explain ZeRO stages 1, 2 and 3." — "All three are data parallel. Stage 1 shards the optimizer states, stage 2 also the gradients, stage 3 also the parameters. Stages 1 and 2 send the same $2\Psi$ as DDP, because they split the all-reduce into a reduce-scatter before the update and an all-gather after it. Stage 3 gathers the weights twice per step: $3\Psi$, 1.5 times DDP, in exchange for $16\Psi/N_d$ memory."
3. "Why does tensor parallelism stay inside a node?" — "It needs 4 all-reduces of the activations per layer, on the critical path. The ratio of communication to compute is $(t-1)F/(3h\,\text{BW})$: about 0.6 on NVLink for $h = 8192$, but about 5.6 at 50 GB/s between nodes."
4. "What is the pipeline bubble and how do you reduce it?" — "With $p$ stages and $m$ micro-batches, the first stage waits $p-1$ slots for the pipeline to fill and the last waits to drain: a fraction $(p-1)/(m+p-1)$. More micro-batches shrink it, $m \ge 4p$ gives under 20%. 1F1B keeps the activations at $p$ micro-batches instead of $m$; interleaving divides the bubble by $v$."
5. "How often should you checkpoint a 16k-GPU job?" — "The job fails about every 3 hours. Young's interval is $\sqrt{2\delta M}$: with a 30-second blocking save, about 14 minutes and 7% overhead. That is too much, so I make the save asynchronous (copy to CPU, write in the background): with 5 seconds, about 6 minutes and 3%."
6. "MFU dropped from 42% to 30% after a restart. What do you check?" — "A slow node first: compare per-rank step times, run GPU and network health checks, exclude the node. Then the data path after the restart (a cold cache), and exposed communication in a profile."

Typical follow-ups. "How does FSDP overlap communication?" (it prefetches the next unit's all-gather and runs reduce-scatters while the backward pass continues). "What happens when you restart on a different number of GPUs?" (sharded checkpoints are resharded at load time; the global batch must stay the same, so change the accumulation steps). "The loss spikes at step 80k?" (roll back to a permanent checkpoint before the spike, skip the batches around it, then lower the learning rate or tighten clipping if it repeats).

10. Practice #

The notebook for this lesson runs in Colab, in NumPy. Each exercise has tests against a reference:

1. ring_allreduce: simulate the ring on $p$ GPUs (a list of arrays); check the sum and that each GPU sends $2(p-1)/p \cdot n$ values.
2. data_parallel_grad: the count-weighted mean of shard gradients equals the full-batch gradient of a linear model, also for uneven shards.
3. zero_memory: the per-GPU model states of DDP and ZeRO-1/2/3, checked against the ZeRO paper's 7.5B example.
4. pipeline_schedule: simulate GPipe and 1F1B and return the step time, the bubble and the peak in-flight micro-batches; compare with $(p-1)/(m+p-1)$.
5. young_interval: find the best interval numerically and by Monte Carlo, and compare with $\sqrt{2\delta M}$.

Timed practice (25 minutes each): Sasha Rush's LLM-Training-Puzzles, in order: "Standard Training", "Gradient Accumulation", "Distributed Data Parallel", "Weight-Sharded Data Parallel", "Fully-Sharded Data Parallel", "Pipeline Parallel", "GPipe Schedule", "Pipeline + FSDP". Predict the time and peak memory of each puzzle before you code it. Then two design drills of 45 minutes: the worked example of section 4(b) from memory, and question 6 of section 8 as a full design.

Say it aloud drill. Without notes, explain in 2 minutes why ZeRO-1 and ZeRO-2 send as much as DDP and ZeRO-3 sends 1.5 times as much. Then derive Young's interval in 1 minute.

11. Watch and read #

EfficientML.ai Lecture 17: Distributed Training (Part I) (MIT 6.5940, Fall 2023)MIT HAN Lab · 61 min

Data parallelism and the communication primitives behind it, including the ring all-reduce; watch it next to sections 2 and 3.2.

How Fully Sharded Data Parallel (FSDP) works?Ahmed Taha · 33 min

How FSDP shards the model states, and when it gathers and frees the weights; watch it next to section 3.3.

Large Model Training and Inference with DeepSpeed // Samyam Rajbhandari // LLMs in Prod ConferenceAAIF Live · 36 min

The first author of the ZeRO paper explains ZeRO, 3D parallelism and the DeepSpeed tools built on them.

• Stanford CS336 (Spring 2025), Lecture 7: Parallelism 1 and Lecture 8: Parallelism 2 (videos on YouTube; embedding is disabled): collectives, data, tensor and pipeline parallelism, with real numbers.
• Hugging Face, The Ultra-Scale Playbook (free web book): "First Steps: Training on One GPU", "Data Parallelism" (ZeRO-1/2/3), "Tensor Parallelism", "Pipeline Parallelism", "Finding the Best Training Configuration".
• Austin et al., How to Scale Your Model (free web book): "All About Rooflines" and "How to Parallelize a Transformer for Training".
• Sasha Rush, LLM-Training-Puzzles: eight puzzles that simulate one training step on several GPUs.
• Rajbhandari et al. (2020), ZeRO: Memory Optimizations Toward Training Trillion Parameter Models: §3.1 (16 bytes per parameter), §5 (the three stages), §7 (communication volume).
• Shoeybi et al. (2019), Megatron-LM, §3 (the column/row split); Korthikanti et al. (2022), Reducing Activation Recomputation in Large Transformer Models, §4 (the activation formula).
• Huang et al. (2019), GPipe; Narayanan et al. (2021), Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, §2 (1F1B, interleaving, the bubble).
• Llama 3 team (2024), The Llama 3 Herd of Models: §3.3.2 and Table 4 (4-D parallelism, MFU), §3.3.4 and Table 5 (interruptions and their causes).
• PyTorch, Asynchronous Saving with Distributed Checkpoint (DCP).
• Young (1974), "A first order approximation to the optimum checkpoint interval", Communications of the ACM 17(9), 530–531: the original derivation of $\tau^* = \sqrt{2\delta M}$.