Data parallelism, ZeRO and FSDP

Before this lesson: lesson 6 (16 bytes per parameter, activation memory), lesson 7 (ring all-reduce = reduce-scatter + all-gather, and the bytes each GPU sends), stochastic gradient descent (the mini-batch gradient is a mean) and Adam and AdamW (the update is elementwise). Time: about 70 minutes, plus the notebook.

---

Why this matters #

Data parallelism (DP) is the outer axis of almost every training run, from a small model on 4 GPUs to Llama 3 405B, which used data parallelism of degree 128 on 16,384 H100 GPUs (with tensor and pipeline parallelism inside). This lesson gives the numbers behind the options of DDP, ZeRO and FSDP: how much memory each stage saves, how many bytes it sends, and when those bytes hide behind compute.

A 7B model trained with mixed-precision AdamW needs $16 \times 7 = 112$ GB of model states (lesson 6). No 80 GB GPU holds that. ZeRO-3 over 8 GPUs needs 14 GB per GPU. The price is 1.5 times the communication of plain data parallelism. Whether you see that price in the step time depends on one ratio: FLOP/s over bandwidth.

1. Predict first #

2. The picture #

Two GPUs, one batch #

Both GPUs hold the same weights $w$. GPU 0 takes the first half of the batch, GPU 1 the second half. Each computes the gradient of the mean loss over its half. The loss is a mean over samples and the gradient is linear, so the mean of the two half-batch gradients is the full-batch gradient (when the halves have the same size). An all-reduce gives this mean to both GPUs. Both apply the same update, and the two copies stay identical. This is DDP (distributed data parallel).

Small numbers. Fit $\hat{y} = w x$ with the loss $\frac{1}{4}\sum_i \frac{1}{2}(w x_i - y_i)^2$ on four samples $x = (1, 2, 3, 4)$, $y = (2, 3, 5, 9)$, at $w = 1$. The gradient of one sample is $(w x_i - y_i)\,x_i$:

sample	1	2	3	4
$w x_i - y_i$	$-1$	$-1$	$-2$	$-5$
$(w x_i - y_i)\,x_i$	$-1$	$-2$	$-6$	$-20$

GPU 0 (samples 1, 2) gets the mean $-1.5$; GPU 1 (samples 3, 4) gets $-13$. The all-reduce averages them: $(-1.5 - 13)/2 = -7.25$. One GPU with all four samples gets $(-1 - 2 - 6 - 20)/4 = -7.25$. The same number.

What is redundant #

With $N_d$ GPUs and $\Psi$ parameters, every GPU stores the same weights, gradients and optimizer states: $N_d$ copies of $16\Psi$ bytes. ZeRO removes the copies. The key fact is that AdamW is elementwise: the update of parameter $j$ needs only its gradient $g_j$, its moments $m_j$ and $v_j$, and its fp32 master weight. So GPU $k$ can own slice $k$ of the parameters. It keeps the optimizer states of that slice only, receives only that slice of the averaged gradient, updates only that slice, and sends the new values to the others.

The communication is already there #

A ring all-reduce has two halves (lesson 7). After the reduce-scatter, GPU $k$ holds slice $k$ of the sum. After the all-gather, every GPU holds every slice. ZeRO puts the optimizer step between the two halves:

	before the update	the update	after the update
DDP	backward, then all-reduce of the gradients (= reduce-scatter + all-gather)	every GPU updates all $\Psi$	—
ZeRO-1, ZeRO-2	backward, then reduce-scatter of the gradients	GPU $k$ updates slice $k$	all-gather of the new weights
ZeRO-3	all-gather the weights (forward), all-gather again (backward), reduce-scatter of the gradients	GPU $k$ updates slice $k$	nothing: the next forward pass gathers

Interactive figure: loading…

3. Definitions and main results #

Notation. $\Psi$: number of parameters. $N_d$: number of data-parallel GPUs (lesson 7 writes $p$). $K$: optimizer bytes per parameter; $K = 12$ for mixed-precision Adam (fp32 master weight, $m$ and $v$, 4 bytes each). Weights and gradients are bf16 (2 bytes). $T$: tokens one GPU processes per optimizer step. $F$: FLOP/s one GPU achieves on the matmuls. $W$: network bandwidth per GPU and direction, in bytes/s. A ring reduce-scatter or all-gather of $n$ bytes sends $\frac{N_d-1}{N_d}\,n$ bytes per GPU; a ring all-reduce sends twice that (lesson 7). 1 GB $= 10^9$ bytes.

Data parallelism is large-batch training #

Proposition (DDP is exact). Let all GPUs start from the same weights (DDP broadcasts them from GPU 0 at the start). Split the global batch into $N_d$ micro-batches of the same size, and let the loss be a mean over samples. Then after every step all GPUs hold the same weights, equal to the weights of one GPU that trains on the whole batch.

Proof.

1. The full-batch loss is $\mathcal{L}(w) = \frac{1}{B}\sum_{i=1}^{B}\ell_i(w)$. Split the $B$ samples into $N_d$ groups $G_k$ of size $B/N_d$.
2. GPU $k$ computes $g_k = \frac{N_d}{B}\sum_{i \in G_k}\nabla\ell_i(w)$, the gradient of its own mean loss.
3. The averaged gradient is $\frac{1}{N_d}\sum_k g_k = \frac{1}{B}\sum_{i=1}^{B}\nabla\ell_i(w) = \nabla\mathcal{L}(w)$.
4. Every GPU applies the same deterministic optimizer step to the same weights, the same gradient and the same optimizer state, so all GPUs end with the same new weights and state.
5. By induction over the steps, the copies stay identical and equal to one-GPU training. $\blacksquare$

In floating point the averaged gradient can differ from the one-GPU gradient in the last bits, because the sum is done in a different order. The copies still stay identical to each other, because all GPUs receive the same all-reduce result.

Theorem (ZeRO memory and communication) #

With one micro-batch per step:

stage	what is sharded	model states per GPU (bytes)	elements sent per GPU per step
DDP (ZeRO-0)	nothing	$(4+K)\Psi$	$2\Psi$
ZeRO-1 ($P_{os}$)	optimizer states	$4\Psi + K\Psi/N_d$	$2\Psi$
ZeRO-2 ($P_{os+g}$)	+ gradients	$2\Psi + (2+K)\Psi/N_d$	$2\Psi$
ZeRO-3 ($P_{os+g+p}$)	+ parameters	$(4+K)\Psi/N_d$	$3\Psi$

To get bytes, multiply the volumes by the bytes per element (2 for bf16) and by $(N_d-1)/N_d$.

Proof.

1. Bytes per parameter. bf16 weight (2), bf16 gradient (2), optimizer states ($K$). DDP keeps everything on every GPU: $(4+K)\Psi = 16\Psi$.
2. ZeRO-1. AdamW is elementwise, so GPU $k$ needs the optimizer states of its slice only: $K\Psi/N_d$. Weights and gradients stay full: $4\Psi + K\Psi/N_d$.
3. ZeRO-2. GPU $k$ uses only slice $k$ of the averaged gradient. So each gradient bucket is reduced to its owners as soon as it is ready (a reduce-scatter), and the rest of the bucket is freed: $2\Psi + (2+K)\Psi/N_d$.
4. ZeRO-3. The weights are sharded too: $(4+K)\Psi/N_d$. Full weights exist only for the few layers that are being computed.
5. DDP volume. An all-reduce of $\Psi$ gradient elements is a reduce-scatter ($\Psi$) plus an all-gather ($\Psi$): $2\Psi$.
6. ZeRO-1 and ZeRO-2 volume. Reduce-scatter the gradients ($\Psi$); GPU $k$ updates slice $k$; all-gather the new weights ($\Psi$). Total $2\Psi$. This is the DDP all-reduce cut in two, with the update in the middle, so there is no extra traffic.
7. ZeRO-3 volume. Before each layer's forward pass, all-gather its weights, and free them after use: $\Psi$ over the whole model. The backward pass needs them again: a second all-gather ($\Psi$). Then reduce-scatter the gradients ($\Psi$). No all-gather follows the update, because the next forward pass does it. Total $3\Psi$, which is 1.5 times DDP. $\blacksquare$

Two remarks.

• For large $N_d$ the model states go to $4\Psi$ (ZeRO-1, 4 times less than DDP), $2\Psi$ (ZeRO-2, 8 times less) and 0 (ZeRO-3).
• ZeRO shards model states only. Every GPU still runs whole layers on its own micro-batch, so the activations are not divided. For them use checkpointing (lesson 6) or tensor and sequence parallelism (lesson 9).

Gradient accumulation #

Gradient accumulation runs $m$ micro-batches one after the other and adds their gradients before one optimizer step. The global batch grows by $m$ without more activation memory, because only one micro-batch's activations are stored at a time. The communication depends on the stage:

stage	what it does per micro-batch	elements sent per GPU per step
DDP	adds into its full local gradient; no communication inside no_sync()	$2\Psi$
ZeRO-1	the same: full local gradient, one reduce-scatter and one all-gather per step	$2\Psi$
ZeRO-2	keeps no full gradient, so it reduce-scatters after every micro-batch	$m\Psi + \Psi$
ZeRO-3	all-gathers the weights twice and reduce-scatters in every micro-batch	$3m\Psi$

So with accumulation, DDP and ZeRO-1 send fewer bytes per token. ZeRO-3 pays its full $3\Psi$ for every micro-batch.

How DDP overlaps the all-reduce with the backward pass #

PyTorch DDP (Li et al. 2020) works in five steps.

1. Construction. GPU 0 broadcasts the parameters and buffers. DDP groups the parameters into buckets (default 25 MiB) in the reverse order of model.parameters(). This is a guess of the order in which the gradients become ready. After the first iteration, DDP rebuilds the buckets in the order it actually saw.
2. Backward. An autograd hook fires when a parameter's gradient is ready and marks it in its bucket. When a bucket is complete, an asynchronous all-reduce of the whole bucket starts on a separate communication stream. Meanwhile autograd computes the gradients of earlier layers. All GPUs launch the buckets in the same order, because collectives must match across GPUs.
3. End of backward. DDP waits for all buckets and writes the averaged gradients into .grad. The optimizer then runs as on one GPU.
4. Why buckets. One all-reduce per tensor sends many small messages, and each pays the latency $\alpha$ (lesson 7). One all-reduce at the end overlaps nothing. Buckets of tens of MB make $\alpha$ small next to the bandwidth term and still start early.
5. What stays exposed. The bucket with the first layers' gradients is ready last. No compute is left to hide it.

Interactive figure: loading…

How FSDP implements ZeRO-3 #

PyTorch FSDP (Zhao et al. 2023) works in six steps.

1. Units. You wrap the model bottom-up, usually one unit per transformer block, plus the root (fully_shard(block) in FSDP2). The parameters of one unit form one communication group. FSDP2 shards each parameter along dimension 0 as a DTensor; FSDP1 flattens a unit's parameters into one FlatParameter split into $N_d$ equal chunks. Each GPU keeps its fp32 shard (the master weights), its gradient shard and the Adam states of its shard. With bf16 compute the shard and its gradient are both fp32: $4 + 4 + 8 = 16$ bytes per parameter, the same $16\Psi/N_d$ as $(4+K)\Psi/N_d$.
2. Forward of a unit. A pre-forward hook all-gathers the unit's parameters (in bf16 if the mixed-precision policy says so). The unit runs. A post-forward hook frees the full copy (reshard_after_forward=True).
3. Backward of a unit (in reverse order). All-gather the parameters again, compute the gradients, free the full parameters, and reduce-scatter the gradients. Each GPU keeps its gradient shard.
4. Optimizer step. A normal optimizer runs on the local shards, with no communication. The next forward pass's all-gathers spread the new weights.
5. Overlap and prefetch. The all-gathers run on a separate stream. Backward prefetch issues the next unit's all-gather before the current unit's reduce-scatter. Forward prefetch issues the next forward all-gather early. FSDP limits how many all-gathers are in flight, so the peak memory is the shards, plus about two units of full parameters, plus the activations.
6. Variants. reshard_after_forward=False keeps the full weights from forward to backward: one all-gather less ($2\Psi$ volume), but the full bf16 weights stay in memory during the step. HSDP (hybrid sharding; FSDP2 with a 2-D device mesh) shards inside groups of $G$ GPUs, usually one node, and replicates across the $R = N_d/G$ groups. The model states are $(4+K)\Psi/G$. The all-gathers and reduce-scatters stay on the fast links inside the node. Only an all-reduce of each GPU's gradient shard ($\Psi/G$ elements) crosses the network.

Interactive figure: loading…

Proposition (tokens that hide the all-reduce) #

Use DDP with bf16 gradients and a ring all-reduce over $N_d$ GPUs, overlapped with the backward pass. The communication is hidden when

Proof.

1. The forward pass costs about $2\Psi$ FLOPs per token and the backward pass $4\Psi$ (one matmul for the input gradient, one for the weight gradient). We ignore the attention-score FLOPs. The backward time is $4\Psi T/F$.
2. The all-reduce of $\Psi$ bf16 gradients ($2\Psi$ bytes) sends $2\cdot\frac{N_d-1}{N_d}\cdot 2\Psi$ bytes per GPU. Its time is $4\Psi\frac{N_d-1}{N_d}/W$.
3. It is hidden if $4\Psi\frac{N_d-1}{N_d}/W \le 4\Psi T/F$, that is, $T \ge \frac{F}{W}\cdot\frac{N_d-1}{N_d}$. $\Psi$ cancels: compute and communication both grow linearly with the model. $\blacksquare$

• ZeRO-3. Forward: all-gather $2\Psi$ bytes against $2\Psi T$ FLOPs. Backward: all-gather and reduce-scatter, $4\Psi$ bytes against $4\Psi T$ FLOPs. The same threshold in both phases, but per micro-batch, because the gathers repeat in every micro-batch.
• fp32 gradient reduction (Llama 3 reduce-scatters gradients in fp32) doubles the bytes: $T \ge 2F/W$.
• Numbers. Llama 3 405B achieved $F = 400$ TFLOP/s per H100, about 40% of the 989 TFLOP/s dense bf16 peak. Across nodes with 400 Gb/s InfiniBand per GPU ($W = 50$ GB/s): $T \ge 8{,}000$ tokens. Inside a node on NVLink ($W = 450$ GB/s): about 900 tokens. With one 200 Gb/s link per GPU ($W = 25$ GB/s): 16,000 tokens.
• Use the bus bandwidth that nccl-tests measures (lesson 7) for $W$; it is often below the link peak. And remember that the last bucket is always exposed.

Interactive figure: loading…

Answers to section 1. (1) The same. A reduce-scatter plus an all-gather is exactly the all-reduce that DDP sends: $2\Psi$ elements. (2) Neither. Both need $T \ge F/W$, because $\Psi$ cancels. (3) 1.5 times: $3\Psi$ against $2\Psi$ elements, for any $N_d$. The bytes per GPU of a ring collective do not grow with the number of GPUs.

4. Worked example #

A 13B model (40 layers, hidden size $h = 5120$, sequence length $s = 4096$), mixed-precision AdamW ($K = 12$), on 4 nodes with 8 H100 80 GB each, so $N_d = 32$. NVLink gives 450 GB/s and InfiniBand 50 GB/s per GPU; the matmuls achieve $F = 400$ TFLOP/s. One micro-batch is one sequence: 4,096 tokens.

Step 1: model states per GPU ($\Psi = 13 \times 10^9$).

	DDP	ZeRO-1	ZeRO-2	ZeRO-3	HSDP ($G = 8$)
formula	$16\Psi$	$4\Psi + 12\Psi/32$	$2\Psi + 14\Psi/32$	$16\Psi/32$	$16\Psi/8$
GB	208	56.9	31.7	6.5	26

Step 2: activations (lesson 6, FlashAttention, no checkpointing): $34sbh = 34 \cdot 4096 \cdot 5120 \approx 0.71$ GB per layer, 28.5 GB for 40 layers. Totals: ZeRO-1 85.4 GB (does not fit), ZeRO-2 60.2 GB, ZeRO-3 35.0 GB, HSDP 54.5 GB.

Step 3: compute per micro-batch. $6\Psi \cdot 4096 = 3.19 \times 10^{14}$ FLOPs, so 0.80 s at 400 TFLOP/s, of which the backward pass takes $4\Psi \cdot 4096/F = 0.53$ s.

Step 4: communication per micro-batch (bf16).

• ZeRO-3 over all 32 GPUs: $3 \cdot 2\Psi \cdot \frac{31}{32} = 75.6$ GB per GPU. The rings cross nodes, so the 50 GB/s links set the pace: 1.51 s, more than the 0.80 s of compute. 4,096 tokens is below $F/W = 8{,}000$: the step is limited by communication.
• ZeRO-2 over 32 GPUs: reduce-scatter $2\Psi \cdot \frac{31}{32} = 25.2$ GB per micro-batch, 0.50 s, just under the 0.53 s backward pass. Plus one all-gather of 25.2 GB (0.50 s) per step.
• HSDP, shards inside each node: $3 \cdot 2\Psi \cdot \frac{7}{8} = 68.3$ GB on NVLink: 0.15 s. Across nodes, each GPU all-reduces its gradient shard ($2\Psi/8 = 3.25$ GB) with its 3 peers in the other nodes: $2 \cdot \frac{3}{4} \cdot 3.25 = 4.9$ GB, 0.10 s.

Verdict. HSDP with $G = 8$: 54.5 GB per GPU plus buffers, no recomputation, and all communication well hidden. ZeRO-3 over 32 GPUs would need 2 sequences per micro-batch to reach $F/W$. Then its activations double to 57 GB (63.5 GB in total), and its communication (1.51 s against 1.60 s of compute) is only just hidden. The same bytes are cheap on NVLink and expensive on the network.

5. In code #

The first block simulates ZeRO-1 on 4 "GPUs" with NumPy and checks two claims: the result equals one-GPU AdamW on the whole batch, and each GPU sends $2\frac{N_d-1}{N_d}\Psi$ elements per step.

import numpy as np

rng = np.random.default_rng(0)
N, P, B = 4, 10, 32                        # GPUs, parameters, global batch
X, y = rng.standard_normal((B, P)), rng.standard_normal(B)
grad = lambda w, Xs, ys: Xs.T @ (Xs @ w - ys) / len(ys)   # gradient of the mean of (x.w - y)^2 / 2

def adamw(w, g, m, v, t, lr=1e-2, b1=0.9, b2=0.999, eps=1e-8, wd=0.1):
    m, v = b1 * m + (1 - b1) * g, b2 * v + (1 - b2) * g * g
    mh, vh = m / (1 - b1**t), v / (1 - b2**t)
    return w - lr * (mh / (np.sqrt(vh) + eps) + wd * w), m, v

w0 = rng.standard_normal(P)
w, m, v = w0.copy(), np.zeros(P), np.zeros(P)        # one GPU, whole batch
for t in range(1, 6):
    w, m, v = adamw(w, grad(w, X, y), m, v, t)

sl = np.array_split(np.arange(P), N)                 # GPU k owns parameters sl[k]
Xs, ys = np.split(X, N), np.split(y, N)              # GPU k gets micro-batch k
W = [w0.copy() for _ in range(N)]                    # every GPU: all weights ...
M = [np.zeros(len(s)) for s in sl]                   # ... but Adam states of its slice only
V = [np.zeros(len(s)) for s in sl]
sent = 0
for t in range(1, 6):
    g_local = [grad(W[k], Xs[k], ys[k]) for k in range(N)]
    for k in range(N):                               # reduce-scatter: GPU k gets the mean of slice k
        gk = np.mean([g[sl[k]] for g in g_local], axis=0)
        W[k][sl[k]], M[k], V[k] = adamw(W[k][sl[k]], gk, M[k], V[k], t)
    for k in range(N):                               # all-gather: everyone gets every new slice
        for j in range(N):
            W[k][sl[j]] = W[j][sl[j]]
    sent += 2 * (N - 1) / N * P                      # elements sent per GPU: reduce-scatter + all-gather
print(max(np.abs(Wk - w).max() for Wk in W) < 1e-12)  # True: ZeRO-1 = one-GPU AdamW
print(sent / 5)                                       # 15.0 = 2 * (3/4) * 10 per step

The second block is PyTorch reading code for gradient accumulation with DDP. Inside no_sync() the hooks add the gradients locally and start no all-reduce; the last micro-batch all-reduces the sum once.

import contextlib
from torch.nn.parallel import DistributedDataParallel as DDP

ddp = DDP(model, bucket_cap_mb=25)            # gradients are grouped into 25 MiB buckets
for i, (x, y) in enumerate(micro_batches):    # m micro-batches per optimizer step
    sync = i == m - 1
    with contextlib.nullcontext() if sync else ddp.no_sync():
        (loss_fn(ddp(x), y) / m).backward()   # all-reduce of each bucket only when sync
opt.step()
opt.zero_grad()

The third block runs FSDP2 (fully_shard) on the CPU with two processes and the gloo backend. It checks three claims: the parameters are sharded, each unit needs two all-gathers and one reduce-scatter, and the result equals one process with the whole batch.

# Save as fsdp_check.py and run: python fsdp_check.py   (CPU, 2 processes)
import os
import warnings
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard
from torch.distributed.tensor.debug import CommDebugMode

warnings.filterwarnings("ignore", message="FSDP2-wrapped module")   # harmless under CommDebugMode

def make_model():
    torch.manual_seed(0)                            # same initial weights on every process
    return torch.nn.Sequential(torch.nn.Linear(16, 64), torch.nn.Tanh(), torch.nn.Linear(64, 4)).double()

def worker(rank, world):
    os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT="29514")
    dist.init_process_group("gloo", rank=rank, world_size=world)
    g = torch.Generator().manual_seed(1)
    X = torch.randn(8, 16, generator=g, dtype=torch.float64)
    Y = torch.randn(8, 4, generator=g, dtype=torch.float64)

    mesh = init_device_mesh("cpu", (world,))        # 1-D mesh: shard over all ranks (2-D gives HSDP)
    model = make_model()
    for layer in (model[0], model[2]):              # one FSDP unit per Linear ...
        fully_shard(layer, mesh=mesh)
    fully_shard(model, mesh=mesh)                   # ... and the root for anything left
    opt = torch.optim.AdamW(model.parameters(), lr=1e-2)   # sees only the local shards
    w = model[0].weight                             # a DTensor; rank r keeps rows 32r ... 32r + 31
    if rank == 0:
        print(type(w).__name__, w.placements, tuple(w.shape), tuple(w.to_local().shape))

    Xr, Yr = X.chunk(world)[rank], Y.chunk(world)[rank]     # each rank: its own 4 samples
    with CommDebugMode() as comm:
        torch.nn.functional.mse_loss(model(Xr), Yr).backward()
    opt.step()                                      # each rank updates its own shard
    full = [p.full_tensor() for p in model.parameters()]    # all-gather, only to compare

    if rank == 0:
        print({str(op).split(".")[-1]: n for op, n in comm.get_comm_counts().items()})
        ref = make_model()                          # reference: one process, the whole batch
        ref_opt = torch.optim.AdamW(ref.parameters(), lr=1e-2)
        torch.nn.functional.mse_loss(ref(X), Y).backward()
        ref_opt.step()
        print(all(torch.allclose(a, b) for a, b in zip(full, ref.parameters())))
    dist.destroy_process_group()

if __name__ == "__main__":
    mp.spawn(worker, args=(2,), nprocs=2)
# DTensor (Shard(dim=0),) (64, 16) (32, 16)
# {'_allgather_base_': 4, '_reduce_scatter_base_': 2}
# True

Two units, each gathered in the forward pass and again in the backward pass (4 all-gathers), and one reduce-scatter per unit: the $3\Psi$ of the theorem. With reshard_after_forward=False on the second unit you would see one all-gather less. (Checked with PyTorch 2.15 on a laptop CPU.)

6. Where it shows up #

• LLM pretraining. Llama 3 405B used FSDP as its data-parallel axis (degree 128 on 16,384 H100s), with tensor and pipeline parallelism inside. It sharded optimizer states and gradients but did not reshard the weights after the forward pass (so $2\Psi$ per step, like ZeRO-2), and it reduce-scattered gradients in fp32 (Llama 3 paper, section 3.3.2). DeepSpeed implements ZeRO-1/2/3 directly.
• Fine-tuning. FSDP with LoRA or QLoRA: the frozen base model is sharded, and only the small adapters have gradients and Adam states. So the gradient and $K\Psi$ terms cover only the adapters.
• ML for science. Earth-system and weather models often have 0.1–5B parameters but long token sequences: a 0.25° global field ($721 \times 1440$ points) with $4 \times 4$ patches gives about 65,000 tokens per time step. ZeRO-1 or ZeRO-2 fits the model states easily (plain DDP only up to about 3B on 80 GB GPUs, because $16 \cdot 5 = 80$ GB already fills one). One sample already gives $T$ about 4 times $F/W$ on a network with 25 GB/s per GPU (65,000 against 16,000 tokens), so the data-parallel communication hides well. The GPU fills with activations, and the limit is the largest useful global batch, not the bandwidth. On nodes with 4 GPUs and a slower network between nodes, shard inside the node: HSDP with $G = 4$.
• Checkpoints. Sharded optimizer states mean sharded checkpoint files (torch.distributed.checkpoint). A restart on a different $N_d$ needs resharding.

7. Common mistakes #

• "ZeRO-3 shards activations." ZeRO-1/2/3 shard model states only.
• "ZeRO-3 costs $N_d$ times the communication." It is $3\Psi$ against $2\Psi$ elements per GPU.
• Using the per-step token count for ZeRO-3. With gradient accumulation its gathers repeat in every micro-batch, so the condition $T \ge F/W$ must hold per micro-batch.
• "DDP traffic per GPU grows with $N_d$." The bytes per GPU stay at about $2n$ for a message of $n$ bytes. What grows is the latency term $2(N_d-1)\alpha$, so small buckets on many GPUs become limited by latency.
• Wrapping only the root with FSDP. Then the whole model is one unit. It is all-gathered at once, the full bf16 weights and gradients (about $4\Psi$ bytes) stay in memory through the backward pass, and nothing overlaps.
• Sum instead of mean, or uneven token counts. With a per-token loss and a different number of tokens on each GPU, weight each GPU's gradient by its token count.
• Local gradient clipping. With sharded gradients, the global norm needs an all-reduce of the squared shard norms. FSDP's clipping does it; a hand-written loop must do it too.

8. Check yourself #

1. A 70B model, mixed-precision AdamW, 128 GPUs of 80 GB. Model states per GPU with ZeRO-1, ZeRO-2 and ZeRO-3? Which ones fit?
   Answer

   ZeRO-1: $4\Psi + 12\Psi/128 = 280 + 6.6 \approx 286.6$ GB. ZeRO-2: $2\Psi + 14\Psi/128 = 140 + 7.7 \approx 147.7$ GB. ZeRO-3: $16\Psi/128 = 8.75$ GB. Only ZeRO-3 fits (or ZeRO combined with tensor or pipeline parallelism); the activations come on top.
2. Why can ZeRO-1 shard the Adam states at no extra communication cost? Which kind of optimizer would break this?
   Answer

   Adam is elementwise, so slice $k$ of the update needs only slice $k$ of the averaged gradient (the reduce-scatter half of the all-reduce) and of the states; the all-gather half then spreads the new weights. An optimizer that needs whole matrices, such as Shampoo or Muon (optimization lesson 12), cannot update a flat slice alone: it must first gather each gradient matrix, which adds communication.
3. A cluster of A100 GPUs achieves 150 TFLOP/s per GPU and has 200 Gb/s InfiniBand per GPU. How many tokens per GPU per step hide the DDP all-reduce? And with fp32 gradient reduction?
   Answer

   $W = 200/8 = 25$ GB/s, so $T \ge F/W = 150 \times 10^{12} / 25 \times 10^9 = 6{,}000$ tokens. With fp32 the bytes double: 12,000 tokens.
4. You set reshard_after_forward=False on every FSDP unit (what Llama 3 did). What changes in memory and communication?
   Answer

   The full bf16 weights stay in memory from forward to backward, up to $2\Psi$ bytes at the end of the forward pass. The backward all-gather disappears, so the volume drops from $3\Psi$ to $2\Psi$ (one all-gather and one reduce-scatter), the same as DDP.
5. HSDP on 16 nodes × 4 GPUs, $\Psi = 2$B, bf16 gradients, 25 GB/s per GPU between nodes. How many bytes does each GPU send across nodes per step, and how long does it take? Compare with DDP over all 64 GPUs.
   Answer

   Each GPU holds a gradient shard of $2\Psi/4 = 1$ GB and all-reduces it over $R = 16$ replicas: $2 \cdot \frac{15}{16} \cdot 1 = 1.875$ GB, 75 ms. DDP over 64 GPUs sends $2 \cdot \frac{63}{64} \cdot 4$ GB $= 7.875$ GB through the same 25 GB/s links: 315 ms.
6. In DDP, which gradients are always exposed, and why?
   Answer

   Those of the first layers (and the input embedding). The backward pass computes them last, so their bucket's all-reduce starts when no compute is left to overlap with.

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

1. Why does data-parallel training give the same update as one GPU with the whole batch?

   Show answer

2. What are the per-GPU model-state bytes of DDP, ZeRO-1, ZeRO-2 and ZeRO-3?

   Show answer

3. Why does ZeRO-2 send no more bytes than DDP?

   Show answer

4. Why does ZeRO-3 send $3\Psi$ elements per step?

   Show answer

5. What does FSDP do before a unit's forward pass and after its backward pass?

   Show answer

6. How many tokens per GPU per step hide the DDP gradient all-reduce?

   Show answer

7. Why does the model size cancel in the overlap condition?

   Show answer

8. What is HSDP?

   Show answer

9. How does gradient accumulation change the communication of DDP and of ZeRO-3?

   Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Why does ZeRO-2 cost no more communication than DDP?" — A ring all-reduce is a reduce-scatter plus an all-gather; ZeRO-2 does the update between the two halves.
• "Why is ZeRO-3 only 1.5 times?" — Two all-gathers of the weights and one reduce-scatter: $3\Psi$ against $2\Psi$.
• "Your FSDP run on 8 nodes is slow. What do you check?" — Tokens per micro-batch against $F/W$; exposed all-gathers in the profiler; the unit size (too small: latency; too large: memory); HSDP; reshard_after_forward=False; prefetching; the measured bus bandwidth.
• "Does the gradient all-reduce get harder to hide for bigger models?" — No: $T \ge F/W$, the model size cancels.
• More practice questions on distributed training are in the interview track.

10. Practice #

The notebook for this lesson runs in Colab, with NumPy only. It has five exercises, each with a check cell and a folded solution:

1. Simulate DDP with 4 "GPUs" on linear regression: show that the copies stay identical and match one-GPU SGD on the whole batch.
2. Sharded AdamW: implement ZeRO-1, -2 and -3 on the simulator; check that the final weights equal DDP's to $10^{-12}$ and count the elements each GPU sends.
3. model_state_bytes(psi, nd, stage, K): the memory formulas, checked against what the simulator stores.
4. A bucketed-overlap timeline: the exposed communication time as a function of the bucket size.
5. tokens_to_hide(F, W, nd, grad_bytes), and the smallest micro-batch for ZeRO-3 on a given network.

11. Watch and read #

Training Models Too Big for One GPU | ZeRO & FSDPJia-Bin Huang · 23 min

Clear drawings of data parallelism, the three ZeRO stages and the FSDP schedule.

Turing-NLG, DeepSpeed and the ZeRO optimizerYannic Kilcher · 21 min

A walk through the ZeRO paper: where the 16 bytes per parameter go and how each stage removes copies.

PyTorch Distributed Data Parallel (DDP) | PyTorch Developer Day 2020PyTorch Foundation · 10 min

How DDP builds buckets and overlaps the all-reduce with the backward pass.

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

A lecture on data parallelism, its communication, and ZeRO, with diagrams.

• Rajbhandari et al. (2019), ZeRO: Memory Optimizations Toward Training Trillion Parameter Models: section 3.1 (16 bytes per parameter), section 5 (the three stages), section 7 (communication volume).
• Zhao et al. (2023), PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel: full and hybrid sharding, overlap, prefetch, the rate limiter.
• Li et al. (2020), PyTorch Distributed: Experiences on Accelerating Data Parallel Training: how DDP builds buckets and overlaps communication.
• Hugging Face, The Ultra-Scale Playbook: data parallelism and ZeRO with measured numbers; picotron has naive and bucketed data parallelism in a few hundred lines.
• Austin et al., How to Scale Your Model, How to Parallelize a Transformer for Training: sections on data parallelism and fully sharded data parallelism.
• Llama Team (2024), The Llama 3 Herd of Models, section 3.3.2: how FSDP is used at 16,384 GPUs.