Training at Scale

Memory accounting for training

For a model with $\Psi$ parameters trained with Adam in mixed precision, the standard accounting (from the ZeRO paper) is: bf16 weights ($2\Psi$ bytes), bf16 gradients ($2\Psi$ bytes), and fp32 optimizer state — a master copy of the weights plus Adam's two moment buffers ($4\Psi +4\Psi +4\Psi =12\Psi$ bytes).