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.
  • Build tensor parallelism (Megatron-LM) from scratch — column-parallel and row-parallel GEMMs, the MLP and attention split, and the f/g conjugate operators — and prove the split layer is bit-exact to one device.
  • Build pipeline parallelism (GPipe / 1F1B) from scratch — the device × time schedule, the (P{-}1)/(M{+}P{-}1) bubble and why micro-batching shrinks it, and why 1F1B keeps the same bubble at O(P) instead of O(M) activation memory.
  • Build ring attention (sequence / context parallelism) from scratch — shard the sequence across devices, rotate K/V blocks around a ring, accumulate with an online softmax, and prove it equals full attention while context scales linearly with devices; plus the causal load-imbalance and the striped fix.
  • Place FSDP and gradient checkpointing in the same map, and reason about tensor parallelism’s memory-for-communication trade.

Prerequisites

This module requires familiarity with:

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 GB

Params 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/device

120 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
  1. 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.
  2. 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.
  3. 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|: 1.1920928955078125e-07
identical step: True

Four 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:

  1. Reduce-scatter (N{-}1 steps): chunks circulate and accumulate until each device owns the fully summed value of exactly one chunk.
  2. 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.0

Bit-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
  1. 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.
  2. 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.

Tensor Parallelism: Split Each Layer

Data parallelism keeps the whole model on each device and splits the work. But if a single layer’s weight matrices don’t fit in one device’s memory, no amount of batch-splitting helps — you have to split the matrix itself. Tensor parallelism (also called model parallelism or intra-layer parallelism) does exactly that: it cuts each layer’s GEMMs across devices and stitches the pieces back together with one collective. This is the scheme Megatron-LM introduced and every large training run uses.

There are only two ways to cut a matrix multiply Y = X A:

  • Column-parallel — split A by its output columns, A = [A_1 \mid A_2 \mid \dots \mid A_N]. Device i computes Y_i = X A_i, a slice of the output columns. No communication: the input X is shared, and each device owns a disjoint block of outputs.
  • Row-parallel — split A by its input rows, A = [A_1; A_2; \dots; A_N], and split the input to match, X = [X_1 \mid \dots \mid X_N]. Device i computes a partial X_i A_i; the full result is \sum_i X_i A_i — an all-reduce.
NoteKey Insight

Column-parallel needs no communication but leaves the output split. Row-parallel consumes a split input and pays one all-reduce to reassemble the output. Chain them — column-parallel then row-parallel — and the split from the first GEMM is exactly the layout the second one wants, so the whole two-layer block costs a single collective. That is the entire trick.

The Math: The Megatron MLP

The transformer’s feed-forward block is two GEMMs with a nonlinearity between them: Z = \text{GeLU}(X A)\,B, where A is d_\text{model}\times d_\text{ff} and B is d_\text{ff}\times d_\text{model}. Megatron splits A column-parallel and B row-parallel:

[\,Y_1 \mid \dots \mid Y_N\,] = \text{GeLU}\big(X\,[\,A_1 \mid \dots \mid A_N\,]\big), \qquad Z = \sum_{i=1}^{N} Y_i B_i .

Why this exact pairing? Because GeLU is elementwise and each output column of XA depends only on the matching column of A, the nonlinearity can be applied independently on each shard — there is no need to gather the columns first:

\text{GeLU}(XA)\ \text{split by columns} \;=\; \big[\,\text{GeLU}(XA_1) \mid \dots \mid \text{GeLU}(XA_N)\,\big].

Splitting A any other way would break this (you’d need to sum partial columns before the nonlinearity — a collective inside the block). Column-then-row is the one split that defers all communication to a single all-reduce at the very end.

Megatron names the two communication points with a conjugate operator pair:

  • f — identity in the forward pass, all-reduce in the backward pass (it sits at the block’s input, where the replicated X enters).
  • g — all-reduce in the forward pass, identity in the backward pass (it sits at the block’s output, summing the row-parallel partials).

Step through one forward pass of the split MLP on two devices:

Only one stage — the all-reduce (g) — moves data between devices. Everything else is a local GEMM or an elementwise op. The input and output boxes collapse to a single lane because X and Z are replicated on every device; the two GEMMs and the GeLU run split.

Code: the Split MLP from Scratch

tensor_parallel.py builds this directly. column_parallel_shards / row_parallel_shards cut the two weights, column_parallel_forward runs the first GEMM with no communication, and row_parallel_forward runs the second and calls the module’s own ring_all_reduce to sum the partials — the same collective data parallelism used, now inside a layer.

import torch
from tensor_parallel import megatron_mlp, single_device_mlp

torch.manual_seed(0)
x = torch.randn(2, 6, 64)      # (batch, seq, d_model)
A = torch.randn(64, 256)       # up-projection   d_model -> d_ff
B = torch.randn(256, 64)       # down-projection d_ff -> d_model

z_ref = single_device_mlp(x, A, B)                 # one device
for n in (1, 2, 4, 8):
    z_tp = megatron_mlp(x, A, B, num_devices=n)     # split across n devices
    rel = (z_tp - z_ref).abs().max() / z_ref.abs().max()
    print(f"N={n}: max relative diff = {rel:.2e}")
N=1: max relative diff = 0.00e+00
N=2: max relative diff = 3.71e-07
N=4: max relative diff = 4.24e-07
N=8: max relative diff = 4.77e-07

The split output equals the single-device MLP for every N — the only gap is the float-point reassociation of summing the partials in a different order, which is machine-epsilon tiny. Distribution is exact, not an approximation — the same anchor data parallelism earned earlier in this module.

Splitting Attention by Heads

Multi-head attention splits even more naturally: the heads are already independent. Give each device a disjoint set of heads and it computes those heads’ attention end-to-end with no communication — the Q/K/V projections are just a column-parallel split whose boundaries fall on head edges, so a head never straddles two devices. Only the output projection W_O is row-parallel, paying the same single all-reduce as the MLP.

import torch
from tensor_parallel import megatron_attention, single_device_attention

torch.manual_seed(0)
x = torch.randn(2, 8, 64)                       # (batch, seq, d_model)
Wq, Wk, Wv, Wo = (torch.randn(64, 64) / 8 for _ in range(4))  # 1/√d init, 8 heads of dim 8

o_ref = single_device_attention(x, Wq, Wk, Wv, Wo, num_heads=8)
o_tp = megatron_attention(x, Wq, Wk, Wv, Wo, num_heads=8, num_devices=4)  # 2 heads/device
print("attention split matches:", torch.allclose(o_ref, o_tp, atol=1e-5, rtol=1e-5))
attention split matches: True

A transformer layer is one attention block plus one MLP block, so it does two all-reduces in the forward pass (g for each) and two more in the backward pass (f for each) — four all-reduces per layer per step. That is the price of tensor parallelism, and it is why it lives inside a node.

The Cost: Memory Down, Communication Up

Unlike data parallelism, tensor parallelism shards the weights, so each device stores only 1/N of the split layers’ parameters — it is a genuine memory cut, not just a batch split. The bill it runs up is communication: those four all-reduces per layer fire on every step and sit on the critical path (a nonlinearity is waiting on the other side), so unlike the data-parallel gradient all-reduce they cannot be hidden behind other compute. Drive the trade-off:

TipTry This
  1. Push N up. Parameters-per-device keep falling as 1/N, but the communication per layer rises toward a fixed ceiling and never goes away. That crossover is why tensor parallelism is kept within a node (fast NVLink) and data/pipeline parallelism spans nodes.
  2. Shrink d_model. Communication scales with the activation (\text{tokens}\times d_\text{model}) while parameters scale with d_\text{model}^2 — so the bigger the layer, the better the memory-for-comm trade. Tensor parallelism pays off exactly where you need it: enormous layers.

Pipeline Parallelism: Split the Stack

Data parallelism splits the batch; tensor parallelism splits each layer’s matrices. The third axis splits the depth: put layers 0–7 on device 0, layers 8–15 on device 1, and so on. A forward pass then flows device 0 → 1 → … → P{-}1 (each stage sends its output activations to the next), and the backward pass flows back. This is pipeline parallelism, and it is the only axis that lets a model deeper than one device’s memory train at all — no single device ever holds more than its slice of the stack.

The catch is idle time. Feed one batch straight through a 4-stage pipeline and the picture is embarrassing: device 1 waits for device 0, device 2 waits for device 1, and while stage 3 finally runs the forward pass, stages 0–2 have nothing to do. Only one device works at a time — a 4\times cluster running at 1\times speed. That wasted wavefront is the pipeline bubble.

The fix (GPipe, Huang et al. 2018) is the assembly line. Split the mini-batch into M micro-batches and stream them through the stages: as soon as stage 0 finishes micro-batch 0 and hands it to stage 1, it starts micro-batch 1. Once the pipeline is full, every stage is busy on a different micro-batch, exactly like cars moving down a factory line where every station works a different car. The bubble never fully closes — someone is idle during the fill and the drain — but its share shrinks as you push more micro-batches through.

NoteKey Insight

Pipeline parallelism trades a fixed bubble of idle time for the ability to place layers on separate devices. Micro-batching hides the bubble by keeping every stage busy on a different micro-batch — the more micro-batches, the smaller the bubble’s share of the total. Unlike tensor parallelism’s per-layer all-reduce, the only communication is a point-to-point send of activations across each stage boundary, so pipeline parallelism happily spans the slow links between nodes.

The Math: The Bubble

Assume, as GPipe does, that a micro-batch’s forward and backward each cost one time slot. With P stages the pipeline needs P{-}1 slots to fill (stage s can’t start until the wavefront reaches it) and P{-}1 slots to drain at the end. During those 2(P{-}1) slots some device is idle. The useful work per device is M forwards +\,M backwards = 2M slots, so the total makespan is 2(M + P - 1) and the idle bubble fraction is

\text{bubble} \;=\; \frac{2(P-1)}{2(M + P - 1)} \;=\; \frac{P-1}{M + P - 1}.

This is GPipe’s O\!\big((P{-}1)/(M{+}P{-}1)\big). Megatron-LM reports the same quantity relative to the ideal compute time 2M, which drops the P{-}1 from the denominator: (P-1)/M. Both say the same thing — more micro-batches, less bubble — and GPipe’s rule of thumb is that the bubble is negligible once M \ge 4P. Drive it:

TipTry This
  1. Start at M = 1. The bubble is (P{-}1)/P — a P-stage pipeline running at nearly 1/P efficiency. That is the naive picture, and it is terrible.
  2. Slide M up to 4P (the blue line). The bubble drops below $$20% and keeps falling. This is why real runs use dozens of micro-batches per step.
  3. Raise P. More stages means a bigger bubble at the same M — deeper pipelines need proportionally more micro-batches to stay efficient.

Code: The Schedule from Scratch

The bubble formula falls out of an actual schedule — an assignment of each micro-batch’s forward and backward to a (device, time-slot). pipeline.py builds one from scratch with a small greedy list-scheduler that respects three data dependencies: a forward can’t run until the previous stage’s forward for that micro-batch is done; a backward starts at the last stage once its forward finishes; and a backward flows from stage P{-}1 back to stage 0.

from pipeline import partition_layers, pipeline_schedule, schedule_makespan, bubble_fraction

# 12 transformer blocks over 4 stages: each device owns a contiguous slice.
print("layer partition:", partition_layers(12, 4))

cells = pipeline_schedule(num_stages=4, num_microbatches=8, policy="gpipe")
print("ops scheduled:", len(cells), "  makespan:", schedule_makespan(cells))
print("bubble fraction:", round(bubble_fraction(4, 8), 4))
layer partition: [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]]
ops scheduled: 64   makespan: 22
bubble fraction: 0.2727

The same scheduler produces both classic schedules — the only difference is the tie-break when a device could run either a ready forward or a ready backward. GPipe prefers forwards (all forwards, then all backwards); 1F1B (one-forward-one-backward) prefers backwards, retiring each micro-batch’s activations as early as it can. Toggle between them and watch the grid — blue is a forward, orange a backward, and an empty cell is a bubble:

Notice the two schedules span the same number of columns — GPipe and 1F1B have the identical bubble. So why does everyone use 1F1B? Look at how long each stage must hold onto activations.

Same Bubble, Less Memory

The backward pass needs the activations saved during the forward pass. In GPipe, stage 0 runs all M forwards before the first backward reaches it — so it must stash the activations of all M micro-batches at once. In 1F1B, each stage starts backpropagating as soon as its warm-up forwards are done and immediately frees that micro-batch’s activations, so it never holds more than P in flight (one per stage ahead of it). Same bubble, but peak activation memory of O(M) vs O(P) — and since good utilization wants M \gg P, that difference is enormous. The scheduler measures it directly:

from pipeline import peak_inflight_microbatches, pipeline_schedule

for policy in ("gpipe", "1f1b"):
    peak = peak_inflight_microbatches(pipeline_schedule(num_stages=4, num_microbatches=16, policy=policy))
    print(f"{policy:5s}: peak micro-batches held at stage 0 = {peak}")
gpipe: peak micro-batches held at stage 0 = 16
1f1b : peak micro-batches held at stage 0 = 4

Exact by Construction

One worry should be nagging: does chopping the batch into micro-batches and running them through separate devices change the answer? No. Splitting the batch and summing the per-micro-batch gradients is exactly the gradient of the summed loss over the whole batch — pipeline parallelism is gradient accumulation along the depth axis. pipeline.py proves it against a single-device reference:

from pipeline import gradient_match

res = gradient_match(num_layers=6, width=8, batch=12, num_stages=3, num_microbatches=4)
print("max relative gradient difference:", f"{res['max_rel_diff']:.2e}")
print("bit-for-bit (to float reassociation):", res["allclose"])
max relative gradient difference: 1.04e-07
bit-for-bit (to float reassociation): True

The gradients agree to machine epsilon — the only gap is the order floating-point adds the micro-batch contributions. Distribution across depth, like distribution across the batch (data parallelism) and across a layer (tensor parallelism), is exact, not an approximation.

Sequence Parallelism: Ring Attention

Data, tensor, and pipeline parallelism all leave one thing whole on every device: the sequence. Each replica still holds the entire length-L activation for its examples. That is fine at L = 2\text{k}; it is fatal at L = 1\text{M}. Attention’s score matrix is L \times L, and even m09’s FlashAttention — which never materializes that matrix — still keeps the whole length-L Q, K, and V resident on one device. Past a few hundred thousand tokens, that alone overflows. None of the three axes above help: they shard the batch, the weights, or the layers, never the tokens.

Ring attention (Liu, Zaharia & Abbeel, 2023) adds the missing fourth axis: shard the sequence itself across P devices. Device p permanently owns its query block Q_p — one L/P-length slice. The key/value blocks are then passed around a ring: at each of P steps, every device holds exactly one K/V block, folds that block’s contribution into its own queries with an online softmax (the running-max / running-denominator trick from m09), then hands the block to its neighbor and receives the next. After P steps every device has seen every K/V block, so every query has attended over the whole sequence — and because the send of the next block overlaps the compute on the current one, the communication is almost free.

NoteKey Insight

Ring attention is the sequence twin of the three axes you just built. It is exact — the online-softmax accumulation is the same identity m09 proves, so the ring output equals full attention bit-for-bit. And it is unboundedly scalable: each device only ever holds O(L/P) of the sequence (its own block plus one in-flight block), so the maximum context length grows linearly with the number of devices, with no change to the math. That is the “near-infinite context” headline.

The Ring, Step by Step

Watch the K/V blocks hop around the ring. Each GPU keeps its own query block fixed (the label Q_p) and, at every step, absorbs whichever K/V block has rotated to it — the colored tile riding on each node tells you which one. The arc inside each node fills as it absorbs more blocks; after P steps it is full and that GPU’s slice of the output is done.

The Math: One Block at a Time

The whole trick is that softmax can be computed incrementally. Keep three running numbers per query row — the max logit m, the denominator \ell, and the unnormalized output o — and fold in one K/V block at a time. When a block with scores s_{ij} = q_i \cdot k_j / \sqrt{d_k} arrives:

m_i \leftarrow \max\!\big(m_i,\ \max_j s_{ij}\big), \qquad \ell_i \leftarrow \ell_i\,e^{\,m_i^{\text{old}} - m_i} + \textstyle\sum_j e^{\,s_{ij} - m_i}

o_i \leftarrow o_i\,e^{\,m_i^{\text{old}} - m_i} + \textstyle\sum_j e^{\,s_{ij} - m_i}\,v_j, \qquad\text{and at the end}\quad \text{out}_i = o_i / \ell_i .

The e^{\,m^{\text{old}} - m} factor rescales the mass already accumulated to the new maximum — the exact correction that makes streaming equal to computing the whole softmax at once (this is m09’s online-softmax identity). Ring attention runs this loop once per device, over the P blocks that rotate past it. Because the loop is order-independent, the blocks can arrive in any ring order and the answer is unchanged. absorb_block in ring_attention.py is exactly this update.

Exact by Construction

ring_attention simulates the P devices: it shards Q, K, V along the sequence, and for each device streams every rotated K/V block through absorb_block. The output is the same tensor a single device would compute — for both full and causal attention.

import torch
from ring_attention import ring_attention

torch.manual_seed(0)
q, k, v = (torch.randn(2, 4, 96, 32) for _ in range(3))   # (batch, heads, seq=96, d=32)
scale = 1.0 / q.size(-1) ** 0.5

# The single-device reference.
ref = torch.softmax(q @ k.transpose(-2, -1) * scale, dim=-1) @ v

# Split the length-96 sequence across 8 ring devices.
ring = ring_attention(q, k, v, num_devices=8)

print("max |ring − full|:", f"{(ring - ref).abs().max():.2e}")
print("exact (bit-for-bit to float reassociation):", torch.allclose(ring, ref, atol=1e-5))
max |ring − full|: 4.17e-07
exact (bit-for-bit to float reassociation): True

The gap is machine epsilon — the ring never approximates. The causal variant skips any K/V block strictly in the future of a device’s queries and masks the diagonal block, matching a causally-masked full attention just as exactly:

seq = q.size(-2)
tri = torch.triu(torch.ones(seq, seq, dtype=torch.bool), diagonal=1)
ref_causal = torch.softmax(
    (q @ k.transpose(-2, -1) * scale).masked_fill(tri, torch.finfo(q.dtype).min), dim=-1
) @ v
ring_causal = ring_attention(q, k, v, num_devices=8, causal=True)
print("causal exact:", torch.allclose(ring_causal, ref_causal, atol=1e-5))
causal exact: True

Context Scales with Devices

Because each device holds only its own block plus one in-flight block, the sequence-activation memory per device is O(L/P), not O(L) — and the L\times L score matrix is never formed. Invert that and the story flips into the headline: with a fixed per-device memory budget, the context length you can attend to grows linearly with the number of ring devices. Drive P and watch the wall move:

TipTry This
  1. Push P to 32. The attainable context grows in lockstep — 32 devices, 32× the context, and the per-device memory is unchanged. This linearity is the whole point: long context becomes a systems dial, not a modeling limit.
  2. Step the ring to the last step. Every GPU’s absorb-arc is full: each has now seen all P key blocks, so every query attended the entire sequence — yet no GPU ever held more than 2/P of the sequence at once.

The Causal Imbalance — and the Striped Fix

There is one honest wrinkle. Under a causal mask, a query at position i only attends to the i{+}1 keys before it. With contiguous shards, the device holding the last block of positions must attend to every prior block, while the device holding the first block attends to almost none — so the late devices do far more work, and the ring’s step time is set by the busiest one. causal_block_work counts the exact causal (q,k) pairs each device computes:

from ring_attention import causal_block_work

contig = causal_block_work(num_devices=8, block_size=16, layout="contiguous")
striped = causal_block_work(num_devices=8, block_size=16, layout="striped")
print("contiguous per-device work:", contig["work"])
print("  imbalance (max/min):", f"{contig['imbalance']:.2f}×")
print("striped    per-device work:", striped["work"])
print("  imbalance (max/min):", f"{striped['imbalance']:.2f}×")
print("same total work:", contig["total"] == striped["total"])
contiguous per-device work: [136, 392, 648, 904, 1160, 1416, 1672, 1928]
  imbalance (max/min): 14.18×
striped    per-device work: [976, 992, 1008, 1024, 1040, 1056, 1072, 1088]
  imbalance (max/min): 1.11×
same total work: True

The fix, from Striped Attention (Brandon et al., 2023), is to stop giving each device a contiguous run of tokens. Assign device p the interleaved positions \{p,\ p+P,\ p+2P,\dots\} instead. Now every device holds some early and some late positions, so under the same causal ring each does nearly identical work — the total number of causal pairs is unchanged, but it is spread evenly. (The same idea, paired up block-by-block, is the widely-used “zigzag” ring schedule.)

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.

WarningTensor parallelism all-reduces every layer — keep it intra-node

The data-parallel all-reduce fires once per step and hides behind the backward pass. Tensor parallelism’s four all-reduces fire per layer and block on the critical path (a GeLU or softmax waits on the sum). Over a slow inter-node link that stall dominates, so tensor parallelism is confined to a single node’s fast interconnect (NVLink); data and pipeline parallelism span the slower links between nodes. And never split below one head or one output column — an empty shard is a device doing nothing but communicating.

WarningToo few micro-batches wastes the pipeline

The bubble is (P{-}1)/(M{+}P{-}1): with M = 1 a P-stage pipeline runs at nearly 1/P efficiency. Micro-batching is not optional — it is pipeline parallelism. Aim for M \ge 4P (GPipe’s rule of thumb) so the bubble drops below $$20%. But micro-batches can’t shrink forever: each one is a smaller GEMM, and past a point the per-kernel overhead and lower arithmetic intensity eat the gain.

WarningGPipe stores every micro-batch’s activations — use 1F1B

The naive all-forward-then-all-backward schedule (GPipe) keeps stage 0 holding the activations of all M micro-batches until the backward sweep arrives — activation memory that grows with the very M you raised to shrink the bubble. 1F1B has the same bubble but bounds in-flight micro-batches to P, so it is the schedule real frameworks ship. (Interleaved 1F1B and activation recomputation/checkpointing cut the bill further still.)

WarningCausal ring attention is load-imbalanced — stripe the tokens

A contiguous causal ring makes the device holding the last query block attend to every prior block while the first device attends to almost none, so its step time is set by the busiest device (~2\times the average at large P). Striped / zigzag token assignment — interleaved positions instead of contiguous runs — spreads the same total causal work evenly and recovers the throughput. The non-causal ring has no such imbalance; this bites only under the causal mask.

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: True

Exercise 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.0

Exercise 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: True

Exercise 4: Build the row-parallel GEMM from the shards

row_parallel_forward splits B by rows, multiplies each partial, and all-reduces the sum. Reconstruct it from the primitives and confirm it matches the full Y B — the second half of the Megatron MLP.

import torch
from tensor_parallel import row_parallel_shards, ring_all_reduce

torch.manual_seed(0)
Y = torch.randn(3, 16)     # column-split activation (full, for the check)
B = torch.randn(16, 8)

def my_row_parallel(y, b, n):
    # Your implementation here: split b by rows AND y by columns to match,
    # form each partial y_i @ b_i, then all-reduce the partials.
    y_shards = list(y.chunk(n, dim=-1))
    b_shards = row_parallel_shards(b, n)
    partials = [ys @ bs for ys, bs in zip(y_shards, b_shards)]
    return ring_all_reduce(partials)[0]

got = my_row_parallel(Y, B, 4)
print("row-parallel matches full:", torch.allclose(got, Y @ B, atol=1e-5))
row-parallel matches full: True

Exercise 5: How many micro-batches to hit a target bubble?

The bubble is (P{-}1)/(M{+}P{-}1). Invert it: for a pipeline of P stages, find the smallest M whose bubble is at most a target fraction, then check it against microbatches_for_bubble.

import math
from pipeline import bubble_fraction, microbatches_for_bubble

def min_microbatches(p, target):
    # Your implementation: solve (p-1)/(M+p-1) <= target for the smallest integer M.
    return max(1, math.ceil((p - 1) * (1 - target) / target))

for p, target in [(4, 0.1), (8, 0.1), (16, 0.05)]:
    m = min_microbatches(p, target)
    print(f"P={p:2d}, target {target:>4}: M={m:3d} → bubble {bubble_fraction(p, m):.3f}",
          "✓" if m == microbatches_for_bubble(p, target) else "✗")
P= 4, target  0.1: M= 27 → bubble 0.100 ✓
P= 8, target  0.1: M= 63 → bubble 0.100 ✓
P=16, target 0.05: M=285 → bubble 0.050 ✓

Summary

Key takeaways:

  1. 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.
  2. 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.
  3. 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.
  4. 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.
  5. Tensor parallelism splits a single layer across devices: the MLP’s first GEMM column-parallel (so GeLU stays independent) and its second row-parallel, reassembled by one all-reduce; attention splits the same way, by heads. The result is bit-exact to one device (bar reassociation) — it shards the weights 1/N at the cost of four all-reduces per layer, which is why it stays inside a node.
  6. Pipeline parallelism splits the stack of layers into stages and streams M micro-batches through them to hide the fill/drain bubble, which shrinks as (P{-}1)/(M{+}P{-}1) (negligible at M \ge 4P). GPipe and 1F1B have the same bubble, but 1F1B bounds peak activation memory to O(P) instead of O(M) — and, like the other axes, pipelining is exact, since micro-batch gradient accumulation equals the full-batch gradient.
  7. The axes compose: data parallelism splits the batch, tensor parallelism splits each layer, pipeline parallelism splits the stack, and gradient checkpointing trades compute for activation memory. Frontier runs combine all of them (3D parallelism).

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 on all four axes — across the batch (data parallelism, ZeRO), within a layer (tensor parallelism), along the depth (pipeline parallelism), and along the sequence (ring attention / context parallelism). The frontier from here: interleaved 1F1B and zero-bubble pipeline schedules, FSDP wired into a real GPTModel, gradient checkpointing built from an autograd hook, and the all-gather context-parallel variant Llama 3 used at 128k — all filed on the roadmap.

Going Deeper

Core Papers:

Practical Resources:

  • DeepSpeed — the library that popularized ZeRO.
  • Megatron-LM (code) — NVIDIA’s reference implementation of tensor + pipeline parallelism.