Distributed training at scale
Training a model that doesn't fit on one GPU is a memory and communication problem before it is a math problem. This lesson builds the three axes of parallelism — data, tensor, and pipeline — plus ZeRO/FSDP sharding, models the all-reduce/all-gather communication cost and the memory math offline in pure Python, composes them into 3D parallelism, and closes on failure and recovery.
Learning objectives
- Distinguish data, tensor, and pipeline parallelism by what each splits and what it communicates.
- Explain the ZeRO stages (and FSDP) — how sharding optimizer state, gradients, and parameters cuts per-GPU memory.
- Model the communication cost of all-reduce and all-gather with the ring-algorithm formula, offline.
- Do the memory math: parameters, gradients, optimizer state, and activations per GPU.
- Compose the axes into 3D parallelism and reason about failure and recovery (checkpointing, straggler/rank loss).
1 · Three axes: data, tensor, pipeline
There is more than one reason a model won't train on one GPU, and each has its own axis of parallelism. Data parallelism (DP) replicates the whole model on every GPU and splits the batch; each GPU computes gradients on its shard and they are averaged with an all-reduce. It scales throughput but every GPU must hold the full model. Tensor parallelism (TP) splits individual layers (e.g. a big matmul's columns) across GPUs that cooperate on one forward pass — heavy communication, so it stays inside a node with fast NVLink. Pipeline parallelism (PP) splits the model by layer stages across GPUs, streaming micro-batches through like an assembly line — cheap communication but a bubble of idle time at the pipe's fill and drain.
| Axis | Splits | Communicates | Best when |
|---|---|---|---|
| Data (DP) | the batch | gradients (all-reduce) | model fits on one GPU, want throughput |
| Tensor (TP) | layers / matmuls | activations, every layer | layer too big for one GPU; stay in-node |
| Pipeline (PP) | layer stages | stage boundaries only | model too deep; cross-node OK |
2 · ZeRO / FSDP: shard the optimizer state
Data parallelism's weakness is that every GPU holds a full copy of parameters, gradients, and optimizer state — and for Adam the optimizer state (two moments in fp32) plus fp32 master weights dwarfs the model itself. ZeRO (Zero Redundancy Optimizer; PyTorch's FSDP is the same idea) removes that redundancy by sharding those tensors across the DP group, reconstructing each shard's neighbours on the fly with communication:
| ZeRO stage | Shards | Per-GPU memory | Extra comms |
|---|---|---|---|
| Stage 1 | optimizer state | big cut | none beyond DP all-reduce |
| Stage 2 | + gradients | bigger cut | reduce-scatter gradients |
| Stage 3 (FSDP) | + parameters | ~1/N of full | all-gather params per layer |
Stage 3 is the extreme: no GPU ever holds the whole model. Each layer's parameters are all-gathered just before its forward/backward and freed after — trading extra communication for a roughly 1/N memory footprint. That trade is what lets a cluster train a model far larger than any single GPU's HBM.
3 · Lab · the per-GPU memory math
Before choosing a strategy, do the arithmetic: what actually has to fit on one GPU? For mixed-precision Adam the bytes-per-parameter are famous — fp16 weights (2) + fp16 grads (2) + fp32 master weights (4) + two fp32 Adam moments (4+4) ≈ 16 bytes/param, plus activations. This lab computes it and shows how ZeRO-3 sharding across N GPUs cuts it.
train_memory.pyBYTES_PER_PARAM = 16 # fp16 w(2) + fp16 grad(2) + fp32 master(4) + Adam moments(4+4)
def per_gpu_gb(params, zero_stage, n_gpus):
full = params * BYTES_PER_PARAM
if zero_stage == 0: # full replica
state = full
elif zero_stage == 1: # shard optimizer state (12 of the 16 bytes)
state = params * 4 + (params * 12) / n_gpus
elif zero_stage == 2: # + shard gradients (2 bytes)
state = params * 2 + (params * 14) / n_gpus
else: # ZeRO-3: shard everything
state = full / n_gpus
return state / 1e9
P, N = 7e9, 8
print(f"Model: {P/1e9:.1f}B params")
full = per_gpu_gb(P, 0, N)
print(f" full replica (DP): {full:6.2f} GB/GPU state -> WON'T FIT on an 80GB GPU")
print(f" ZeRO-1 (opt state sharded, N={N}): {per_gpu_gb(P,1,N):6.2f} GB/GPU")
print(f" ZeRO-2 (+ grads, N={N}): {per_gpu_gb(P,2,N):6.2f} GB/GPU")
print(f" ZeRO-3 (+ params, N={N}): {per_gpu_gb(P,3,N):6.2f} GB/GPU "
f"-> fits, room for activations")
print("sharding turns a model that won't fit into one that does")
Model: 7.0B params
full replica (DP): 112.00 GB/GPU state -> WON'T FIT on an 80GB GPU
ZeRO-1 (opt state sharded, N=8): 50.00 GB/GPU
ZeRO-2 (+ grads, N=8): 36.00 GB/GPU
ZeRO-3 (+ params, N=8): 14.00 GB/GPU -> fits, room for activations
sharding turns a model that won't fit into one that doesThe full replica needs 112 GB just for training state — impossible on an 80 GB GPU. ZeRO-3 across 8 GPUs brings it to ~14 GB each, leaving headroom for activations. This is the calculation that decides your parallelism plan.
4 · Lab · communication cost of all-reduce
Communication is the tax on parallelism. The workhorse collective is all-reduce (averaging gradients across GPUs). The efficient ring all-reduce has a beautiful property: each GPU sends and receives about 2·(N−1)/N × message_size bytes — essentially independent of N for large N (it approaches 2× the message). All-gather (used by ZeRO-3 to reconstruct parameters) moves about (N−1)/N × total. This lab models both and the resulting time.
comms_cost.pydef ring_allreduce_gb(msg_gb, n):
return 2 * (n - 1) / n * msg_gb # bytes sent+received per GPU
def ring_allgather_gb(msg_gb, n):
return (n - 1) / n * msg_gb
def ms(gb, link_gb_s):
return gb / link_gb_s * 1000
MSG, LINK = 14.0, 200.0 # 14 GB gradient buffer, 200 GB/s link
print(f"gradient buffer: {MSG:.2f} GB across N=8 GPUs, link={LINK:.0f} GB/s")
ar = ring_allreduce_gb(MSG, 8); ag = ring_allgather_gb(MSG, 8)
print(f" ring all-reduce: moves {ar:5.2f} GB/GPU -> {ms(ar,LINK):5.1f} ms")
print(f" ring all-gather: moves {ag:5.2f} GB/GPU -> {ms(ag,LINK):5.1f} ms")
print("N=64 GPUs:")
ar64 = ring_allreduce_gb(MSG, 64)
print(f" ring all-reduce: moves {ar64:5.2f} GB/GPU -> {ms(ar64,LINK):5.1f} ms (~flat vs N)")
print("all-reduce time is ~independent of N -- that's why ring scales")
gradient buffer: 14.00 GB across N=8 GPUs, link=200 GB/s
ring all-reduce: moves 24.50 GB/GPU -> 122.5 ms
ring all-gather: moves 12.25 GB/GPU -> 61.2 ms
N=64 GPUs:
ring all-reduce: moves 27.56 GB/GPU -> 137.8 ms (~flat vs N)
all-reduce time is ~independent of N -- that's why ring scalesThe key result: ring all-reduce cost per GPU barely grows with N (8 vs 64 GPUs move nearly the same bytes). That flatness is why data parallelism scales to thousands of GPUs — the gradient sync doesn't blow up. TP is the opposite: it all-reduces activations every layer, so it only pays off over fast in-node links.
5 · 3D parallelism: composing the axes
At frontier scale you use all three at once — 3D parallelism. The standard recipe layers them by communication cost onto the hardware topology: tensor parallelism within a node (fast NVLink, since TP talks every layer), pipeline parallelism across a few nodes (cheap stage-boundary comms), and data parallelism across the remaining replicas (flat-scaling all-reduce). The total GPU count is the product: world_size = DP × TP × PP.
three_d_plan.pydef plan_3d(world, gpus_per_node, tp):
assert tp <= gpus_per_node, "TP must fit within one node (NVLink)"
assert world % tp == 0, "world size must divide by TP"
# pick PP = number of nodes spanned per replica; DP absorbs the rest
pp = world // gpus_per_node # one pipeline stage group per node here
remaining = world // (tp * pp)
dp = remaining
assert tp * pp * dp == world, "layout must multiply to world size"
return tp, pp, dp
world, gpn, tp = 256, 8, 8
tp, pp, dp = plan_3d(world, gpn, tp)
print(f"world size = {world} GPUs")
print(f" TP={tp} (within a node, over NVLink)")
print(f" PP={pp} (across {pp} nodes)")
print(f" DP={dp} (replicas of the TPxPP group)")
print(f" check: {tp} x {pp} x {dp} = {tp*pp*dp} OK")
print(f" each DP replica spans {tp*pp} GPUs (TPxPP)")
world size = 256 GPUs
TP=8 (within a node, over NVLink)
PP=4 (across 4 nodes)
DP=8 (replicas of the TPxPP group)
check: 8 x 4 x 8 = 256 OK
each DP replica spans 32 GPUs (TPxPP)6 · Failure & recovery at scale
At thousands of GPUs for weeks, hardware will fail — the question is only how gracefully. The bedrock is checkpointing: periodically save model + optimizer state so a crash costs only the work since the last save. The tradeoff is frequency: checkpoint too often and the I/O stalls training; too rarely and a crash is expensive. The optimal interval balances checkpoint cost against expected lost work given the failure rate — the classic MTBF calculation.
Beyond crashes: a straggler (one slow GPU) throttles every synchronous collective to its pace, so you detect and evict it; losing a rank in a TP/PP group is fatal to that group until it's replaced and re-loaded from checkpoint. Elastic training frameworks reshape the world size around lost nodes. This lab models the checkpoint-interval tradeoff.
checkpoint_interval.pyimport math
def waste_fraction(interval_min, ckpt_cost_min, mtbf_min):
ckpt_overhead = ckpt_cost_min / interval_min # time spent checkpointing
lost_work = interval_min / (2 * mtbf_min) # avg work lost per crash
return ckpt_overhead, lost_work
MTBF = 50 * 60 # 50 hours in minutes
CKPT = 3.0 # 3 minutes to write a checkpoint
print(f"failure every ~{MTBF/60:.1f} h, checkpoint costs {CKPT:.1f} min:")
for interval in (15, 60, 240, 480):
ov, lost = waste_fraction(interval, CKPT, MTBF)
print(f" interval={interval:>3} min: overhead {ov*100:4.1f}% ckpt + "
f"{lost*100:.1f}% lost = {(ov+lost)*100:4.1f}% wasted")
opt = math.sqrt(2 * CKPT * MTBF)
print(f" ~optimal interval near sqrt(2 * ckpt_cost * MTBF) = {opt:.0f} min")
failure every ~50.0 h, checkpoint costs 3.0 min:
interval= 15 min: overhead 20.0% ckpt + 0.1% lost = 20.1% wasted
interval= 60 min: overhead 5.0% ckpt + 0.5% lost = 5.5% wasted
interval=240 min: overhead 1.2% ckpt + 2.0% lost = 3.2% wasted
interval=480 min: overhead 0.6% ckpt + 4.0% lost = 4.6% wasted
~optimal interval near sqrt(2 * ckpt_cost * MTBF)Your team wants to scale training from 8 to 512 GPUs. A colleague worries the gradient all-reduce will become the bottleneck as GPU count grows. Is that the right worry for data parallelism, and where should the real concern be?
Show answer
2·(N−1)/N times the message per GPU, which approaches a constant (~2× the gradient buffer) and is essentially independent of N. That flatness is exactly why DP scales to thousands of GPUs. The real concern at 512 GPUs is memory (a full replica per GPU may not fit — use ZeRO/FSDP sharding) and, if you add tensor parallelism, its per-layer all-reduce, which does not scale over slow links and must stay inside a node on NVLink.Why is tensor parallelism kept inside a single node while pipeline and data parallelism span nodes?
Show answer
✓ Checkpoint — you can move on when you can…
- Distinguish DP, TP, and PP by what each splits and communicates, and where each belongs on the topology.
- Explain the ZeRO stages / FSDP and how sharding optimizer state, grads, and params cuts per-GPU memory.
- Compute per-GPU training memory (16 bytes/param for Adam) and how sharding across N reduces it.
- Model ring all-reduce cost and explain why it is ~independent of N.
- Lay out a 3D-parallel world (DP×TP×PP) and reason about checkpointing and rank loss.
🪜 Practice ladder beginner → industry
Six graded exercises, easy to real-world. Try each before opening its solution.
Context: The three axes are told apart by what they split and what they send. Mixing them up leads to putting the chatty axis on a slow link.
Your task: Write a lookup that, given a parallelism axis, returns what it splits and what it communicates.
Requirements:
- Data → splits the batch, communicates gradients (all-reduce)
- Tensor → splits layers/matmuls, communicates activations every layer
- Pipeline → splits layer stages, communicates at stage boundaries
- Print each axis with its split and its comms
💡 Hint: Data = batch, Tensor = layer internals, Pipeline = layer stages.
Show solution
A direct lookup of split and comms per axis. Runnable:
AXES = {
"data": ("the batch", "gradients (all-reduce)"),
"tensor": ("layers / matmuls", "activations, every layer"),
"pipeline": ("layer stages", "stage boundaries only"),
}
for axis, (splits, comms) in AXES.items():
print(f"{axis:9} splits {splits:20} communicates {comms}")Data splits the batch and averages gradients; tensor splits inside layers and talks every layer; pipeline splits by stage and talks only at boundaries. Keeping these straight is what tells you the chatty axis (tensor) belongs on the fastest link.
Context: Mixed-precision Adam costs ~16 bytes per parameter of state, which is why big models don't fit even when the fp16 weights alone would.
Your task: Compute the per-GPU training-state memory for a model, then show ZeRO-3 sharding across N GPUs.
Requirements:
- Bytes/param ≈ 2 (fp16 w) + 2 (fp16 grad) + 4 (fp32 master) + 8 (Adam moments) = 16
- Full state = params × 16 bytes
- ZeRO-3 divides that by N
- Show a 7B model won't fit an 80GB GPU as a full replica but fits when sharded
💡 Hint: The optimizer state, not the weights, is the memory hog.
Show solution
Adam state is ~16 bytes/param; ZeRO-3 divides it by N. Runnable:
def per_gpu_gb(params, sharded_over=1):
bytes_per = 16 # fp16 w(2)+grad(2)+fp32 master(4)+Adam moments(8)
return params * bytes_per / sharded_over / 1e9
P = 7e9
print(f"full replica: {per_gpu_gb(P):.0f} GB/GPU (80GB GPU: WON'T FIT)")
print(f"ZeRO-3 / N=8: {per_gpu_gb(P, 8):.0f} GB/GPU (fits, room for activations)")The 7B model needs 112 GB of training state as a full replica — impossible on an 80 GB GPU — but only ~14 GB when ZeRO-3 shards it across 8 GPUs. The optimizer state, not the weights, is the memory hog, and sharding it is what makes big-model training fit.
Context: Ring all-reduce moves about 2(N-1)/N times the message per GPU — nearly constant in N. That formula is why data parallelism scales.
Your task: Model the bytes each GPU moves in a ring all-reduce and the time given a link bandwidth, for several N.
Requirements:
- Per-GPU bytes =
2 × (N−1)/N × message_bytes - Time = bytes / link_bandwidth
- Report for N = 8, 64, 512
- Show the per-GPU cost approaches a constant (~2× message) as N grows
💡 Hint: The (N-1)/N factor approaches 1, so the cost approaches 2x the message regardless of N.
Show solution
Ring all-reduce per-GPU bytes are ~independent of N. Runnable:
def ring_allreduce_gb(msg_gb, n):
return 2 * (n - 1) / n * msg_gb
MSG, LINK = 14.0, 200.0 # GB, GB/s
for n in (8, 64, 512):
gb = ring_allreduce_gb(MSG, n)
print(f"N={n:>3}: {gb:5.2f} GB/GPU -> {gb/LINK*1000:5.1f} ms")
print("cost approaches 2x the message (~28 GB) regardless of N")From 8 to 512 GPUs the per-GPU traffic barely moves (the (N-1)/N factor just creeps toward
1). Because the gradient sync is flat in N, data parallelism scales to thousands of GPUs — the reason it is
the outermost axis.
Context: Pipeline parallelism idles GPUs while the pipe fills and drains — the bubble. More micro-batches amortize it; the bubble fraction is (stages-1)/(stages-1+microbatches).
Your task: Compute the pipeline bubble fraction and show how increasing micro-batches shrinks it.
Requirements:
- Bubble fraction =
(P−1) / (P−1 + M)for P stages, M micro-batches - Report the fraction for a few values of M
- Show more micro-batches drive the bubble toward zero
- State the tradeoff: more micro-batches means more activation memory
💡 Hint: With P stages, the first output appears after P-1 warmup steps; M micro-batches amortize that.
Show solution
The bubble fraction is (P-1)/(P-1+M). Runnable:
def bubble_fraction(stages, microbatches):
return (stages - 1) / (stages - 1 + microbatches)
P = 4
for M in (1, 4, 16, 64):
print(f"P={P} stages, M={M:>2} micro-batches: bubble = {bubble_fraction(P,M)*100:4.1f}%")
print("more micro-batches shrink the bubble -- but cost more activation memory")With 4 stages and 1 micro-batch, 75% of the pipeline is idle warmup/drain; at 64 micro-batches the bubble is under 5%. That is why pipeline parallelism uses many micro-batches — bounded only by the extra activation memory each in-flight micro-batch holds.
Context: A frontier run is DP x TP x PP = world size, with TP inside a node, PP across nodes, DP outermost. The layout must multiply out and respect the node size.
Your task: Write a planner that, given world size, GPUs-per-node, and a desired TP degree, chooses PP and DP so the product matches and TP fits within a node.
Requirements:
DP × TP × PP == world_size- TP must be ≤ GPUs per node (stays on NVLink)
- Derive PP and DP from the remaining factors
- Print the layout and verify the product, rejecting impossible requests
💡 Hint: TP is bounded by the node; PP and DP absorb the rest of the world size.
Show solution
TP fits the node; PP and DP absorb the rest. Runnable:
def plan_3d(world, gpus_per_node, tp, pp):
if tp > gpus_per_node:
return None, "TP exceeds node size -- would cross NVLink"
if world % (tp * pp) != 0:
return None, "TP*PP must divide world size"
dp = world // (tp * pp)
return (tp, pp, dp), f"{tp} x {pp} x {dp} = {tp*pp*dp}"
layout, msg = plan_3d(world=256, gpus_per_node=8, tp=8, pp=4)
print("layout (TP,PP,DP):", layout, "check:", msg)
bad, why = plan_3d(world=256, gpus_per_node=8, tp=16, pp=4)
print("rejected:", why)TP is capped at the node's 8 GPUs so it stays on NVLink; PP spans a few nodes; DP absorbs the remaining factor so the product hits the world size. A TP request larger than the node is rejected because it would force the chattiest axis onto a slow cross-node link.
Context: At scale, hardware fails on a schedule. Checkpoint too often and I/O dominates; too rarely and a crash is expensive. The optimum is roughly sqrt(2 x ckpt_cost x MTBF).
Your task: Model total wasted-time fraction as a function of checkpoint interval and find the interval that minimizes it.
Requirements:
- Checkpoint overhead fraction =
ckpt_cost / interval - Expected lost-work fraction ≈
interval / (2 × MTBF) - Total waste = sum of the two; sweep intervals and report the minimum
- Compare against the closed-form
sqrt(2 × ckpt_cost × MTBF)
💡 Hint: The two terms trade off; their sum is minimized where they are balanced.
Show solution
Balance checkpoint overhead against expected lost work. Runnable:
import math
def waste(interval, ckpt_cost, mtbf):
return ckpt_cost / interval + interval / (2 * mtbf) # overhead + lost work
CKPT, MTBF = 3.0, 50 * 60 # minutes
best = min(range(5, 600, 5), key=lambda i: waste(i, CKPT, MTBF))
print(f"swept optimum interval: {best} min (waste {waste(best,CKPT,MTBF)*100:.2f}%)")
print(f"closed form sqrt(2*ckpt*MTBF): {math.sqrt(2*CKPT*MTBF):.0f} min")The swept minimum matches the closed form sqrt(2 × ckpt_cost × MTBF): checkpoint
too often and write-I/O dominates, too rarely and each crash burns hours. Balancing the two is the whole
game of resilient large-scale training.