/**
* Segmented step control for visualization stepping.
* @param {Object} options
* @param {number} options.min - Minimum step value (default 0)
* @param {number} options.max - Maximum step value
* @param {number} options.value - Initial value (default min)
* @param {string} options.label - Optional label text
* @returns {number} Current step value (reactive)
*/
stepControl = function({min = 0, max, value, label = null} = {}) {
const initialValue = value ?? min;
const steps = Array.from({length: max - min + 1}, (_, i) => min + i);
const container = htl.html`<div class="step-control">
${label ? htl.html`<span class="step-control-label">${label}</span>` : ''}
<div class="step-control-segments" role="group" aria-label="${label || 'Step control'}">
${steps.map(step => htl.html`<button
class="step-control-segment ${step === initialValue ? 'active' : ''}"
data-step="${step}"
aria-pressed="${step === initialValue}"
tabindex="${step === initialValue ? 0 : -1}"
>${step}</button>`)}
</div>
</div>`;
const segments = container.querySelectorAll('.step-control-segment');
let currentValue = initialValue;
function updateActive(newValue) {
currentValue = newValue;
segments.forEach(seg => {
const isActive = parseInt(seg.dataset.step) === newValue;
seg.classList.toggle('active', isActive);
seg.setAttribute('aria-pressed', isActive);
seg.tabIndex = isActive ? 0 : -1;
});
container.value = newValue;
container.dispatchEvent(new Event('input', {bubbles: true}));
}
// Click handler
segments.forEach(seg => {
seg.addEventListener('click', () => {
updateActive(parseInt(seg.dataset.step));
});
});
// Keyboard navigation
container.addEventListener('keydown', (e) => {
if (e.key === 'ArrowRight' || e.key === 'ArrowDown') {
e.preventDefault();
const next = Math.min(currentValue + 1, max);
updateActive(next);
segments[next - min].focus();
} else if (e.key === 'ArrowLeft' || e.key === 'ArrowUp') {
e.preventDefault();
const prev = Math.max(currentValue - 1, min);
updateActive(prev);
segments[prev - min].focus();
} else if (e.key === 'Home') {
e.preventDefault();
updateActive(min);
segments[0].focus();
} else if (e.key === 'End') {
e.preventDefault();
updateActive(max);
segments[max - min].focus();
}
});
container.value = initialValue;
return container;
}Module 23: Distributed Training
Introduction
Every training module so far (m02, m07) quietly assumed one device holds the whole model. Frontier LLMs do not fit. A 7-billion-parameter model needs 120 GB just to hold its training state — more than any single GPU — and that is before a single activation. Distributed training is how the same math from m07 runs across a cluster.
There are two orthogonal problems, and this module builds a from-scratch answer to each:
- Speed → data parallelism. Replicate the model on \(N\) devices, feed each a different slice of the batch, then average the gradients so every replica takes the exact step one big-batch device would. The averaging is a ring all-reduce, which we build as a real, bandwidth-optimal algorithm.
- Memory → ZeRO. Plain data parallelism keeps a full copy of the parameters, gradients, and optimizer states on every device — \(N\)-way redundant. ZeRO (Zero Redundancy Optimizer) partitions them across devices, cutting per-device memory from \(16\Psi\) bytes toward \(16\Psi/N\).
Why it matters for LLMs:
- Nothing else fits. Every model past a few billion parameters is trained distributed; the memory math below is the reason the field talks in “GPU-hours”.
- The step is unchanged. Done right, \(N\) devices produce bit-for-bit the same gradient as one — distribution is a systems trick, not a new algorithm.
- Memory is a budget you can rebalance. ZeRO turns “the model doesn’t fit” into “add devices,” on a precise, linear curve you will drive yourself.
What You’ll Learn
After this module, you can:
- Compute the training memory budget of any model: the \(2+2+12 = 16\) bytes-per-parameter model of mixed-precision Adam.
- Explain data parallelism and why averaging gradients equals one big batch.
- Build a ring all-reduce from scratch (reduce-scatter → all-gather) and prove it equals the naive sum, at \(2(N{-}1)/N \cdot M\) bytes per device.
- Derive the three ZeRO stages (\(P_{os}\), \(P_{os+g}\), \(P_{os+g+p}\)) and their per-device memory, and see the \(4\times\)/\(8\times\)/\(N\times\) savings.
- Place tensor parallel, pipeline parallel, FSDP, and gradient checkpointing in the same map.
Prerequisites
This module requires familiarity with:
- Module 02: Autograd — gradients are the thing we average across devices.
- Module 07: Training — the optimizer states (Adam’s momentum and variance) that dominate the memory budget.
- Module 09: Efficient Attention — the inference memory wall (the KV cache); this module is its training twin.
Intuition: One GPU Is Not Enough
m09 hit a memory wall at inference — the KV cache grows with the context. Training hits a wall long before that, and for a different reason: to take one Adam step you must keep several full-model-sized tensors resident at once.
Think of the model’s parameters \(\Psi\) as the unit. Modern training runs in mixed precision: the forward and backward passes use fast 16-bit floats, but the optimizer keeps a high-precision 32-bit copy so tiny updates don’t vanish. Counting bytes per parameter:
- an fp16 copy of the parameters — used in the forward/backward pass → 2 bytes
- an fp16 copy of the gradients → 2 bytes
- the fp32 optimizer states: a master copy of the parameters (4), Adam’s momentum (4), and Adam’s variance (4) → 12 bytes
That is \(2 + 2 + 12 = \mathbf{16}\) bytes for every parameter, every step. A 7.5B model therefore needs \(7.5 \times 10^9 \times 16 = 120\) GB — and an 80 GB A100 cannot hold it. The parameters you actually compute with are only \(\tfrac{1}{8}\) of the bill; the optimizer states are \(\tfrac{3}{4}\) of it. That imbalance is what ZeRO exploits.
NoteKey Insight
The thing that doesn’t fit is rarely the parameters — it’s the optimizer state. Adam’s momentum and variance, plus the fp32 master copy, are \(12\) of the \(16\) bytes. Halving parameter precision barely helps; partitioning the optimizer state is where the memory is.
The Math: The Training Memory Budget
The byte-count above is exactly the memory model from the ZeRO paper. In code it is training_memory_bytes in distributed.py:
import sys
sys.path.insert(0, ".")
from distributed import training_memory_bytes, zero_memory_per_device
mem = training_memory_bytes(7_500_000_000) # 7.5B parameters
for k, v in mem.items():
print(f"{k:>10}: {v/1e9:6.1f} GB") params: 15.0 GB
grads: 15.0 GB
optimizer: 90.0 GB
total: 120.0 GBParams and grads are 2 bytes each; the optimizer states are 12 (Adam’s \(K=12\)). Now the key move — ZeRO data parallelism partitions that state across \(N\) devices in three stages, each partitioning one more kind of state:
| Stage | Name | What is partitioned | Per-device bytes |
|---|---|---|---|
| 0 | baseline | nothing (full replica) | \(16\Psi\) |
| 1 | \(P_{os}\) | optimizer states | \(4\Psi + 12\Psi/N\) |
| 2 | \(P_{os+g}\) | + gradients | \(2\Psi + 14\Psi/N\) |
| 3 | \(P_{os+g+p}\) | + parameters | \(16\Psi/N\) |
Watch a 7.5B model land on 64 GPUs:
psi, n = 7_500_000_000, 64
for stage in range(4):
gb = zero_memory_per_device(psi, n, stage) / 1e9
print(f"stage {stage}: {gb:7.2f} GB/device")stage 0: 120.00 GB/device
stage 1: 31.41 GB/device
stage 2: 16.64 GB/device
stage 3: 1.88 GB/device120 GB collapses to under 2 — the same numbers the ZeRO paper reports (a \(4\times\), \(8\times\), and \(N\times\) reduction). The model didn’t shrink; the redundancy did.
Drive the memory budget
Slide the model size, the device count, and the ZeRO stage. The bar is per-device memory; the ghost behind it is the \(16\Psi\) baseline you started from.
TipTry This
- Set params to 70B. Even stage 3 needs many devices to fit an 80 GB GPU — this is why 70B models train on hundreds of GPUs.
- Hold the model fixed and slide \(N\). Stage 1 flattens out (it can never beat \(4\Psi\)); stage 3 keeps falling as \(16\Psi/N\). That gap is the difference between partitioning some state and partitioning all of it.
- Compare stage 1 vs stage 3 at \(N=8\). Small clusters still get most of the win from partitioning just the optimizer state — the cheapest stage to run.
Code: Data Parallelism & Ring All-Reduce
Memory is one half; speed is the other. Data parallelism runs \(N\) replicas of the model, each on a different shard of the batch. Replica \(i\) computes the gradient of its own mini-batch; then all replicas average their gradients so every one applies the same update. Because each shard’s loss is a mean over an equal slice, that average is exactly the gradient of the loss over the whole batch:
\[ \frac{1}{N}\sum_{i=1}^{N} \nabla_i \;=\; \nabla \Big(\tfrac{1}{B}\textstyle\sum_{b} \ell_b\Big) \]
Let’s prove it on a linear model. data_parallel_gradients shards the batch, computes each shard’s gradient with autograd, sums them with a ring all-reduce, and divides by \(N\):
import torch
from distributed import data_parallel_gradients, single_device_gradient
torch.manual_seed(0)
W = torch.randn(3, 5) # a shared (out=3, in=5) weight
x_shards = [torch.randn(4, 5) for _ in range(4)] # 4 devices, 4 samples each
y_shards = [torch.randn(4, 3) for _ in range(4)]
g_dp = data_parallel_gradients(W, x_shards, y_shards)
g_sd = single_device_gradient(W, torch.cat(x_shards), torch.cat(y_shards))
print("max |g_dp - g_single|:", (g_dp - g_sd).abs().max().item())
print("identical step:", torch.allclose(g_dp, g_sd, atol=1e-5))max |g_dp - g_single|: 5.960464477539063e-08
identical step: TrueFour devices, four separate gradients, one all-reduce — and the result is the single-device gradient over all 16 samples. Distribution changed nothing about the math. The only new machinery is the collective that did the averaging.
The all-reduce, and why it must be a ring
An all-reduce leaves every device holding the sum of all devices’ tensors. The obvious way — send everything to device 0, add, broadcast back — pushes \(O(N \cdot M)\) bytes through one poor GPU’s link while the rest sit idle. The ring all-reduce fixes both: it uses every link at once and moves only \(2(N{-}1)/N \cdot M\) bytes per device, independent of \(N\) as it grows.
The trick is to arrange the devices in a logical ring (each sends only to its right neighbour) and split every tensor into \(N\) chunks, then run two phases:
- Reduce-scatter (\(N{-}1\) steps): chunks circulate and accumulate until each device owns the fully summed value of exactly one chunk.
- All-gather (\(N{-}1\) steps): those finished chunks circulate again until every device has all of them.
This is ring_all_reduce in distributed.py — a real implementation, snapshotting each step’s sends before applying receives so it matches true simultaneous communication. Its correctness anchor is exact:
from distributed import ring_all_reduce, ring_all_reduce_bytes
tensors = [
torch.tensor([1., 2., 3., 4.]),
torch.tensor([10., 20., 30., 40.]),
torch.tensor([100., 200., 300., 400.]),
torch.tensor([1000., 2000., 3000., 4000.]),
]
out = ring_all_reduce(tensors)
print("device 0 result:", out[0].tolist())
print("== naive sum: ", torch.equal(out[0], torch.stack(tensors).sum(0)))
print("all devices agree:", all(torch.equal(o, out[0]) for o in out))
print("bytes/device (M=4KB, N=4):", ring_all_reduce_bytes(4000, 4))device 0 result: [1111.0, 2222.0, 3333.0, 4444.0]
== naive sum: True
all devices agree: True
bytes/device (M=4KB, N=4): 6000.0Bit-for-bit equal to the sum, on every device. And the communication cost \(2(N{-}1)/N \cdot M\) approaches \(2M\) — a fixed budget no matter how many GPUs you add, which is exactly why the ring, not a central reducer, is what real systems (NCCL, Horovod) run.
Step through the ring
Four devices, each starting with one row of values; watch reduce-scatter accumulate one finished chunk per device (the ringed cell), then all-gather spread the finished chunks around. By the last step every cell equals the column sum.
TipTry This
- Step to the end of reduce-scatter (step 3). Exactly one cell per device is ringed — the finished chunk. No device has the whole answer yet.
- Finish the all-gather. The ringed chunks copy around the ring until every row equals the column sums. That second half is why the cost is \(2\times\) the scatter, not \(1\times\).
ZeRO: Stop Replicating What You Can Partition
Data parallelism is fast but memory-wasteful: all \(N\) devices hold the same 16 bytes/param. ZeRO keeps data parallelism’s compute pattern but removes the redundancy — each device stores only its slice of the state and fetches the rest with collectives (an all-gather for parameters just before they’re used, a reduce-scatter for gradients) exactly when needed.
- Stage 1 — \(P_{os}\): partition the optimizer states (the 12 fp32 bytes). Each device updates only its slice of the master weights, then all-gathers the updated fp16 params. Per-device: \(4\Psi + 12\Psi/N\).
- Stage 2 — \(P_{os+g}\): also partition the gradients — a device only needs the gradient slice for the optimizer slice it owns. Per-device: \(2\Psi + 14\Psi/N\).
- Stage 3 — \(P_{os+g+p}\): also partition the parameters themselves, all-gathering each layer’s weights just-in-time for its forward/backward, then discarding them. Per-device: \(16\Psi/N\) — a true \(N\times\) split. (This is what PyTorch ships as FSDP.)
The step through the stages, with the \(N=64\) / 7.5B numbers as the caption:
NoteKey Insight
ZeRO is not a different parallelism from data parallelism — it is data parallelism, with the redundant storage removed and paid back with communication. Stage 3 partitions everything and is what PyTorch’s FSDP implements; you trade a modest amount of extra all-gather traffic for an \(N\times\) memory cut.
Beyond Data Parallelism
Data parallelism (and ZeRO) shards the batch and, optionally, the state — but every device still runs the whole model graph. When a single layer is too big, or the pipeline of layers is too deep, two other axes come in. They compose: real systems run 3D parallelism (data × tensor × pipeline).
| Kind | What it splits | The cost it pays |
|---|---|---|
| Data parallel (this module) | the batch | an all-reduce of gradients per step |
| Tensor parallel | each layer’s matrices (split \(QKV\)/FFN across devices) | an all-reduce inside every layer — needs fast intra-node links |
| Pipeline parallel | the stack of layers into stages | a “bubble” of idle time; hidden by micro-batching |
| ZeRO / FSDP | the optimizer/grad/param state | extra all-gathers to reassemble params just-in-time |
One more lever is orthogonal to all of these — gradient checkpointing (a.k.a. activation recomputation). The backward pass needs the activations from the forward pass; storing them all costs memory that grows with depth × batch × sequence. Checkpointing keeps only a few and recomputes the rest during the backward pass — trading roughly one extra forward pass ($$33% more compute) for a large activation-memory cut. It stacks on top of everything above.
Common Pitfalls
WarningUneven shards break the “one big batch” equality
The gradient average equals the full-batch gradient only when the shards are the same size. With ragged shards you’d need a weighted average by sample count; data_parallel_gradients rejects uneven shards rather than silently bias the step.
WarningScaling \(N\) scales the effective batch — retune the LR
\(N\) devices with per-device batch \(b\) is one global batch of \(Nb\). Larger batches usually need a larger learning rate (the linear-scaling rule) and a warmup, or the first steps diverge. Distribution changes the batch, and the batch changes the optimizer (m07).
WarningCommunication can dominate if you don’t overlap it
The all-reduce moves \(\sim 2M\) bytes every step. Naively you’d compute, then communicate, then wait. Real frameworks overlap the gradient all-reduce of early layers with the backward pass of later ones, so most of the transfer is free. Without overlap, distributed training can be slower than one device.
WarningZeRO trades memory for communication, not for free
Higher ZeRO stages add all-gather traffic to reassemble what they partitioned. Stage 3 on a slow interconnect can be communication-bound. Pick the lowest stage that makes the model fit — stage 1 is often enough and the cheapest.
Exercises
Exercise 1: All-reduce mean from all-reduce sum
ring_all_reduce returns the sum. Wrap it to return the mean (what gradient averaging actually needs), and confirm it matches torch.stack(...).mean(0).
import torch
from distributed import ring_all_reduce
def all_reduce_mean(tensors):
# Your implementation here: sum, then divide by the device count.
summed = ring_all_reduce(tensors)
n = len(tensors)
return [s / n for s in summed]
xs = [torch.randn(6) for _ in range(4)]
got = all_reduce_mean(xs)[0]
print("matches mean:", torch.allclose(got, torch.stack(xs).mean(0), atol=1e-6))matches mean: TrueExercise 2: Derive the stage-2 memory formula
Stage 2 partitions the optimizer states (12 bytes) and the gradients (2 bytes) across \(N\), but replicates the fp16 params (2 bytes). Write the per-device byte formula and check it against zero_memory_per_device(psi, N, 2).
from distributed import zero_memory_per_device
def stage2_bytes(psi, n):
# Your implementation here: 2 bytes replicated + (2 + 12) bytes partitioned.
return 2 * psi + (2 + 12) * psi / n
psi, n = 1_000_000, 8
print("mine:", stage2_bytes(psi, n))
print("ref: ", zero_memory_per_device(psi, n, 2))mine: 3750000.0
ref: 3750000.0Exercise 3: Gradient accumulation — data parallelism on one device
Gradient accumulation gets a big effective batch on one device by summing gradients over several micro-batches before stepping. Show it equals the data-parallel gradient over the same shards.
import torch
from distributed import data_parallel_gradients
torch.manual_seed(1)
W = torch.randn(2, 3)
shards_x = [torch.randn(4, 3) for _ in range(3)]
shards_y = [torch.randn(4, 2) for _ in range(3)]
def accumulate(W, xs, ys):
# Your implementation here: mean-grad per micro-batch, averaged.
grads = []
for x, y in zip(xs, ys):
w = W.detach().clone().requires_grad_(True)
loss = ((x @ w.T - y) ** 2).mean()
grads.append(torch.autograd.grad(loss, w)[0])
return sum(grads) / len(grads)
acc = accumulate(W, shards_x, shards_y)
dp = data_parallel_gradients(W, shards_x, shards_y)
print("accumulation == data parallel:", torch.allclose(acc, dp, atol=1e-6))accumulation == data parallel: TrueSummary
Key takeaways:
- Training memory is \(16\) bytes/param — under mixed-precision Adam: 2 (fp16 params) + 2 (fp16 grads) + 12 (fp32 master + momentum + variance). The optimizer state, not the parameters, is what doesn’t fit.
- Data parallelism replicates the model and averages gradients, and that average is bit-for-bit the gradient one big-batch device would compute — a distributed step is not an approximation.
- The ring all-reduce does the averaging in two phases (reduce-scatter → all-gather), bit-exact to the naive sum, at \(2(N{-}1)/N \cdot M\) bytes per device — a fixed budget as \(N\) grows, which is why it, not a central reducer, is the standard.
- ZeRO removes data parallelism’s redundancy in three stages — partitioning the optimizer states (\(4\Psi + 12\Psi/N\)), then gradients (\(2\Psi + 14\Psi/N\)), then parameters (\(16\Psi/N\)) — up to an \(N\times\) memory cut. Stage 3 is FSDP.
- Other axes compose: tensor parallel splits each layer, pipeline parallel splits the stack, and gradient checkpointing trades compute for activation memory. Frontier runs combine all of them.
What’s Next
You can now build the model (m01–m06), train it (m07), align it (m12), serve it (m08–m09), and scale its training across a cluster (this module). The frontier from here: tensor & pipeline parallelism as runnable code, FSDP wired into a real GPTModel, and gradient checkpointing built from an autograd hook — all filed on the roadmap.
Going Deeper
Core Papers:
- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models — Rajbhandari et al. (2020): the \(2+2+12\) memory model and the three-stage partitioning built here.
- Bandwidth Optimal All-reduce Algorithms for Clusters of Workstations — Patarasuk & Yuan (2009): the ring all-reduce and its \(2(N{-}1)/N \cdot M\) cost.
- Horovod: fast and easy distributed deep learning in TensorFlow — Sergeev & Del Balso (2018): ring all-reduce brought to deep learning.
- PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel — Zhao et al. (2023): ZeRO-3 as a production PyTorch API.
- Training Deep Nets with Sublinear Memory Cost — Chen et al. (2016): gradient checkpointing.
Practical Resources:
- DeepSpeed — the library that popularized ZeRO.
- Megatron-LM — Shoeybi et al. (2019): tensor parallelism for transformers.