Training Memory & ZeRO

The problem, from first principles

Training a model needs far more memory than the weights alone. With mixed-precision Adam the usual per-parameter bill is 2 bytes for the bf16 weights, 2 bytes for the bf16 gradients and 12 bytes of optimizer state (a 4-byte fp32 master copy of the weights plus two 4-byte moment estimates): 16 bytes per parameter. A 7-billion-parameter model therefore needs over 100 GB before counting a single activation. ZeRO removes the redundancy of data parallelism by sharding these three components across n GPUs: stage 1 shards the optimizer state, stage 2 also shards the gradients and stage 3 also shards the weights.

From theory to code

Implement bytes_per_param and training_memory_gib.

Constraints

• The three components are weights 2, gradients 2 and optimizer state 12 bytes per parameter.
• bytes_per_param(zero_stage, n_gpus) returns the bytes per parameter held by one GPU. Stage 0 shards nothing. Stage 1 divides the optimizer state by n_gpus. Stage 2 also divides the gradients. Stage 3 also divides the weights. Any other stage raises ValueError.
• training_memory_gib(n_params, zero_stage, n_gpus) is n_params * bytes_per_param(...) / 2**30, ignoring activations.
• n_gpus >= 1.

Hints