Does mixed precision training really reduce memory usage?
"TLDR: This article explores whether mixed precision training (AMP) truly reduces GPU memory usage. By analyzing the storage requirements of model parameters, gradients, optimizer states, and activations under FP32 and FP16, it points out that although forward inference halves memory usage, backpropagation requires converting back to FP32 due to numerical range limitations and retaining a copy of the original FP32 model, resulting in total memory usage becoming 1.5 times the original. However, in actual training, gradients and activations are the primary overhead, and comprehensive calculations show that mixed precision can still reduce overall memory consumption (e.g., a 7B model dropping from 112G to 84G) while also improving computational speed."
Mixed precision training (AMP) refers to converting model parameters from fp32 to fp16 for computation. Since the bit width is halved, GPU operations run faster, which can accelerate model training. But does mixed precision training always reduce memory usage?
The parameters involved in model training generally fall into the following categories:
- The model parameters themselves (e.g., a 7B model at fp32 precision requires 28 GB of GPU memory just to load)
- Gradients (every parameter being optimized carries a gradient; if all parameters are optimized, this is equivalent in size to the model parameters, also 28 GB)
- Optimizer states (some optimizers, to make gradient updates more accurate and faster, compute additional momentum terms, such as AdamW, which adds another 2 * gradient size, i.e., 56 GB of memory)
- Activations — during backpropagation, computing gradients requires the activation values output by the current node during the forward pass. This part is highly flexible and depends on batch size, number of layers, etc.
Forward Pass
With model parameters converted to fp16, all memory usage during the forward pass is halved (gradients, extra optimizer states, and activations).
Backward Pass
Since backpropagation involves multiplication and division of gradients during chain-rule derivation, it is very easy to overflow in fp16 (fp16 has a relatively small representable range).
Therefore, during gradient backpropagation, in order to represent a wider range of values, fp16 is converted back to fp32.
However, since the model parameters are in fp16, to avoid losing precision in computation, a copy of the original fp32 model parameters must be retained.
As a result, mixed precision causes the model parameters to be stored twice, increasing memory usage to of the original (a 7B model now requires 42 GB).
Does Memory Usage Go Down or Up?
During model training, gradients and activations are the main memory consumers. For example, a 7B model (28 GB) with an fp32 optimizer takes 28 + 56 = 84 GB, totaling 112 GB.
Under mixed precision, the model parameters take 42 GB and the optimizer takes 42 GB, totaling 84 GB.
Therefore, mixed precision can reduce memory usage, and fp16 computation is also faster.