Home / Fundamentals / Distributed Training

Distributed TrainingSeniorQuestion 3 of 7

Explain ZeRO stages 1–3 and FSDP. What do they save, and what do they cost?

30-second answer

Data parallelism replicates 16 bytes per parameter on every GPU. ZeRO-1 shards the 12 bytes of fp32 optimizer state, ZeRO-2 also shards gradients (reduce-scatter instead of all-reduce), both at the same communication volume as DDP. ZeRO-3/FSDP also shards parameters, all-gathering each layer group just before use in forward and backward and freeing it afterwards, for ~1.5× DDP's communication and 16/N bytes per parameter. A 7B model on 64 GPUs drops from 112 GB to ~29 GB (ZeRO-1) or ~1.75 GB (ZeRO-3) of state per GPU. Prefetching hides communication; hybrid sharding keeps all-gathers within nodes.

Try answering out loud before reading on. The walkthrough below is paced for a 40-minute interview round.

0–5 min Clarify and frame

Framing: "I'll start from the redundancy in data parallelism, go through ZeRO stages with their memory and communication, explain how FSDP executes a step, and finish with when to use what."

5–15 min Core explanation

In plain data parallelism, every GPU holds identical copies of 16 bytes per parameter. With N_d replicas, that's N_d − 1 redundant copies. ZeRO partitions them:

Stage Sharded Bytes/param/GPU Comm per step
DDP — 16 all-reduce (2S)
ZeRO-1 Optimizer state 4 + 12/N_d reduce-scatter + all-gather (2S)
ZeRO-2 + gradients 2 + 14/N_d same (2S)
ZeRO-3 / FSDP + parameters 16/N_d all-gather fwd + all-gather bwd + reduce-scatter (3S)

(S = size of the bf16 parameter or gradient tensor.)

ZeRO-1/2 for free: an all-reduce is already a reduce-scatter followed by an all-gather. ZeRO reduce-scatters gradients so each rank gets the summed gradient for its shard, updates only that shard's parameters and optimizer state, then all-gathers the updated parameters. Same bytes, less memory.

Example: 7B, 64 GPUs → DDP 112 GB, ZeRO-1 ~29 GB, ZeRO-2 ~15.5 GB, ZeRO-3 ~1.75 GB per GPU.

15–25 min Deep dive: an FSDP step

Scroll sideways to see the whole diagram.

One FSDP training step (per layer unit)[forward]; Rank i → Collectives: all-gather unit k params; Collectives → Unit k: full params; Rank i → Collectives: prefetch: all-gather unit k+1; Unit k: forward, then free non-owned params; Collectives → Unit k+1: full params; Unit k+1: forward, free; [backward (reverse order)]; Rank i → Collectives: all-gather unit k+1 again; Unit k+1: backward → full gradients; Unit k+1 → Collectives: reduce-scatter gradients; Collectives → Rank i: summed gradient shard; Note: unit k repeats: all-gather, backward, reduce-scatter; [step]; Rank i → Optimizer: update own shard (AdamW); Optimizer → Rank i: new bf16 shard.Rank iholds 1/N shardCollectivesNCCLUnit ktransformer blockUnit k+1Optimizershard of fp32 + AdamFORWARDall-gather unit k paramsfull paramsprefetch: all-gather unit k+1forward, then free non-owned paramsfull paramsforward, freeBACKWARD (REVERSE ORDER)all-gather unit k+1 againbackward → full gradientsreduce-scatter gradientssummed gradient shardunit k repeats: all-gather,backward, reduce-scatterSTEPupdate own shard (AdamW)new bf16 shard
One FSDP training step (per layer unit). Each rank stores only its shard of every unit. Parameters are all-gathered just before use and freed afterwards, in both forward and backward. Gradients are reduce-scattered so each rank keeps the summed gradient for its own shard, and each rank updates only that shard and its optimizer state.

The model is wrapped into units (typically one transformer block each). Each rank stores a flat shard of every unit's parameters.

  • Forward: for each unit, all-gather its full parameters, compute, then free the non-owned parts. Prefetch the next unit's all-gather while computing the current one.
  • Backward: all-gather the unit's parameters again (they were freed), compute gradients, reduce-scatter them so each rank keeps the summed gradient shard, free the full gradients.
  • Optimizer: each rank updates its own shard of fp32 master weights and Adam moments.

Peak memory ≈ sharded state + one or two full units + activations.

Tuning:

  • Unit size: too small means many small collectives with poor bandwidth; too large raises peak memory.
  • reshard_after_forward=False (ZeRO-2-like) keeps parameters gathered between forward and backward: less communication, more memory.
  • Hybrid sharding (HSDP): shard within a node of 8, replicate across nodes; all-gathers stay on NVLink, and only gradient reduction crosses nodes.

25–32 min Trade-offs

Option When it fits
DDP Model state fits easily; simplest, fastest
ZeRO-1 Large models with TP/PP; standard in 3D parallel setups
ZeRO-2 When gradients matter too, no extra communication
FSDP/ZeRO-3 Up to tens of billions of parameters without TP/PP; simpler code
FSDP + TP Large models where per-layer activations or communication need TP too

ZeRO-3's communication is per layer and per step, so at very large DP degrees across slow links it can become exposed; then TP/PP with ZeRO-1 or HSDP is better.

Offload (ZeRO-Infinity): move optimizer state and even parameters to CPU or NVMe. Enables huge models on few GPUs, but throughput is limited by PCIe.

32–37 min Production example

A team trained a 13B model with DDP + ZeRO-1 on 128 GPUs but couldn't fit sequence length 8k. Switching to FSDP with full sharding freed ~20 GB per GPU for activations; throughput dropped 8% due to all-gathers across nodes. HSDP (shard within nodes) recovered most of it, and selective recomputation made 8k fit with room for micro-batch 2.

37–40 min Likely follow-ups

  • How does FSDP interact with mixed precision? Shards are fp32 masters; all-gathers happen in bf16 to halve traffic.
  • How do you checkpoint FSDP? Save sharded state per rank (fast) and convert to a full checkpoint offline when needed.
  • Can ZeRO-3 replace tensor parallelism? For memory, often; but it doesn't reduce per-layer activation size or matmul size per GPU, so very large models still use TP.

What interviewers listen for

  • The memory formula per stage and a worked example.
  • Why ZeRO-1/2 add no communication.
  • FSDP's all-gather/free/reduce-scatter flow and hybrid sharding.

Red flags

  • Thinking ZeRO shards activations.
  • Not knowing ZeRO-3's extra communication.