TL;DR: Mixed-precision Adam costs 16 bytes/param (2 BF16 weights + 2 BF16 grads + 4+4+4 FP32 master/momentum/variance), and ZeRO/FSDP shard those pieces in order of size-to-pain: stage 1 optimizer states (free on the wire), stage 2 gradients (still ~2K traffic), stage 3 parameters (~3K traffic, near-linear memory scaling). Gradient checkpointing attacks the separate activation budget, trading ~30% recompute for memory.
How to approach it
Start from the memory anatomy of mixed-precision Adam training, the famous 16 bytes per parameter, because once that is on the table the three ZeRO stages are just "shard each piece in order of size-to-pain ratio." Then give the communication cost of each stage, because that is the half candidates skip and interviewers do not.
A strong answer
Plain data parallelism stores on every GPU, per parameter, with mixed-precision Adam:
BF16 weights 2 bytes
BF16 gradients 2 bytes
FP32 master weights 4 bytes
FP32 Adam momentum 4 bytes
FP32 Adam variance 4 bytes
16 bytes/param (+ activations on top)
So a 7B model needs ~112 GB of state per GPU before a single activation, already over an 80 GB H100. That redundancy across N data-parallel ranks is pure waste, and ZeRO (DeepSpeed) / FSDP (PyTorch) remove it in stages:
- Stage 1, shard optimizer states (the 12 FP32 bytes): per-GPU drops to 4 + 12/N bytes/param. Communication unchanged from DP (gradients still all-reduced, implemented as reduce-scatter plus the optimizer step happening on shards).
- Stage 2, also shard gradients: 2 + 14/N bytes/param. Still ~2K bytes on the wire per step, same as DP; gradients are reduce-scattered to their owner rank instead of all-reduced.
- Stage 3 / full FSDP, also shard parameters: 16/N bytes/param. Now each layer's weights must be all-gathered just-in-time for forward, re-gathered (or kept) for backward, and gradients reduce-scattered. Total traffic ≈ 3K vs DP's 2K, a 1.5x communication tax in exchange for near-linear memory scaling. FSDP overlaps the all-gathers with compute via prefetch, so on NVLink plus fast fabrics the tax is mostly hidden; on slow interconnects it is not, and stage 2 with a smaller model per rank can beat stage 3.
The stages side by side, with the per-param byte math:
| ZeRO stage | What is sharded | Bytes per param | Comm cost |
|---|---|---|---|
| Baseline / DP | nothing (full copy per GPU) | 16 | ~2K (all-reduce) |
| Stage 1 | optimizer states | 4 + 12/N | unchanged from DP |
| Stage 2 | + gradients | 2 + 14/N | ~2K (reduce-scatter) |
| Stage 3 / FSDP | + parameters | 16/N | ~3K (1.5x DP) |
Numbers that land well: 7B with Adam on 8 GPUs goes from 112 GB/GPU (impossible) to ~14 GB/GPU of state under full sharding, so fine-tuning a 7B on a single 8×A100 node stops being a memory problem at all.
Gradient checkpointing attacks the other consumer: activations, which scale with batch × sequence length × depth and are not touched by ZeRO. Instead of storing every layer's activations for backward, you keep checkpoints every k layers and recompute the rest during backward. Storage drops roughly by the checkpoint interval (the classic √L scheme stores O(√L) layers); cost is ~30% extra compute (one extra forward through most of the net). It composes with FSDP and is the standard trade when sequence length or batch pushes you over, because recompute is cheap and HBM is not. Same philosophy as FlashAttention: spend FLOPs to avoid memory.
Decision rule I would state: fits with DP, use DP; optimizer states are what kill you (they usually are), do stage 1/2 first since they are communication-free; params themselves do not fit, full FSDP/stage 3; activations are the binding constraint, checkpointing (and sequence parallelism at extreme lengths). Hybrid-shard (shard within a node, replicate across nodes, HSDP) when the fabric cannot afford stage-3 traffic at full width.
What interviewers probe next
- "Why is stage 1 nearly free?" Optimizer states are only read/written during the optimizer step, which is already local; sharding them changes no inter-step communication. It is 12 of the 16 bytes for zero wire cost, so always take it.
- "FSDP vs TP, both shard parameters, what is the difference?" FSDP gathers whole layers transiently and computes locally (memory technique); TP computes on partial layers and syncs activations (compute partitioning). Different traffic patterns, different latency sensitivity.
- "Where does CPU offload fit?" ZeRO-Offload/Infinity push optimizer states (or params) to host RAM/NVMe; viable for fine-tuning where step time is long, brutal for pretraining where PCIe becomes the step clock.
- "How would you fine-tune 70B on one 8×80GB node?" Run the arithmetic first: full-Adam state is 70B × 16 = 1,120 GB, so even sharded across 8 GPUs that is 140 GB/GPU and does not fit. So you drop the optimizer-state cost: LoRA/QLoRA (train adapters, base weights frozen and 4-bit) or a memory-lean optimizer, plus activation checkpointing. And the same byte math shows why QLoRA closes the gap rather than just asserting it (verified): the frozen base at 4-bit is 70B × 0.5 = 35 GB total, and if the trainable adapters are ~1% of parameters, their full 16-byte Adam state is 0.7B × 16 ≈ 11 GB, so the entire training state is ~46 GB against the node's 640 GB, versus 1,120 GB for full fine-tuning. The 24x collapse comes from one structural fact: the 16-byte tax applies only to trainable parameters, so shrinking what trains beats sharding what doesn't. Showing you do the byte math before naming a technique is the actual test, and finishing the technique's own math is the flourish.
Common mistakes
- Reciting "stage 1, 2, 3" without the byte accounting that makes them meaningful.
- Claiming ZeRO-3 has "the same communication as DP." It is ~1.5x, and on Ethernet-class fabrics that difference is visible in step time.
- Confusing FSDP with tensor parallelism (both "split the model" in slide-speak).
- Forgetting activations exist, then sharding 16 bytes/param beautifully and being unable to explain why long-context fine-tuning still OOMs.
Key takeaways
- The 16 bytes/param breakdown makes the three stages obvious: shard biggest-and-cheapest-to-shard first.
- Stage 1 is free on the wire, stage 3 costs ~1.5x DP traffic for near-linear memory scaling.
- ZeRO never touches activations; gradient checkpointing is the separate, ~30%-recompute lever for those.