Training models too big for one GPU: ZeRO & FSDP

The video starts with how a computer stores a real number and ends with a model whose weights live in pieces across many GPUs. This guide follows the same path: number formats, mixed precision, the memory bill of Adam, and how ZeRO and FSDP remove the copies that plain data parallelism keeps on every GPU.

Video by Jia-Bin Huang · Watch on YouTube

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. Watch

Chapter guide

Timestamps come from the video's description. The right-hand column is this course's reading guide, written from the chapter titles; it doesn't describe what the video shows frame by frame.

2. The ideas, explained

Adam, and why it's expensive

Adam (Kingma and Ba, 2014) keeps an exponential moving average of the gradient (first moment) and of the squared gradient (second moment) for every parameter, and divides one by the square root of the other:

m_t = β1·m_(t−1) + (1 − β1)·g_t
v_t = β2·v_(t−1) + (1 − β2)·g_t²
m̂ = m_t / (1 − β1^t),   v̂ = v_t / (1 − β2^t)
θ_t = θ_(t−1) − η · m̂ / (√v̂ + ε)

So Adam stores two extra numbers per parameter, on top of the parameter and its gradient. The optimiser itself, its variants and AdamW's decoupled weight decay are the subject of Video 07; here only the memory matters.

How a computer stores real numbers

A floating-point number is (−1)^sign × 1.mantissa × 2^(exponent − bias). The exponent field sets the range (how big and how small a number can be); the mantissa field sets the precision (how close together neighbouring numbers are). The formats that matter for training:

BF16 ("brain float", from Google Brain's TPU work) is FP32 with the mantissa cut short: same range, much less precision. FP16 keeps more precision but its range is tiny, so small gradients underflow to zero and large activations overflow to infinity. The two FP8 encodings were proposed together by Micikevicius et al. (2022): E5M2 follows IEEE conventions, while E4M3 drops infinities and keeps a single NaN pattern to stretch its range to 448. FP8 training and inference on newer GPUs is covered in Module 09 · Quantisation & hardware.

Mixed-precision training

Micikevicius et al. (2017, "Mixed Precision Training") showed you can store weights, activations and gradients in FP16 and still match FP32 accuracy, with two tricks:

• An FP32 master copy of the weights. Updates are often much smaller than the weight itself. Added in FP16, they round away to nothing (with a gap of about 0.001 after 1.0, adding 0.0001 changes nothing). So the optimiser updates an FP32 copy, which is rounded to half precision for the next forward pass.
• Loss scaling. Multiply the loss by a factor S before the backward pass so small gradients stay above FP16's underflow threshold, then divide the gradients by S before the update. Modern frameworks pick S dynamically and skip a step if an overflow appears.

BF16 has FP32's range, so it usually needs no loss scaling, which is why it became the default for large-model training. The FP32 master weights stay either way.

The memory bill: 16 bytes per parameter

The ZeRO paper (Rajbhandari, Rasley, Ruwase and He, 2019) counts the "model states" for mixed-precision Adam with Ψ parameters:

fp16 weights          2Ψ bytes
fp16 gradients        2Ψ
fp32 master weights   4Ψ  ┐
fp32 Adam m           4Ψ  ├ optimiser states: K·Ψ with K = 12
fp32 Adam v           4Ψ  ┘
total                16Ψ bytes

7.5B parameters  →  16 × 7.5e9 = 120 GB, before any activations

Activations, temporary buffers and fragmentation come on top. Activation memory is reduced separately, mainly with checkpointing (recomputing activations in the backward pass), which Module 08 · Distributed training covers.

Data parallelism and the all-reduce

In data parallelism each of N GPUs runs the full model on its own micro-batch, and the gradients are averaged with an all-reduce so every replica applies the same update. A ring all-reduce has two halves:

reduce-scatter:  after it, GPU i holds the SUM of chunk i over all GPUs
all-gather:      every GPU sends its finished chunk to all the others

data each GPU sends ≈ 2·(N − 1)/N × gradient size   (nearly independent of N)

The catch: all N GPUs hold the same 16Ψ bytes of model states. Adding GPUs adds throughput but never lets you fit a bigger model.

ZeRO stages 1, 2 and 3

ZeRO ("Zero Redundancy Optimizer") removes the duplication in three steps, and the per-GPU memory for model states becomes:

baseline DP      (2 + 2 + K)·Ψ          = 16Ψ
ZeRO-1  P_os     (2 + 2)·Ψ + K·Ψ/N
ZeRO-2  P_os+g   2Ψ + (2 + K)·Ψ/N
ZeRO-3  P_os+g+p (2 + 2 + K)·Ψ/N       = 16Ψ/N

7.5B params, N = 64:  120 GB → 31.4 GB → 16.6 GB → 1.9 GB

Stage 1 shards the optimiser states: GPU i updates only its 1/N of the parameters. Stage 2 also shards gradients: instead of an all-reduce, a reduce-scatter leaves each GPU with the summed gradient for exactly the slice it updates, and the rest is freed. After the update, an all-gather shares the new weights. Reduce-scatter plus all-gather is what an all-reduce already does, so the paper reports no extra communication for stages 1 and 2. Stage 3 shards the weights too. Before each layer's forward pass, its weights are all-gathered, used and freed; the same happens again in the backward pass, followed by a reduce-scatter of the gradients. That's about 1.5× the communication volume of plain data parallelism, in exchange for memory that falls linearly with N.

The ZeRO memory calculator in Module 08 · Distributed training lets you vary Ψ, N and K and see each stage's bar; this guide won't repeat it.

FSDP: the same idea, one unit at a time

PyTorch's Fully Sharded Data Parallel (Zhao et al., 2023) is the native implementation of the ZeRO-3 schedule. You wrap the model into units, typically one transformer block each. Each unit's parameters are flattened and sharded across GPUs. During the forward pass, FSDP all-gathers one unit, runs it, frees the full parameters, and moves on, prefetching the next unit's all-gather while the current one computes. The backward pass gathers each unit again, computes its gradients and reduce-scatters them. Peak memory is roughly the sharded states plus one or two fully gathered units, instead of the whole model.

for unit in model:                      # forward
    full = all_gather(unit.shard)       # overlapped with the previous unit's compute
    x = unit(x, full); free(full)
for unit in reversed(model):            # backward
    full = all_gather(unit.shard)
    grads = backward(unit, full); free(full)
    unit.grad_shard = reduce_scatter(grads)
optimizer.step()                        # each GPU updates its own shard only

The original FSDP API exposes sharding strategies that map onto ZeRO: FULL_SHARD is stage 3, SHARD_GRAD_OP is close to stage 2, NO_SHARD is ordinary DDP, and HYBRID_SHARD shards inside a node and replicates across nodes to keep the chatty all-gathers on fast links. The paper reports throughput comparable to DDP on models that fit, with near-linear scaling on models that don't.

3. Where the course goes deeper

4. Why it matters for robotics

Self-check

References

From the video's description:

• Paper: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (Rajbhandari, Rasley, Ruwase, He, 2019)
• Code: DeepSpeed
• Docs: PyTorch FSDP

Primary papers cited in this guide:

• Adam: A Method for Stochastic Optimization (Kingma, Ba, 2014)
• Mixed Precision Training (Micikevicius et al., 2017)
• FP8 Formats for Deep Learning (Micikevicius et al., 2022)
• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (Zhao et al., 2023)