ZeRO-1/2/3 & Offload

ZeRO (Zero Redundancy Optimizer, DeepSpeed) fits large-model training into limited GPU memory by eliminating redundant copies of state. With $\Psi$ total parameters (FP16, 2 bytes/element) trained in mixed precision with Adam, training must simultaneously hold the FP16 params $2\Psi$, the FP16 gradients $2\Psi$ produced by backprop, and the FP32 optimizer states — a master weight copy $4\Psi$ + first moment $4\Psi$ + second moment $4\Psi$ = $12\Psi$ — giving a static total of $16\Psi = 2\Psi + 2\Psi + 12\Psi$. For a 70B model: $16 \times 70$ GB $= 1120$ GB, which one H100 (80 GB) cannot hold, and naive DDP replicating the full state $N_d$ times across $N_d$ ranks makes it worse. ZeRO shards in three stages (per-rank memory with $N_d$ ranks): ① ZeRO-1 shards only optimizer states: $2\Psi + 2\Psi + \tfrac{12\Psi}{N_d}$; ② ZeRO-2 also shards gradients: $2\Psi + \tfrac{2\Psi + 12\Psi}{N_d}$; ③ ZeRO-3 shards parameters as well: $\tfrac{16\Psi}{N_d}$ — 70B needs $N_d = 16$ H100s ($1120/16 = 70$ GB each). Still not enough, offload: move optimizer states and gradients to CPU RAM while the GPU keeps only the local FP16 param shard $2\Psi/N_d = 140/8 = 17.5$ GB (8 GPUs), paging to NVMe on demand (ZeRO-Infinity, scaling to trillion-parameter models).

the canonical memory-budget calculation for pre-training and fine-tuning (7B–175B); favorite interview follow-ups are computing 70B training memory by hand, what ZeRO-2 saves over ZeRO-1, and what offload costs.

naive DDP stores $N_d$ full copies of state so per-rank memory never shrinks as GPUs scale; ZeRO cuts the redundancy from $O(N_d)$ to $O(1)$ — per-rank memory drops by an order of magnitude (70B: 1120 GB → 70 GB/rank) while communication stays close to DDP (~1.5× for ZeRO-3: ~3Φ vs DDP's 2Φ); combined with CPU/NVMe offload, even 175B-class models train on a single node.

By hand, compute the per-rank static memory for training a 70B model (Adam + FP16/FP32 mixed precision): write $16\Psi = 2\Psi + 2\Psi + 12\Psi$, name each term (FP16 params / FP16 grads / FP32 master weights + first + second moments), arrive at 1120 GB, and explain why naive DDP still costs 1120 GB per rank (N redundant copies).

Write the per-rank memory formulas for ZeRO-1/2/3: $2\Psi + 2\Psi + \tfrac{12\Psi}{N_d}$, $2\Psi + \tfrac{2\Psi + 12\Psi}{N_d}$, $\tfrac{16\Psi}{N_d}$; state what each stage shards and what remains replicated.

What does ZeRO-2 save over ZeRO-1, and by how much (gradients drop from $2\Psi$ to $2\Psi/N_d$ per rank)? Why can gradients be sharded — Reduce-Scatter right after backprop instead of a full All-Reduce.

Why is ZeRO-3's communication ~1.5× DDP (forward All-Gather Φ + backward Reduce-Scatter Φ + a second All-Gather for recompute Φ ≈ 3Φ, vs DDP's 2Φ), and is the memory win $16\Psi \to 16\Psi/N_d$ worth it?

What does ZeRO-Offload move to the CPU and at what cost: optimizer states and the Adam step run on the CPU while the GPU keeps only the layer's param shard $2\Psi/N_d$; memory drops dramatically at the price of PCIe-bound, slower steps; how does ZeRO-Infinity extend to NVMe?