FSDP / ZeRO — how do they save memory?

hard

Answer

$Standard DDP: each GPU holds full model + gradients + optimizer state.$
$Adam optimizer states = 2x model params.$
$FSDP (PyTorch) / ZeRO-3 (DeepSpeed): shard model params + gradients + optimizer states across GPUs.$
$All-gather params before forward, discard after.$
$All-reduce gradients + reshard.$
$(\mathrm{params}\; + \;\mathrm{grads}\; + \;\mathrm{optimizer})\; / \;N_{\mathrm{gpus}}$
$Enables 10-100x larger models.$
$ZeRO stages: 1 (shard optimizer), 2 (+gradients), 3 (+params).$
$Standard for LLM training.$

Check yourself — multiple choice

Random
Standard DDP: full model+grads+optim per GPU; FSDP/ZeRO-3: shard params+grads+optim across GPUs; all-gather forward + discard + all-reduce grads + reshard; enables 10-100x larger; ZeRO stages 1/2/3
Same
Not real

FSDP/ZeRO: shard params+grads+optim; enables large LLM training.

#infrastructure#advanced

215 interview questions in this topic.