What are optimizer states?

In this article

1. Before you start
2. Why interviewers ask this
3. Optimizer states from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page belongs to the deep learning guide, which places the optimizer between the loss and the weights.

Why interviewers ask this

A candidate who knows only that Adam is "a better stochastic gradient descent (SGD)" cannot explain why a model that fits on one graphics processing unit (GPU) for inference runs out of memory the moment training starts. Optimizer states are the usual answer, and the topic comes up as prompts such as "how much GPU memory is needed to train a large language model (LLM)", "what does momentum change" and "which ZeRO stage and why".

A weak answer says "Adam stores momentum". A strong answer counts the bytes: weights, gradients, two moments and, under mixed precision, a full-precision copy of the weights, and then names what can be shrunk.

Optimizer states from first principles

In plain words

Optimizer state

An optimizer state is a number an optimizer stores for each model weight between training steps so that the next update can use the history of earlier gradients. Plain SGD stores none; Adam stores two per weight.

First moment

The first moment is a running average of a weight's recent gradients, kept so that an update follows the direction the gradient has pointed in lately instead of reacting to one noisy batch.

Second moment

The second moment is a running average of a weight's recent squared gradients, kept so that an update can be scaled by how large that weight's gradients typically are.

Master weights

Master weights are a full-precision fp32 copy of the model weights that mixed-precision training keeps, because the update added each step is often too small to register in an fp16 number.

Picture a hiker crossing hilly ground in fog. A gradient is what the hiker feels under one foot right now. Optimizer states are the hiker's notes: which way the last few steps went, and how bumpy each stretch was. Notes let the hiker keep going in a steady direction and take smaller steps where the ground is rough. The analogy breaks in one place. A hiker keeps one set of notes, while an optimizer keeps its own notes for every one of billions of weights, which is why they cost so much memory.

In one precise sentence, optimizer states are per-weight arrays of numbers, the same shape as the weights, that an optimizer updates at every step. Everything below is about what is in them and what they cost.

Step 1: Each weight gets a gradient

Backpropagation is the procedure that runs the model backward from the loss to compute how each weight affected it, and it produces one gradient for every weight, and plain gradient descent subtracts the learning rate times that gradient. Nothing is remembered between steps, so plain SGD needs no optimizer state. The input is the loss on one batch (a small group of training examples), and the output is one gradient number per weight. We follow one weight with an illustrative first gradient of 0.5¹.

The weakness is noise. A single batch is a small sample, so its gradient jumps around, and a weight whose gradient flips sign from step to step makes little progress. The next two steps are the two pieces of memory that fix that.

Step 2: The first moment remembers the direction

The first moment is an exponential moving average of the gradients: each step it keeps 90 percent of the old value and adds 10 percent of the new gradient, using Adam's default ${\beta }_1=0.9$. The PyTorch Adam reference writes it as $m_t={\beta }_1{m}_{t-1}+(1-{\beta }_1)g_t$, starting from zero. Kingma and Ba (2015) introduced the method.

With a first gradient of 0.5 the first moment becomes $m_1=0.9\times 0+0.1\times 0.5=0.05$¹. This is the same idea as classical momentum. In PyTorch, SGD with a momentum setting keeps one buffer per weight, the momentum_buffer, and the SGD reference shows that PyTorch keeps that buffer only when momentum is not 0, and the default is 0.

The input is the new gradient and the old first moment. The output is a smoothed direction, and it is the first of Adam's two stored arrays.

Step 3: The second moment sets each weight's step size

The second moment is a moving average of the squared gradient, $v_t={\beta }_2{v}_{t-1}+(1-{\beta }_2)g_t^2$ with the default ${\beta }_2=0.999$ (PyTorch Adam reference), and the update divides by its square root. Dividing by the typical size of a weight's gradient means every weight moves by roughly the learning rate, whether its gradients are tiny or huge. A weight with large, noisy gradients takes small steps, and one with rare small gradients takes relatively larger ones.

For the first gradient of 0.5 the second moment is $v_1=0.001\times 0.25=0.00025$¹. It needs a second array of the same size as the weights, and that is exactly why Adam costs more memory than momentum: two states per weight against one.

The input is the squared gradient. The output is a running estimate of gradient size, stored as Adam's second array.

Step 4: Bias correction and the update

Both moments start at zero, so early values are biased low, and Adam divides each by one minus its decay raised to the step number. This is bias correction, a division that rescales early averages to the size a long-run average would have. In the PyTorch pseudocode the corrected values are ${\hat{m}}_t=m_t/(1-{\beta }_1^t)$ and ${\hat{v}}_t=v_t/(1-{\beta }_2^t)$, and the weight moves by the learning rate times ${\hat{m}}_t/(\sqrt{{\hat{v}}_t}+\epsilon )$, where epsilon is a tiny constant (default $\epsilon ={10}^{-8}$) that keeps the division from hitting zero.

On step 1 the corrected values are 0.5 and 0.25, so with an illustrative learning rate of 0.1 the weight moves by $\mathrm{\Delta }=0.1\times 0.5/0.5=0.1$. Without the correction the step would be 0.316, three times too large¹. Bias correction costs no memory, since it uses only the step counter.

1 · Gradient2 · First moment3 · Second moment4 · Correct and stepReplay

Step-by-step description

1. Backpropagation gives one gradient per weight.
2. State one remembers the direction.
3. State two remembers the size.
4. Bias correction then the update.

• The first moment after the first gradient of 0.5 is 0.05.
• The second moment after the first gradient is 0.00025.
• After bias correction the first step is 0.1.

Worked example

The memory bill for training is the sum of weights, gradients and optimizer states, and the optimizer states are the largest part. The Zero Redundancy Optimizer (ZeRO) paper, Rajbhandari et al., counts it for mixed-precision training of a model with $\mathrm{\Psi }$ parameters using Adam: fp16 weights take two bytes per parameter, fp16 gradients take two more, and the optimizer states are an fp32 copy of the parameters, the momentum and the variance at four bytes each, a multiplier it calls $K=12$. The total is $B=2+2+K=16$ bytes per parameter, or $B\mathrm{\Psi }$ for the whole model¹.

The totals use decimal gigabytes, one billion bytes each, and 7 billion parameters is an illustrative size¹. The ZeRO paper's own example checks the formula: a 1.5 billion parameter generative pre-trained transformer (GPT) model, GPT-2, needs at least 24 GB, against 3 GB for its fp16 weights alone¹. Rows 3 and 4 are equal on purpose: fp32 Adam stores 4 bytes of weights, 4 of gradients and 8 of moments, while mixed-precision Adam stores 2 and 2 but adds a 4-byte fp32 master copy to its 12 bytes of states, so the 4 bytes saved on weights and gradients are paid back by the master copy¹.

Mixed precision is used anyway because Micikevicius et al. report that it cuts memory consumption by nearly 2x overall and speeds up arithmetic; by the ZeRO accounting that saving cannot come from weights, gradients and states, so it comes from other items such as activations stored in half precision. Activations, the intermediate values each layer produces and the backward pass needs, come on top of every row.

The last row is our arithmetic from two sources. The 8-bit optimizers paper reports that 32-bit momentum and Adam use 4 and 8 bytes per parameter, which is 4 GB and 8 GB for a model with one billion parameters, and that 8-bit quantisation brings that to 1 GB and 2 GB. It also states that optimizer states use 33 to 75 percent of the total memory footprint during training. Keeping the fp32 master copy of 4 bytes and adding 2 bytes of 8-bit states gives 6 optimizer bytes per parameter in the table¹.

Text description of the diagram

• Mixed-precision Adam needs 4 bytes for weights and gradients and 12 bytes for optimizer states.
• Plain SGD keeps no optimizer state.

Trade-offs and edge cases

More state buys faster, steadier training, and the price is memory that scales with the weight count. Three levers cut it, and each trades something.

Text description of the diagram

• With 8 workers, sharding only the optimizer states leaves 5.5 bytes per parameter on each GPU.
• Sharding states, gradients and weights leaves 2 bytes per parameter on each GPU.

First, choose a lighter optimizer. SGD with momentum halves the state of Adam. Second, shrink the states: the 8-bit optimizers paper reports that block-wise 8-bit statistics keep 32-bit performance on the tasks it tested, with states a quarter of the size. Third, shard them: ZeRO partitions the optimizer states across data-parallel workers, the GPUs that each train on a different slice of the data, so each holds a fraction instead of a full copy. The total byte count is unchanged and is divided across the machines.

The chart redraws the three stages of the ZeRO paper's Figure 1 (drawn there for 64 GPUs) for 8 workers: sharding only the states leaves 4 + 12/8 = 5.5 bytes per parameter on each GPU, sharding the gradients too leaves 2 + 14/8 = 3.75, and sharding the weights as well leaves 16/8 = 2, which is 38.5, 26.25 and 14 GB for 7 billion parameters¹. On communication, the ZeRO paper reports that sharding the states and gradients adds no extra communication volume over standard data parallelism, and that only sharding the weights as well costs more (about 1.5 times).

The master copy is an edge case worth knowing. The mixed precision paper keeps a single-precision copy of the weights that accumulates the gradients after each optimizer step, because half precision cannot represent small updates to large weights. Training in the 16-bit bf16 or 8-bit fp8 formats changes the arithmetic; see the fp16 versus bf16 and fp8 training pages for the formats.

A further case is parameter-efficient training. When only a small adapter is trained, as in low-rank adaptation (LoRA), covered on the fine-tuning with LoRA page, the optimizer keeps states only for the trainable weights, so the states shrink with them. The frozen weights still occupy memory at their own precision; quantisation is the lever for that part.

How to answer it in an interview

The spoken structure is: list what training stores, give the byte count, then name the levers.

Common mistakes

• "Optimizer states are the model weights" (wrong). They are extra arrays the optimizer keeps in addition to the weights (right: weights, gradients and states are three separate costs).
• "SGD and Adam use the same memory" (wrong). Plain SGD keeps no state and Adam keeps two arrays per weight (right: count the states per optimizer).
• "Mixed precision halves the memory of weights, gradients and Adam states" (wrong). In the ZeRO accounting both fp32 Adam and mixed-precision Adam cost 16 bytes per parameter for those three items, because the fp32 master copy cancels the saving (right: the roughly 2x saving Micikevicius et al. report comes from storing activations in half precision, not from model states).
• "Bias correction is a memory cost" (wrong). It uses only the step number, so it adds no per-weight state (right).
• "Activations are optimizer states" (wrong). Activations are intermediate values kept for the backward pass, a separate cost that comes on top (right).

Follow-up questions

How much GPU memory is needed to train an LLM?

At least 16 bytes per parameter for mixed-precision Adam before activations, so about 112 GB for a 7 billion parameter model, of which 84 GB is optimizer states. Activations and any batch-size-dependent buffers come on top, so the real figure is higher¹.

What does momentum change?

Momentum adds one stored buffer per weight, a running average of past gradients, so each update follows the recent direction instead of one noisy batch. In PyTorch SGD it is off by default (momentum 0), and turning it on costs 4 bytes per weight in fp32.

Which ZeRO stage, and why?

It depends on what no longer fits. Sharding the optimizer states alone takes a 7 billion parameter model from 112 to 38.5 GB per GPU on 8 GPUs, because the states are the largest item; adding gradients gives 26.25 GB and adding the weights gives 14 GB, and the ZeRO paper reports that only the weight-sharding stage adds communication, about 1.5 times the baseline volume, while the first two add none¹. Stages 1 and 2 add no communication volume, so stage 2 is the usual default; stage 3 is worth its 1.5 times traffic only when 3.75 bytes per parameter per GPU still does not fit.

What changes in the update if the gradient is much larger than usual?

Far less than the gradient would suggest: the step stays at a few times the learning rate instead of growing with the spike. Kingma and Ba (section 2.1) state an upper bound of $(1-{\beta }_1)/\sqrt{1-{\beta }_2}\approx 3.16$ times the learning rate for the default betas, for the most extreme sparse case: a gradient that was zero at every earlier step. That bound is not exact for every history: in our computation, a 100 times spike after steady gradients of 1 moved the weight by 3.29 times the learning rate, while plain SGD would have taken a 100 times larger step¹. Plain SGD has no such limit, because its step is the gradient times the learning rate; see gradient descent.

Sources

1. Adam: A Method for Stochastic Optimization (arXiv 1412.6980v9) — arXiv (Kingma and Ba), 30 Jan 2017
2. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (arXiv 1910.02054v3) — arXiv (Rajbhandari et al.), 13 May 2020
3. torch.optim.Adam, PyTorch 2.14 documentation — PyTorch (page shows version only; read 2026-10-24), 24 Oct 2026
4. torch.optim.SGD, PyTorch 2.14 documentation — PyTorch (page shows version only; read 2026-10-24), 24 Oct 2026
5. 8-bit Optimizers via Block-wise Quantization (arXiv 2110.02861v2) — arXiv (Dettmers et al.), 20 Jun 2022
6. Mixed Precision Training (arXiv 1710.03740v3) — arXiv (Micikevicius et al.), 15 Feb 2018