Note: FSDP / ZeRO, Sharding Optimizer State and Parameters Across GPUs

The problem, from first principles

01-data-parallelism-gradient-averaging's DDP-style data parallelism has a hidden cost this question makes explicit: every single GPU keeps a full, redundant copy of the model's parameters, gradients, AND optimizer state (for Adam, that's fp32 copies of the parameters, momentum, and variance — far bigger than the parameters themselves). ZeRO (Zero Redundancy Optimizer) asks a simple question: if every GPU needs the FULL model to do a forward pass, but only needs its OWN SHARE of the optimizer's bookkeeping, why replicate the bookkeeping at all?

From theory to code

Implement per_gpu_memory_bytes(num_params, num_gpus, zero_stage), following the ZeRO paper's exact mixed-precision-Adam memory model for stages 0 (no sharding) through 3 (shard everything), and memory_reduction_factor, quantifying the savings against stage 0.

Constraints

• Using bytes-per-parameter units: fp16 params = 2, fp16 gradients = 2, fp32 Adam optimizer state = 12 (fp32 param copy + momentum + variance, 4 bytes each).
• zero_stage=0: (2 + 2 + 12) * num_params — full replication, no sharding at all (16Ψ bytes, the ZeRO paper's baseline figure).
• zero_stage=1 (P_os): shard only the optimizer state: (2 + 2) * num_params + 12 * num_params / num_gpus.
• zero_stage=2 (P_os+g): also shard gradients: 2 * num_params + (2 + 12) * num_params / num_gpus.
• zero_stage=3 (P_os+g+p): shard everything: (2 + 2 + 12) * num_params / num_gpus.
• Any other zero_stage raises ValueError.

Hints

Open one at a time. Each gives away a little more than the last.