ZeRO: Memory per GPU
Intermediate
The problem, from first principles
ZeRO (Rajbhandari et al., 2020) removes the redundancy in data-parallel training. Mixed-precision Adam needs 16 bytes per parameter: fp16 weights and gradients, plus fp32 master weights, momentum and variance. ZeRO partitions those pieces across GPUs in stages.
From theory to code
Implement zero_memory_per_gpu(psi, N, stage), returning the per-GPU memory for the chosen stage.
Constraints
- Stage 0 is ordinary data parallelism with no partitioning.
Hints
Hint 1
Use the 2 bytes of fp16 weights and 2 of gradients plus 12 of optimizer state; each stage divides the parts it shards by N.
Keywords
research-paperssystemstrainingmemory
Python
Run your code to see results
No results yet
Run your code to check it against the examples, or submit to run every test.