Adam Optimizer › Memory Implications for Large Model Training

At large scale, Adam's memory requirements become a significant engineering constraint. Consider a language model with $N$ parameters. A standard training setup in float32 requires:

• $N$ floats for the model parameters themselves
• $N$ floats for the parameter gradients (computed during backpropagation)
• $2N$ floats for the Adam first and second moments

The total is $4N$ floats, four times the model's parameter count. For float32 training, each float is 4 bytes, so a 7 billion parameter model requires $4 \times 7 \times 10^9 \times 4 = 112$ GB of memory just for the training state. This far exceeds the capacity of any single GPU (which typically has 40-80 GB of HBM memory), requiring distributed training across many devices.

Several strategies reduce this memory footprint:

ZeRO (Zero Redundancy Optimizer) from Microsoft DeepSpeed shards the optimizer state, gradients, and parameters across all GPUs in a distributed training job. In the most aggressive mode (ZeRO-3), each GPU holds only $1/N_{\text{GPU}}$ of the model weights, gradients, and Adam moments. The memory per GPU scales inversely with the number of GPUs, making it possible to train models far larger than any single GPU's capacity.

8-bit Adam (Dettmers et al., 2022, available via the bitsandbytes library) quantizes the optimizer states to 8-bit integers, reducing the 8 bytes per parameter for Adam state (two float32 tensors) to 2 bytes per parameter. The quantization is done with a dynamic block-wise scheme that preserves the optimizer's behavior almost exactly, with negligible accuracy loss. This is perhaps the most practical approach for single-GPU training of large models.

Mixed precision training uses float16 for the forward and backward passes (saving memory on activations and intermediate computations) while keeping the optimizer state and master copy of the weights in float32. The Adam step itself is computed in float32 and then cast back to float16 for the next forward pass. This is the standard approach for training transformers on modern hardware and roughly halves memory usage compared to full float32 training.