Distributed training at scale

1 · Three axes: data, tensor, pipeline

Start with the moment it stops fitting. You have a training script that ran perfectly on a 7B model on one 80 GB GPU. Your team bumps to 70B, you launch the same script on the same card, and it dies before the first step with CUDA out of memory. Your instinct is "buy a bigger GPU" — but there is no bigger GPU; 80 GB is roughly the ceiling. The only way forward is to spread one training job across many GPUs. The mistake almost everyone makes first is assuming that's a single technique. It isn't — and picking the wrong one wastes a cluster.

Here is the mental model that unlocks the whole lesson: think of a busy restaurant kitchen that suddenly gets ten times the orders. There are three genuinely different ways to add cooks, and they are not interchangeable. You can clone the whole kitchen — ten identical kitchens, each handling one-tenth of the tickets (that's data parallelism: replicate the model, split the batch). You can put several cooks on one dish that's too big for one person — one chops, one sears, one plates, all on the same plate at the same time (that's tensor parallelism: split a single layer across GPUs that cooperate on one forward pass). Or you can run an assembly line — station 1 does prep, hands to station 2, which cooks, hands to station 3, which plates (that's pipeline parallelism: split the model into stages, stream micro-batches through). Same goal — more throughput — but the communication each style demands is wildly different, and that difference is what decides where each axis is allowed to live in your cluster.

Read the paragraph and table below through that lens. The one thing to internalise: the question is never "should I parallelise?" but "which axis solves my bottleneck?" A model that fits but trains too slowly wants data parallelism; a single layer that won't fit wants tensor parallelism; a model too deep to fit even one copy wants pipeline parallelism. Real frontier runs use all three at once (§5), but you only understand the combination once you can name what each one splits and what it must send over the wire — because the thing it sends is the tax you pay, and the whole art is keeping that tax on the fastest link you have.

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.

2 · ZeRO / FSDP: shard the optimizer state

Here is the number that surprises everyone the first time they do the arithmetic. You picture a 7B model as "7 billion weights in fp16, so ~14 GB — easy, it fits." Then you try to train it with plain data parallelism and it dies. Where did the memory go? The weights are the small part. To train with mixed-precision Adam, every GPU is also holding fp16 gradients, an fp32 master copy of the weights, and Adam's two fp32 momentum buffers — about 16 bytes per parameter, not 2. That's 112 GB of training state for a 7B model, on a card that maxes out at 80. The weights were never the problem; the optimizer state was.

Now the key observation that makes ZeRO obvious. In plain data parallelism, all N GPUs hold the exact same copy of that 112 GB. That is pure redundancy — you are storing the same optimizer state N times and gaining nothing from the duplication except convenience. The analogy: eight librarians who each keep a full identical copy of the same encyclopedia set, so any of them can answer any question instantly. It works, but you've bought eight encyclopedias to serve what one collection could, if the librarians were willing to pass volumes to each other when needed.

ZeRO (Zero Redundancy Optimizer — PyTorch ships the same idea as FSDP, Fully Sharded Data Parallel) is exactly that willingness to share. Instead of every GPU keeping the whole encyclopedia, each keeps only its slice (one-Nth of the optimizer state, gradients, and eventually parameters) and fetches a neighbour's slice over the interconnect at the precise moment it's needed, then throws it away. You trade extra communication for a dramatic cut in per-GPU memory. The beauty is that it comes in graded stages, so you can pay only as much communication as you must: shard just the optimizer state (Stage 1, almost free), add gradients (Stage 2), or shard parameters too (Stage 3 / FSDP, the most aggressive). The table lays out what each stage shards, what memory it saves, and what it costs you in return — read it as a dial you turn up only as far as your memory pressure demands.

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:

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

Every distributed-training decision you will ever make starts with one back-of-envelope sum, and skipping it is the most expensive mistake in this whole field. Teams routinely burn a day spinning up a 64-GPU job, watching it OOM on step one, guessing at a fix, and re-launching — when ninety seconds of arithmetic beforehand would have told them exactly which parallelism to use and how many GPUs it needs. The memory math is not a formality; it is the plan.

The sum has two parts and a famous constant. First, the fixed per-parameter cost: for mixed-precision Adam it's ~16 bytes per parameter — memorise the breakdown because you'll recite it in interviews and design reviews alike. fp16 weights (2 bytes) + fp16 gradients (2) + an fp32 master copy of the weights (4) + Adam's first and second moment estimates in fp32 (4 + 4) = 16. Notice that the weights you think of as "the model" are only 2 of those 16 bytes; the other 14 exist purely to train it, which is why inference fits on hardware that training never could. Second, activations — the intermediate tensors saved during the forward pass so the backward pass can compute gradients — which scale with batch size and sequence length rather than parameter count, and which this lab deliberately leaves as headroom rather than modelling exactly (they're workload-dependent and easy to get wrong).

The lab below turns that into runnable Python: it computes the full-replica cost, then shows how each ZeRO stage divides the shardable portion by N. Watch the specific transition it prints — a 7B model's 112 GB of state (impossible on an 80 GB card) collapsing to ~14 GB once ZeRO-3 spreads it across 8 GPUs. That single number flip — "won't fit" to "fits, with room to spare" — is the entire justification for sharding, and once you can reproduce it by hand you can size any run. One honest caveat as you read: these figures are the theoretical state budget; real frameworks add communication buffers, memory fragmentation, and the activation term above, so treat the result as a floor, not a promise, and always leave margin.

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 does

The 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

Here's the fear that stops teams from scaling — and it's mostly wrong. "If we go from 8 GPUs to 512, surely averaging gradients across 512 machines will drown the network and everything grinds to a halt?" It's a reasonable worry, and for the naive way of averaging it would be true. But the algorithm the whole field actually uses — the ring all-reduce — has a near-magical property: the bytes each GPU must move barely change whether you have 8 GPUs or 8,000. Understanding why is what tells you data parallelism is safe to scale and tensor parallelism is not.

First, what an all-reduce even is. After a training step, each GPU has computed gradients on its own slice of the batch. To take one synchronised optimizer step they must agree on the average gradient — so every GPU needs the sum of everyone's gradients. "All-reduce" is the collective operation that does exactly this: reduce (sum) across all ranks, and give the result back to all of them. The naive implementation — everyone sends everything to one GPU, it sums, it sends the answer back — makes that one GPU a bottleneck whose traffic grows linearly with N. The ring algorithm is the clever fix: arrange the GPUs in a circle, chop each gradient buffer into N chunks, and pass chunks around the ring so that every link carries a steady trickle and no single GPU is a hub. The arithmetic falls out as each GPU sending and receiving about 2·(N−1)/N × message_size bytes — and the (N−1)/N factor creeps toward 1 as N grows, so the cost flattens at roughly twice the gradient buffer, no matter how many GPUs you add.

The related collective, all-gather (which ZeRO-3 uses to rebuild each layer's full parameters from the shards just before it runs), moves about (N−1)/N × total — same flattening behaviour. The lab models both and converts bytes to milliseconds using a link bandwidth, so you can feel the two forces that actually govern a run: message size and interconnect speed, not GPU count. Keep this contrast in your pocket for §5: because DP's all-reduce is flat in N and can be overlapped with backward-pass compute (you sync layer L's gradients while still computing layer L−1's), it scales to the whole cluster — whereas tensor parallelism must all-reduce activations at every single layer, high-frequency traffic that only survives on the fastest link in the building.

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 scales

The 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

You now know three axes — and the natural next question trips up almost everyone: "which one do I pick?" The answer, at frontier scale, is a surprise: you don't pick, you layer all three at once. A single 70B-plus run on hundreds of GPUs is simultaneously tensor-, pipeline-, and data-parallel. The reason isn't greed for complexity; it's that each axis solves a different limit — TP makes an oversized layer fit, PP makes an oversized depth fit, DP turns the surviving replica into throughput — and a real model hits all three limits at once. Combining them is called 3D parallelism.

The insight that makes it tractable — rather than a nightmare of tangled communication — is that your hardware is not flat. A GPU cluster is a hierarchy of link speeds: within a single node, 8 GPUs talk over NVLink at hundreds of GB/s; between nodes, they talk over a network (InfiniBand or Ethernet) that is often an order of magnitude slower. That hierarchy is the whole game. Picture a company's communication: teammates at adjacent desks (NVLink) can whisper constantly at no cost; different buildings (cross-node) exchange far fewer, more deliberate messages. You would never put two people who must talk every second in different buildings — and you would never put your chattiest parallelism axis on your slowest link.

So the placement rule writes itself, and it is the one thing to carry out of this section: rank the axes by how often they communicate, and map the chattiest onto the fastest link. Tensor parallelism talks every layer (activations back and forth) → keep it inside a node on NVLink. Pipeline parallelism talks only at stage boundaries (a few tensors per micro-batch) → it tolerates a slower cross-node hop. Data parallelism's all-reduce is flat in N and overlaps with compute (you proved this in §4) → it sits outermost, spanning the whole cluster. The three degrees multiply to your total GPU count — world_size = DP × TP × PP — and the planner lab below turns that identity into a concrete layout, refusing any request that would shove TP across a node boundary. The common mistake it guards against is picking round numbers that "multiply out" while quietly stranding the highest-frequency traffic on the wrong wire.

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

Do the grim arithmetic once and you'll never treat failure as an edge case again. Suppose a single GPU has a mean time between failures of, say, a few years — reassuringly rare. Now run 4,000 of them together for three weeks. The collective failure rate is the sum, and it collapses that "few years" into a failure every few hours. A frontier training run is not a program that might crash; it is a program that is guaranteed to crash repeatedly before it finishes. The engineering question is never "how do I prevent failure" — you can't — but "how little work do I lose each time, and how fast do I recover."

The bedrock answer is checkpointing: periodically write the full model and optimizer state to durable storage, so a crash costs only the work done since the last save, not the whole run. (Note "and optimizer state" — a checkpoint that saves weights but forgets Adam's momentum buffers resumes into a training disturbance, a classic and painful bug.) But checkpointing is itself expensive — writing hundreds of gigabytes stalls the whole cluster — so you cannot just save constantly. This is a genuine optimization with a sweet spot, and it's worth feeling the two opposing forces: checkpoint too often and write-I/O overhead eats your throughput; checkpoint too rarely and every crash burns hours of recomputation. The total wasted fraction is the sum of those two, and — like many such trade-offs — it's minimised where they roughly balance, at an interval near sqrt(2 × checkpoint_cost × MTBF). The lab makes that curve concrete so you can see the U-shape rather than take the formula on faith.

Two failure modes beyond a clean crash deserve names because they behave differently. A straggler is a GPU that hasn't died but has gone slow (thermal throttling, a flaky link) — and because every synchronous collective waits for the last participant, one straggler drags the entire cluster down to its pace; the fix is to detect and evict it, not wait politely. Losing a whole rank inside a tensor- or pipeline-parallel group is worse: that group cannot complete a forward pass at all until the rank is replaced and reloaded from the last checkpoint, which is why elastic training frameworks exist to reshape the world size around lost nodes and resume. The through-line for this entire section: at scale, resilience isn't a feature you bolt on at the end — the checkpoint interval, the straggler detector, and the elastic reshape are as much a part of the training design as the parallelism plan itself.

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)

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}")

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)")

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")

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")

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)

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")