Learn: ZeRO & FSDP

What you're learning

Why a 70B model won’t fit on an 80 GB GPU even before you touch a single activation, and the family of techniques — DeepSpeed ZeRO and PyTorch FSDP — that fix it by sharding the model states across the data-parallel group instead of replicating them. By the end you should be able to derive the canonical 16 bytes/param accounting, write the per-GPU memory formula for ZeRO-1/2/3, reason about the memory↔communication tradeoff each stage buys, and explain exactly why FSDP FULL_SHARD is ZeRO-3 wearing a different API.

This lesson builds directly on 08-data-parallelism-and-collectives — you need all-gather, reduce-scatter, and all-reduce (and the identity all-reduce = reduce-scatter + all-gather) to be automatic. It also leans on mixed-precision training (Tier 2) and activation recomputation from 06-memory-wall-and-flashattention. If the collectives don’t feel reflexive, reread lesson 08 first.

---

1. Learning map

graph TD
    A["Data parallelism<br/>replicates the model<br/>(lesson 08)"] --> B["Memory problem:<br/>16 bytes/param model states<br/>under mixed-precision Adam"]
    B --> C["Redundancy insight:<br/>every rank holds an<br/>identical copy — waste"]
    C --> D["ZeRO: shard the states<br/>across N ranks (arXiv:1910.02054)"]
    D --> D1["ZeRO-1<br/>shard optimizer states"]
    D --> D2["ZeRO-2<br/>+ shard gradients"]
    D --> D3["ZeRO-3<br/>+ shard parameters"]
    D3 --> E["FSDP FULL_SHARD<br/>= ZeRO-3 in PyTorch<br/>(arXiv:2304.11277)"]
    E --> E1["FSDP1: FlatParameter"]
    E --> E2["FSDP2: per-parameter<br/>DTensor sharding"]
    D3 --> F["Offload: trade<br/>bandwidth for capacity"]
    F --> F1["ZeRO-Offload → CPU<br/>(2101.06840)"]
    F --> F2["ZeRO-Infinity → NVMe<br/>(2104.07857)"]
    E --> G["Limit: comms volume grows,<br/>per-GPU params bounded"]
    G --> H["→ Tensor / pipeline parallelism<br/>(lesson 10)"]

    style B fill:#a44,color:#fff
    style D fill:#44a,color:#fff
    style E fill:#4a4,color:#fff
    style H fill:#a83,color:#fff

---

2. Why this matters

Plain data parallelism (DDP) is beautiful: replicate the model on every GPU, each processes a different micro-batch, all-reduce the gradients, everyone steps in lockstep. It scales throughput almost linearly and the only collective is one gradient all-reduce per step.¹ It has exactly one fatal flaw: it assumes the whole model — and all of training’s bookkeeping — fits on one GPU. For anything past ~a few billion parameters, it doesn’t.

The trap is that people budget for “the model” (the parameters) and forget that training a model in mixed precision with Adam carries 8× the parameter count in extra state. Let’s count bytes, because the count is the whole lesson.

The number that surprises people

A 70B model is 140 GB just in bf16 parameters ($2\times 70\text{B}$) — already larger than an 80 GB H100 before a single gradient, optimizer moment, or activation exists.² Add the full training state and you need 1.12 TB for model states alone. DDP cannot train this model on any number of GPUs, because DDP replicates all 1.12 TB on every GPU. The redundancy is the enemy.

---

3. The memory problem: deriving 16 bytes/param

Modern LLM training uses mixed precision (bf16/fp16 compute with an fp32 master copy) and the Adam optimizer. Walk through every per-parameter tensor that must live in memory across a step. This is the canonical “model states” accounting from the ZeRO paper.³

For a model with $\Psi$ parameters, per parameter you hold:

Tensor	Precision	Bytes/param	Why it exists
Parameters (working copy)	bf16/fp16	2	Used in fwd/bwd matmuls (tensor cores want 16-bit)
Gradients	bf16/fp16	2	Produced by backward, consumed by optimizer
Master weights	fp32	4	fp32 copy the optimizer actually updates (16-bit steps lose precision)
Adam momentum $m$	fp32	4	1st moment (running mean of grad)
Adam variance $v$	fp32	4	2nd moment (running mean of grad²)
Total model states		16

$$\underset{\text{fp16 }W}{\underbrace{2}}+\underset{\text{fp16 }\nabla }{\underbrace{2}}+\underset{\text{fp32 }{W}_{\text{master}}}{\underbrace{4}}+\underset{\text{fp32 }m}{\underbrace{4}}+\underset{\text{fp32 }v}{\underbrace{4}}=16\text{ bytes/param}$$

The ZeRO paper factors this as $(2+2+K)\Psi$ where $K=12$ is the optimizer multiplier for Adam — the 12 bytes of fp32 optimizer state (master + $m$ + $v$).³ So model states = 16Ψ bytes.

Figure: the 16 bytes/param decomposition. The 12 fp32 optimizer bytes (master weights + Adam’s two moments) dominate — that’s why ZeRO-1 alone, which shards only those, already cuts most of the redundancy.³

Why an fp32 master copy at all?

Adam’s update is $\theta \leftarrow \theta -\alpha \hat{m}/(\sqrt{\hat{v}}+ϵ)$. The step size per parameter can be tiny relative to the parameter’s magnitude; in bf16 (8-bit mantissa, ~2–3 decimal digits) that update rounds to zero — “swamping.” Keeping the master weights, $m$, and $v$ in fp32 preserves the small updates. The bf16 working copy is a rounded-down view refreshed each step for the fast matmuls.⁴

What 16 bytes/param is not: activations. Forward activations stashed for backward are a separate, batch- and sequence-dependent budget, attacked by activation recomputation / checkpointing (lesson 06) and — for the attention score matrix — FlashAttention. The DeepSpeed docs warn that for long sequences activations can rival params+grads+optim combined.⁵ ZeRO/FSDP shard model states; activations need their own strategy.

Quiz: Derive the memory to train a 13B model in fp32 (no mixed precision) with Adam, then in bf16 mixed precision with Adam. Which is larger?

Answer

Pure fp32 + Adam: params 4B + grads 4B + momentum 4B + variance 4B = 16 bytes/param → $16\times 13\text{B}=208$ GB. (No separate master copy — the fp32 params are the master.)
bf16 mixed precision + Adam: 2 (bf16 W) + 2 (bf16 ∇) + 4 (fp32 master) + 4 (m) + 4 (v) = 16 bytes/param → $16\times 13\text{B}=208$ GB.
They’re the same — 16 bytes/param either way. This is the punchline that trips people up: mixed precision does not shrink your model-state memory, because you pay for a second fp32 master copy that pure-fp32 training folds into the parameters themselves. Mixed precision buys you speed (16-bit tensor-core matmuls) and activation memory (bf16 activations), not optimizer-state memory. The optimizer state is the immovable 12 bytes/param in both cases.³⁴

---

4. ZeRO: shard the states across the DP group

The redundancy insight: in DDP with $N$ data-parallel ranks, all $N$ ranks hold identical copies of all 16Ψ bytes. That’s $N-1$ wasted replicas. ZeRO (Zero Redundancy Optimizer)³ partitions the states so each rank owns only $1/N$ of them, materializing the full tensor only transiently when a computation needs it. It comes in three cumulative stages — each shards one more category.

graph LR
    subgraph DDP["Baseline DDP (replicated)"]
        d0["rank 0<br/>P G Os"]
        d1["rank 1<br/>P G Os"]
        d2["rank N<br/>P G Os"]
    end
    subgraph Z1["ZeRO-1 (P_os)"]
        a0["rank 0<br/>P G os₀"]
        a1["rank 1<br/>P G os₁"]
        a2["rank N<br/>P G os_N"]
    end
    subgraph Z2["ZeRO-2 (P_os+g)"]
        b0["rank 0<br/>P g₀ os₀"]
        b1["rank 1<br/>P g₁ os₁"]
        b2["rank N<br/>P g_N os_N"]
    end
    subgraph Z3["ZeRO-3 (P_os+g+p)"]
        c0["rank 0<br/>p₀ g₀ os₀"]
        c1["rank 1<br/>p₁ g₁ os₁"]
        c2["rank N<br/>p_N g_N os_N"]
    end
    DDP --> Z1 --> Z2 --> Z3

Uppercase P G Os = full replica; lowercase subscripted = only this rank’s shard. Each stage lowercases one more column.³

Let $N_d$ be the data-parallel degree (number of GPUs sharing the shard). Per-GPU model-state memory:³

$$\begin{array}{rl}\text{DDP (baseline):}& (2+2+K)\Psi =16\Psi \\ \text{ZeRO-1 }(P_{os}):& 2\Psi +2\Psi +\frac{K\Psi }{N_d}=4\Psi +\frac{12\Psi }{N_d}\\ \text{ZeRO-2 }(P_{os+g}):& 2\Psi +\frac{(2+K)\Psi }{N_d}=2\Psi +\frac{14\Psi }{N_d}\\ \text{ZeRO-3 }(P_{os+g+p}):& \frac{(2+2+K)\Psi }{N_d}=\frac{16\Psi }{N_d}\end{array}$$

The pattern: whatever a stage shards contributes $(\cdot )/N_d$; whatever it still replicates contributes the full per-param bytes. Only ZeRO-3 drives per-GPU memory to zero as $N_d\to \infty$ — ZeRO-1 and ZeRO-2 have an irreducible replicated floor (4Ψ and 2Ψ respectively) that no amount of GPUs removes.

Why ZeRO-1 already gets you ~4×

With Adam, 12 of the 16 bytes are fp32 optimizer state. Sharding just that (ZeRO-1) at large $N_d$ takes you from 16Ψ toward 4Ψ — a 4× reduction for essentially free (no extra communication over DDP; see below). It’s the highest-leverage, lowest-cost stage, and the right first move when you’re only slightly over budget.³

The communication cost of each stage

This is the crux of the memory↔communication tradeoff. Recall from lesson 08 that a DDP gradient all-reduce moves $2\Psi$ worth of data per step (reduce-scatter Ψ + all-gather Ψ).¹ ZeRO’s communication analysis:³

• DDP: all-reduce grads = 2Ψ total volume.
• ZeRO-1 & ZeRO-2: same 2Ψ. ZeRO decomposes the all-reduce into a reduce-scatter of gradients (each rank ends up with the grad shard matching its optimizer shard) plus, after the local optimizer step, an all-gather of updated parameters. That’s still Ψ + Ψ = 2Ψ. You get 4×–8× memory savings for zero extra bytes on the wire. This is why ZeRO-2 is often the sweet spot.
• ZeRO-3: 3Ψ = 1.5× DDP. Because parameters are now sharded, each rank no longer has the full weights for forward/backward — it must all-gather them just-in-time per layer (Ψ in forward) and again in backward (Ψ), then reduce-scatter gradients (Ψ). $\Psi +\Psi +\Psi =3\Psi$, a 50% communication increase for potentially $N_d\times$ memory reduction.³

The tradeoff is not free, and it bites at scale

ZeRO-3’s extra all-gather turns forward and backward into a stream of collectives. On a fast intra-node fabric (NVLink/NVSwitch, ~900 GB/s on H100²) this overlaps with compute and hides well; across slower inter-node links it can dominate. That growing communication volume — not memory — is what eventually caps sharded data parallelism and motivates tensor/pipeline parallelism in lesson 10.

Quiz: You're training a 30B model. Optimizer states are killing you but bf16 params+grads (4 bytes/param = 120 GB) fit fine when replicated across your 32×A100-80GB. Which ZeRO stage do you pick, and what does it cost you in communication?

Answer

ZeRO-1. Your problem is only the 12 bytes/param of optimizer state ($12\times 30\text{B}=360$ GB replicated). ZeRO-1 shards exactly that: per-GPU model states become $4\Psi +12\Psi /32=120+11.25=131.25$ GB… which still doesn’t fit 80 GB — so actually you’d need ZeRO-2 ($2\Psi +14\Psi /32=60+13.1=73.1$ GB, fits) here. The instructive point: ZeRO-1 leaves a 4Ψ = 120 GB replicated floor that 32 GPUs can’t reduce. Communication cost of either ZeRO-1 or ZeRO-2 is identical to DDP (2Ψ) — you pay nothing extra on the wire until you reach ZeRO-3. Always push to the lowest stage that fits; the memory floor (4Ψ for Z1, 2Ψ for Z2) decides whether it does.

---

5. FSDP: ZeRO-3, the PyTorch way

FSDP FULL_SHARD is architecturally equivalent to ZeRO-3.¹ It shards parameters, gradients, and optimizer state across the DP group; per-rank memory drops from $O(\Psi )$ to $O(\Psi /N)$ at the cost of one extra all-gather per unit per pass — exactly the 3Ψ communication story above. Terminology differs; mechanism is the same. (FSDP also exposes SHARD_GRAD_OP ≈ ZeRO-2 and NO_SHARD = DDP as intermediate modes.)

The unit of sharding and the just-in-time dance

FSDP groups the model into units (typically a transformer block per unit) via a wrapping policy (e.g. transformer_auto_wrap_policy or, in FSDP2, calling fully_shard on each block). Each unit’s parameters are sharded across ranks. At runtime, per unit:¹

sequenceDiagram
    participant C as Compute (default stream)
    participant K as Comm (separate stream)
    Note over C,K: FORWARD, unit i
    K->>K: all-gather unit i params (→ full weights)
    C->>C: run unit i forward
    C->>C: free unit i params (keep only shard)
    Note over C,K: (prefetch all-gather unit i+1 during unit i compute)
    Note over C,K: BACKWARD, unit i
    K->>K: all-gather unit i params again
    C->>C: run unit i backward → grads
    K->>K: reduce-scatter grads (each rank keeps its shard)
    C->>C: free params; local optimizer step on shard

The key moves: all-gather params just before compute, free them right after (so only one unit’s full weights are ever resident, not the whole model), and reduce-scatter gradients so each rank keeps only the shard matching its optimizer state. Because parameters are re-gathered in backward, backward pays a second all-gather — hence 3Ψ total.

Overlap via prefetch is what makes it fast

FSDP issues all-gathers on a separate CUDA stream so the collective for unit $i+1$ overlaps with the compute of unit $i$ (forward prefetch), and analogously in backward (backward prefetch).¹ Without this overlap you’d serialize comm and compute and eat the full 3Ψ latency. In eager mode FSDP can’t reorder the all-gather before its consuming compute, so it deliberately launches on the side stream to break the false dependency — the paper calls this out as a core implementation subtlety.¹

FSDP1 vs FSDP2 — current status (Sept 2026)

This is actively evolving, so pin the distinction:

• FSDP1 shards a FlatParameter: it flattens, concatenates, and chunks a group of a unit’s tensors into one 1-D buffer, then splits that buffer across ranks. Efficient, but the flat buffer makes per-original-parameter reasoning awkward — mixed dtypes/requires_grad within a unit, sharded state dicts, and composing with other parallelisms all get hairy.⁶
• FSDP2 (the fully_shard API) replaces the FlatParameter with per-parameter sharding backed by DTensor: each parameter is chunked on dim-0 (torch.chunk(dim=0)) and represented as a DTensor with Shard(dim=0) placement. This gives a more intuitive representation, relaxes the frozen-parameter constraints, enables communication-free (already-sharded) state dicts, and makes optimizer/grad-clipping/checkpointing “just work” on DTensors — while preserving FSDP1-level throughput.⁶ As of PyTorch 2.x, FSDP2 (fully_shard) is the recommended path and PyTorch is actively steering new code toward it; FSDP1 (FullyShardedDataParallel) remains available but is effectively in maintenance. The APIs are not backward-compatible (though checkpoints are largely portable).⁶⁷

Why per-parameter DTensor sharding matters beyond aesthetics

DTensor gives every parameter a first-class description of how it’s distributed (device mesh + placement). That’s the substrate for composing FSDP with tensor parallelism (2-D parallelism / HSDP): a parameter can be Shard along the FSDP mesh dimension and Shard/Replicate along the TP mesh dimension simultaneously. FlatParameter’s opaque 1-D buffer made this composition painful; DTensor makes it declarative.⁶ This is the plumbing that lesson 10’s parallelisms plug into.

Quiz: A colleague says "we use FSDP, so we're not using ZeRO." Is that a meaningful distinction? Where do the two genuinely differ?

Answer

At the level of what memory ends up where and what collectives run, no meaningful distinction: FSDP FULL_SHARD = ZeRO-3 (shard P+G+Os; all-gather params in fwd/bwd, reduce-scatter grads; 3Ψ ≈ 1.5× DDP communication).¹ The genuine differences are implementation and ecosystem: (1) FSDP is native PyTorch, integrates with autograd/DTensor/torch.compile; ZeRO is a DeepSpeed engine with its own launcher, config JSON, and optimizer wrappers. (2) FSDP2 shards per-parameter (DTensor, dim-0 chunk) whereas ZeRO shards a flattened buffer (like FSDP1’s FlatParameter).⁶ (3) DeepSpeed’s offload/Infinity stack (CPU+NVMe, §6) is more mature; FSDP has CPU offload but not an equivalent NVMe engine. So “FSDP vs ZeRO” is a framework choice, not an algorithm choice — both are sharded data parallelism.

---

6. Offloading: trade bandwidth for capacity

Sharding across $N_d$ GPUs still requires $N_d$ GPUs’ worth of aggregate HBM. When you don’t have enough GPUs, push state down the memory hierarchy — CPU DRAM, then NVMe SSD — accepting far lower bandwidth in exchange for far more capacity.

• ZeRO-Offload offloads the optimizer states and the fp32 optimizer step to CPU (built on ZeRO-2). The GPU does forward/backward in bf16; gradients go to CPU, the Adam update runs on CPU DRAM, updated params come back.⁸ The bottleneck is the single-GPU↔CPU PCIe link (~12 GB/s on PCIe Gen3), so ZeRO-Offload needs large per-GPU batch sizes (~2K) to amortize the transfer and keep the GPU busy.⁹
• ZeRO-Infinity generalizes this to CPU and NVMe, built on ZeRO-3, and can offload all model states (params, grads, optimizer) plus activations.⁹ Two ideas make the slow links tolerable: bandwidth-centric partitioning (retrieve a tensor’s pieces in parallel across all GPUs’ PCIe/NVMe links so aggregate bandwidth scales with device count, instead of bottlenecking on one link) and an overlap-centric design (overlap NVMe→CPU and CPU→GPU transfers with GPU compute). It also adds memory-centric tiling to run arbitrarily large operators without model-parallel refactoring.⁹

The capacity numbers are staggering — and so are the caveats

A cluster has ~2–3× more CPU DRAM than GPU HBM and ~50× more NVMe than HBM.⁹ ZeRO-Infinity reported fitting 1 trillion params per single V100 DGX-2 node and showed a 100T-param model’s states fit in the NVMe of a 96-node (1536-GPU) cluster.⁹ But this trades bandwidth for capacity: params/grads need >~70 GB/s effective bandwidth to keep GPUs fed⁹, vs. HBM’s ~900 GB/s+². Offloading is a capacity escape hatch for memory-constrained setups, not a throughput win — expect lower MFU. Use it when you’d otherwise be unable to run at all.

Quiz: You can fit a model with ZeRO-3 on 64 GPUs at 30% MFU, or run it with ZeRO-Infinity NVMe offload on 8 GPUs at maybe 10% MFU. When is offload the right call?

Answer

Offload wins when GPUs are the scarce/expensive resource and wall-clock throughput is secondary — e.g. you only have 8 GPUs, or you’re fine-tuning where a slow run is acceptable, or you’re doing single-node experimentation on a model that can’t otherwise fit. The 8-GPU offload run uses 8× fewer accelerators; even at ⅓ the MFU it can be far cheaper per experiment if GPU-hours are the budget. ZeRO-3-on-64 wins when you have the fleet and throughput/time-to-result matters — the extra GPUs give you both more aggregate HBM (no slow-link tax) and more compute. Rule of thumb: offload trades time and MFU for the ability to run on fewer/smaller GPUs; reach for it only when you can’t satisfy the HBM budget any other way.⁹

---

7. Composing with TP/PP and DP (teaser for lesson 10)

Sharded data parallelism is one axis. Real frontier runs use N-D parallelism, and ZeRO/FSDP is the “DP” leg:

• DP × ZeRO/FSDP: the model states are sharded across the data-parallel dimension; different DP groups process different data. This is sharded data parallelism.
• + Tensor Parallelism (TP): shard individual matmuls (e.g. attention/MLP weight matrices) within a layer across a small, high-bandwidth group (usually intra-node over NVLink). FSDP2’s DTensor placements compose with TP’s — a weight is Shard on the FSDP mesh dim and Shard on the TP mesh dim.⁶ This is the standard 2-D (FSDP+TP) recipe (e.g. torchtitan).⁶
• + Pipeline Parallelism (PP): split the model’s layers across GPU groups (stages) and pipeline micro-batches through them to keep every stage busy.
• Hybrid Sharded Data Parallel (HSDP): shard within a node, replicate across nodes — bounds the all-gather to fast intra-node links while keeping inter-node traffic to a periodic all-reduce.

The mental model

ZeRO/FSDP shards along the batch/data axis (every rank sees the whole computation graph, just a slice of each tensor and a different micro-batch). TP shards inside an operator; PP shards across layers. They’re orthogonal and multiply: a 405B run might be FSDP × TP=8 × PP=… — lesson 10 makes this precise.

---

8. The GPU/LLM connection — and why lesson 10 exists

FSDP / ZeRO-3 is the default way to train a model that won’t fit on one GPU. Shard the 16Ψ bytes of model states across $N_d$ GPUs, all-gather each layer’s params just-in-time, free them after, reduce-scatter the grads. Per-GPU model-state memory falls as $16\Psi /N_d$ — arbitrarily small with enough GPUs.

But it has a hard limit, and the limit is the whole reason lesson 10 exists:

1. Communication volume grows. ZeRO-3/FSDP pays 3Ψ per step (1.5× DDP), and every layer’s forward and backward now blocks on an all-gather.³¹ Beyond the fast intra-node fabric, this comm can’t fully overlap and MFU craters.
2. Per-GPU transient params are still bounded by one unit. Even fully sharded, at the moment you compute a layer you must materialize that layer’s full (un-sharded) parameters and activations on one GPU. A single layer whose weights or activations exceed one GPU’s HBM cannot be trained by sharding the batch axis alone — no $N_d$ helps.

Both pressures point the same way: stop sharding only the data axis. Tensor parallelism splits the operator itself; pipeline parallelism splits the layer stack. Lesson 10 shows how to combine them with FSDP into the N-D parallelism used for genuinely frontier-scale models.

---

9. Practice problems

Practice problem 1: per-GPU memory for a 70B model at ZeRO-1/2/3 on 64 GPUs

Problem: You have 64× H100-80GB. You want to train a 70B model in bf16 mixed precision with Adam. Compute per-GPU model-state memory under DDP, ZeRO-1, ZeRO-2, and ZeRO-3. Which stages fit in 80 GB (ignoring activations/workspace for now)?

Worked solution: $\Psi =70\text{B}$, $K=12$, $N_d=64$. Use GB $=10^9$ bytes.

• DDP (baseline): $16\Psi =16\times 70=1120\mathbf{GB}$ per GPU. Absurd — 14× an H100. DDP cannot train this on any GPU.
• ZeRO-1: $4\Psi +12\Psi /N_d=280+\frac{840}{64}=280+13.1=293.1\mathbf{GB}$. Does not fit — the 4Ψ = 280 GB replicated floor (bf16 params + bf16 grads) dominates.
• ZeRO-2: $2\Psi +14\Psi /N_d=140+\frac{980}{64}=140+15.3=155.3\mathbf{GB}$. Still does not fit — the 2Ψ = 140 GB replicated bf16 params exceed 80 GB.
• ZeRO-3: $16\Psi /N_d=\frac{1120}{64}=17.5\mathbf{GB}$. Fits comfortably, leaving ~60 GB per GPU for activations, comm buffers, and the transient all-gathered layer.

Takeaway: For a 70B model on 64 GPUs, only full parameter sharding (ZeRO-3/FSDP) fits — because ZeRO-1 and ZeRO-2 each leave a large replicated floor (4Ψ and 2Ψ) that no number of GPUs removes. This is the concrete reason “just use ZeRO-3/FSDP” is the default for large models. (Sanity check the floors: bf16 params alone = 140 GB > 80 GB, so any stage that replicates params is dead on arrival here.)

Practice problem 2: how many GPUs does ZeRO-2 need for 70B, and what does that reveal?

Problem: Ignoring activations, what data-parallel degree $N_d$ would ZeRO-2 need to fit the 70B model’s model states in 80 GB/GPU? Then compute ZeRO-3’s minimum $N_d$. Comment.

Worked solution:

• ZeRO-2: need $2\Psi +14\Psi /N_d\le 80$ GB. But $2\Psi =140$ GB $>80$ GB regardless of $N_d$ (the $14\Psi /N_d$ term is non-negative). So there is no finite $N_d$ that makes ZeRO-2 fit — the replicated bf16-param floor alone exceeds one GPU. ZeRO-2 is fundamentally incapable of training a model whose bf16 parameters don’t fit on a single GPU.
• ZeRO-3: need $16\Psi /N_d\le 80\Rightarrow N_d\ge 16\times 70/80=1120/80=14$ GPUs (for model states only; in practice you need headroom for activations + the transient all-gathered layer + comm buffers, so more).
  Takeaway: The dividing line between “ZeRO-1/2 can help” and “you must shard parameters” is exactly whether the bf16 parameters (2Ψ) fit on one GPU. Once $2\Psi >$ single-GPU HBM, only ZeRO-3/FSDP (or model parallelism, lesson 10) can train it. For 70B: $2\Psi =140$ GB, so ZeRO-3 is mandatory.

---

10. Practice & resources

Reading order — paper → docs → code → hands-on:

1. Primary source (ZeRO): ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, arXiv:1910.02054 — read §5 (memory analysis, the $(4+K)\Psi$ accounting) and §7 (communication analysis, the 2Ψ vs 3Ψ result). This is the canonical reference for everything in §3–4.³
2. FSDP paper: PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel, arXiv:2304.11277 — the FlatParameter design, stream-based overlap, and the ZeRO-3 equivalence.¹
3. FSDP2 docs & tutorial: torch.distributed.fsdp.fully_shard and the Getting Started with FSDP2 tutorial — per-parameter DTensor sharding, wrapping via fully_shard, migration from FSDP1.⁷ Cross-read torchtitan’s FSDP notes for 2-D (FSDP+TP) composition.⁶
4. DeepSpeed ZeRO tutorial: deepspeed.ai/tutorials/zero — stage config, and the memory estimators estimate_zero{2,3}_model_states_mem_needs_all_live to check your arithmetic against the framework.⁵
5. Offloading: ZeRO-Offload, arXiv:2101.06840 and ZeRO-Infinity, arXiv:2104.07857.⁸⁹

Hands-on:

• Train with FSDP2. Wrap a small transformer with fully_shard per block, run on 2+ GPUs (or simulate with torchrun --nproc_per_node). Toggle reshard_after_forward (keep-gathered vs. free-after-forward) and watch the memory/throughput tradeoff.
• Inspect memory with torch.cuda. Bracket a step with torch.cuda.reset_peak_memory_stats() / torch.cuda.max_memory_allocated(), and dump torch.cuda.memory_summary(). Confirm the per-GPU model-state footprint matches $16\Psi /N_d$ for FULL_SHARD vs. the replicated $16\Psi$ for NO_SHARD. Use torch.cuda.memory._dump_snapshot() + the memory visualizer to see the transient all-gathered-layer spike during forward.
• Reproduce the stage table. For a model you have, print DeepSpeed’s estimator output at ZeRO-1/2/3 and check it against the formulas in §4.

---

11. What’s next

Three threads, in order:

1. 10-model-parallelism (next lesson) — tensor and pipeline parallelism: what to do when communication volume caps FSDP and when a single layer won’t fit on one GPU. This is the direct sequel motivated by §8.
2. 08-data-parallelism-and-collectives (back-reference) — if the all-gather/reduce-scatter/all-reduce decomposition or the 2Ψ all-reduce cost felt shaky, this is the foundation ZeRO/FSDP is built on.
3. Activation memory & recompute — revisit 06-memory-wall-and-flashattention: ZeRO/FSDP shard model states but not activations; the two budgets are attacked by orthogonal techniques and you must plan both.

Suggested concept stubs to write: zero, fsdp, optimizer-state-memory.

---

Part of gpu-systems-for-llms · Filed: 2026-09-02

---

References

Footnotes

1. Zhao et al. (Meta), “PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel,” arXiv:2304.11277 (2023). https://arxiv.org/abs/2304.11277 — FSDP FULL_SHARD is architecturally equivalent to ZeRO-3; FlatParameter unit sharding; all-gather params per unit just-in-time in fwd and bwd then free, reduce-scatter grads; overlap via a separate CUDA stream for all-gathers (forward/backward prefetch) to break the false dependency in eager mode. [established] ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶ ↩⁷ ↩⁸ ↩⁹

2. NVIDIA H100 SXM: 80 GB HBM3, ~3.35 TB/s HBM bandwidth; intra-node NVLink/NVSwitch ~900 GB/s aggregate per GPU. NVIDIA H100 datasheet. https://resources.nvidia.com/en-us-hopper-architecture/nvidia-h100-datasheet — used here for the 80 GB per-GPU budget and the intra-node vs. HBM bandwidth contrast. [established] ↩ ↩² ↩³

3. Rajbhandari et al., “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models,” arXiv:1910.02054 (2019/2020). https://arxiv.org/abs/1910.02054 — defines “model states” as $(2+2+K)\Psi$ bytes with $K=12$ for Adam (fp16 params 2B + fp16 grads 2B + fp32 master 4B + fp32 momentum 4B + fp32 variance 4B = 16 bytes/param); per-device memory $4\Psi +K\Psi /N_d$ (ZeRO-1), $2\Psi +(2+K)\Psi /N_d$ (ZeRO-2), $(4+K)\Psi /N_d$ (ZeRO-3); communication analysis: DDP and ZeRO-1/2 = 2Ψ, ZeRO-3 = 3Ψ (1.5× DDP). [established] ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶ ↩⁷ ↩⁸ ↩⁹ ↩¹⁰ ↩¹¹ ↩¹²

4. Micikevicius et al., “Mixed Precision Training,” arXiv:1710.03740 (ICLR 2018). https://arxiv.org/abs/1710.03740 — the fp32 master-weight copy exists because small bf16/fp16 weight updates swamp (round to zero); 16-bit working copy is used for the tensor-core matmuls. ↩ ↩²

5. DeepSpeed “Memory Requirements” docs and ZeRO tutorial. https://deepspeed.readthedocs.io/en/latest/memory.html and https://www.deepspeed.ai/tutorials/zero/ — confirms the byte constants (Adam momentum 4B + variance 4B, fp16 params 2B / fp32 4B), provides estimate_zero{2,3}_model_states_mem_needs_all_live, and warns activations can rival params+grads+optimizer states for long sequences. ↩ ↩²

6. PyTorch FSDP2 design: per-parameter dim-0 sharding via torch.chunk(dim=0) represented as DTensor with Shard(dim=0) placement (vs. FSDP1’s flatten-concat-chunk FlatParameter); benefits include communication-free sharded state dicts, relaxed frozen-parameter constraints, and out-of-the-box DTensor compatibility with optimizer/grad-clip/checkpoint; composes with tensor parallelism for 2-D parallelism. Per-Parameter-Sharding FSDP RFC, pytorch/pytorch#114299 (https://github.com/pytorch/pytorch/issues/114299); torchtitan FSDP docs (https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md). [recent] ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶ ↩⁷ ↩⁸

7. PyTorch docs, torch.distributed.fsdp.fully_shard (https://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html) and “Getting Started with Fully Sharded Data Parallel (FSDP2)” tutorial (https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html). As of PyTorch 2.x (docs through 2.13, Sept 2026), fully_shard (FSDP2) is the recommended API with an active migration path from the legacy FullyShardedDataParallel (FSDP1); APIs are not backward-compatible though checkpoints are largely portable. [recent] ↩ ↩²

8. Ren et al., “ZeRO-Offload: Democratizing Billion-Scale Model Training,” arXiv:2101.06840 (2021). https://arxiv.org/abs/2101.06840 — offloads optimizer states and the fp32 optimizer step to CPU (built on ZeRO-2); GPU↔CPU PCIe bandwidth (~12 GB/s Gen3) is the bottleneck. [established] ↩ ↩²

9. Rajbhandari et al., “ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning,” arXiv:2104.07857 (SC’21). https://arxiv.org/abs/2104.07857 — built on ZeRO-3; offloads all model states (and activations) to CPU DRAM and NVMe; bandwidth-centric partitioning aggregates PCIe/NVMe bandwidth across all devices; overlap-centric design; memory-centric tiling. Reports NVMe ~25 GB/s/node, PCIe single-GPU ~12 GB/s, HBM ~900 GB/s; clusters have ~2–3× more CPU DRAM and ~50× more NVMe than HBM; ZeRO-Offload needs ~2K batch/GPU for efficiency; fits ~1T params/DGX-2 node and 100T-param states on a 96-node cluster. [established] ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶ ↩⁷ ↩⁸