ZeRO Memory Arithmetic

The problem

Before training a large model, you need to know whether it fits in GPU (graphics processing unit) memory, and with how many GPUs. You will compute it for the four stages of ZeRO (Zero Redundancy Optimizer).

With mixed-precision Adam, each parameter costs 16 bytes of model states:

• 2 bytes for its weight in FP16 (16-bit floating point);
• 2 bytes for its gradient in FP16;
• 12 bytes of optimizer states in FP32 (32-bit floating point): a master copy of the weight, the momentum and the variance, 4 bytes each.

Plain data parallelism stores all of it on every GPU. ZeRO shards it: it cuts one more piece per stage into $N$ parts, one per GPU (Rajbhandari et al., 2020).

The four stages

where $\Psi$ (psi) is the number of parameters and $N$ the number of GPUs. Each row adds the pieces that stay whole, at full size, to the pieces that are sharded, divided by $N$. At stage 2, for example, only the 2 bytes of weights stay whole, and the other $2 + 12 = 14$ bytes are sharded.

Example. $\Psi = 7 \times 10^9$ parameters on $N = 8$ GPUs:

• stage 0: $16 \cdot 7 \times 10^9 = 112$ GB (gigabytes) per GPU;
• stage 1: $28 + 10.5 = 38.5$ GB;
• stage 2: $14 + 12.25 = 26.25$ GB;
• stage 3: $112 / 8 = 14$ GB.

The other way round: the largest model that fits

Every row is $\Psi$ times a number of bytes per parameter that depends only on $N$ and the stage: $16$, $4 + 12/N$, $2 + 14/N$, $16/N$. So the largest model that fits in a given memory is the memory divided by that number, rounded down.

Example. 80 GB per GPU and $N = 8$:

• stage 0: $80 \times 10^9 / 16 = 5 \times 10^9$ parameters;
• stage 1: $80 \times 10^9 / 5.5 \approx 14.5 \times 10^9$;
• stage 3: $80 \times 10^9 / 2 = 40 \times 10^9$.

What to implement

• model_state_bytes(n_params, n_gpus, stage): the bytes per GPU, for stage in {0, 1, 2, 3}. Raise ValueError for any other stage.
• max_params(gpu_bytes, n_gpus, stage): the largest number of parameters, an int rounded down, whose model states fit in gpu_bytes per GPU. Activations are ignored here.

Hint

Write one helper that returns the bytes per parameter for a stage and $N$, and use it in both functions.