Memory Required for Gradients and Optimizer States

Before you run a large model, you need an estimate of how much memory it will occupy. Large language models have billions of parameters, and each parameter has to be held in memory while the model works. A good estimate tells you whether a model will run at all, and whether you need more hardware or a smaller model.

The good news is that the estimate is simple arithmetic built from two numbers: how many parameters the model has, and how many bytes each parameter takes to store. This article builds that arithmetic step by step, first for inference (running the model) and then for training (updating its parameters).

A model's weights are stored as a long list of numbers, one entry per parameter. The storage cost of the weights is therefore the number of parameters multiplied by the number of bytes used to store each one. The result is usually expressed in gigabytes, which is why the parameter count is written with a B for billions.

In this formula, N is the number of parameters and b is the number of bytes each parameter occupies. The choice of b matters a great deal. Halving b halves the memory needed for the weights, so the same model can have very different footprints depending on how its values are stored.

The formula is linear in the parameter count, so doubling the model roughly doubles its weight memory. Consider a 70B-parameter model with 2 bytes per parameter. Its weights alone need 70B × 2 bytes = 140 GB, which is far more than the 26 GB needed by the 13B model. Memory quickly becomes a bottleneck when operating large models.

Inference only needs the weights and their overhead. To estimate the total inference memory, you take the weight memory and multiply it by an overhead factor. The source uses an overhead factor of 1.2 in its example.

Training is different from inference. When you update the parameters, the system must also keep information about each trainable parameter, in addition to the weights. The source uses the Adam optimizer as its example. For each trainable parameter, three values are stored for its gradient and optimizer states.

If each of those three values takes 2 bytes, then each trainable parameter needs 3 × 2 = 6 bytes of additional storage on top of its weight. Written as a formula, the gradient and optimizer storage depends on three quantities: the number of trainable parameters, the number of values stored per parameter, and the bytes per value.

Here N_{trainable} is the number of trainable parameters, k is the number of values stored per parameter for gradients and optimizer states (three in the Adam example), and b is the number of bytes per value. This term is separate from the weight memory calculated earlier, so training memory is the weight memory plus this gradient and optimizer term.

The number of trainable parameters is the lever that controls gradient and optimizer memory. The same per-value rule applies whatever the trainable count is. If only 1B of the parameters are trainable, the calculation gives 1B × 3 × 2 bytes = 6 GB for gradients and optimizer states. That is a large reduction from the 78 GB needed when the full 13B model trains.

The source states this effect qualitatively: reducing trainable parameters from 13B to 1B significantly reduces the memory needed for gradients and optimizer states. The calculation above gives the size of that reduction.