#zero

A model with 2 billion parameters is trained with mixed-precision Adam on 8 GPUs. Per parameter it needs 2 bytes for the fp16 weights, 2 bytes for the fp16 gradients and 12 bytes of optimizer state (fp32 master weights, momentum and variance), so plain data parallelism stores 32 GB of these model states on every GPU. ZeRO stage 2 splits the gradients and the optimizer states evenly across the 8 GPUs, while every GPU still keeps all the fp16 weights. How many GB ($10^9$ bytes) of model states does each GPU hold?