Optimizer States Memory Requirements

Keeping a trained model in memory and training that same model are different jobs. Training has to keep track of more information while it runs, and that extra information is what makes training expensive. This article focuses on one part of that extra cost: the memory used by gradients and optimizer states, which grows with the size of the model.

The total memory needed for training is built from four pieces: the model's weights, its activations, its gradients, and its optimizer states. The weights and activations are parts you have already seen. The gradients and optimizer states are the new pieces during training, and their size is tied to how many trainable parameters the model has.

Read the diagram as a sum rather than a chain. Each of the four pieces is a separate contribution to the same total. The gradient and optimizer-state pieces are the ones this article calculates, because their size depends directly on the number of trainable parameters.

During the backward pass, every trainable parameter needs one value for its gradient. On top of that, it may need zero, one, or two extra values for optimizer states, depending on which optimizer is in use. The number of values per parameter is called V, and for the optimizers discussed here V ranges from 1 to 3.

The count of values per parameter is the key quantity. It does not describe the model's architecture or its weights. It describes the bookkeeping that the training process must hold for each parameter while it runs.

The gradient and optimizer slice of training memory is computed with a single expression. In it, M is the memory for gradients and optimizer states, P is the number of trainable parameters, V is the number of values stored per parameter, and B is the number of bytes used for each value.

Each factor plays a separate role. The parameter count $P$ sets how many things need storage. The values per parameter $V$ set how many numbers each thing needs. The bytes per value $B$ set how much space each number takes. Multiplying the three gives the size of the slice.

The optimizer changes only the value V in the expression. Vanilla SGD keeps no state, so its only per-parameter value is the gradient, which gives V = 1. A momentum optimizer keeps one state value per parameter in addition to the gradient, which gives V = 2. An Adam optimizer keeps two state values per parameter in addition to the gradient, which gives V = 3.

Notice that the optimizer's choice changes the result by a factor of three in this comparison, even though the model and the bytes per value stay the same. This is why the same model can need very different amounts of training memory depending on the optimizer.

Now keep the same optimizer and the same bytes per value, but reduce the trainable parameters to 1 billion. Only the P factor changes, so the calculation becomes 1 billion × 3 × 2 bytes, which comes to 6 GB for gradients and optimizer states.

Because memory here is proportional to the parameter count, the ratio of two results matches the ratio of the two parameter counts. Going from 13 billion to 1 billion parameters reduces this slice from 78 GB to 6 GB, roughly a thirteenfold reduction. Checking this ratio is a quick way to catch arithmetic errors.