Optimizer State Explosion: Why Your Training Memory Budget Collapses at Scale

Most engineers budget GPU memory for training by looking at the model. Count the parameters, multiply by bytes per element, add a buffer for activations. The number feels manageable. Then training starts and the cluster falls over with OOM errors before the first gradient step completes.

Photo by panumas nikhomkhai on Pexels.

The culprit, almost always, is optimizer state.

AdamW (and Adam generally) maintains two additional tensors per trainable parameter: the first moment (mean of past gradients) and the second moment (uncentered variance). Both are stored in fp32 by default, regardless of the precision you're using for weights and activations. For a 7B parameter model, that means roughly 28 GB for the weights in fp16, and then another 56 GB for the optimizer moments. You've nearly tripled your memory budget before you've allocated a single activation buffer.

This isn't a small tax. At 70B parameters, optimizer states alone occupy around 560 GB. On a node with 8×H100s (640 GB aggregate HBM), you've already consumed most of the cluster before a single forward pass.

Where the Memory Actually Lives

It helps to be precise about what's in GPU memory at any given step.

graph TD
    A[Parameter Tensor] --> B(Gradient Buffer)
    B --> C{Optimizer Step}
    C --> D[First Moment m]
    C --> E[Second Moment v]
    D --> F(Updated Weight)
    E --> F

Gradients are transient: they accumulate during the backward pass and get consumed by the optimizer step. But m and v persist for the entire training run, one full copy per parameter. If you're using gradient checkpointing to reduce activation memory, you've bought back some headroom on that side of the ledger while the optimizer states quietly hold the floor.

Mixed-precision training compounds this asymmetry. Weights live in fp16 or bf16 for the forward and backward passes, but a fp32 master copy is maintained for numerical stability during the weight update. Add that master copy to the two Adam moments, and you have four bytes times three tensors times the parameter count, all in fp32. That's 12 bytes per parameter in optimizer overhead alone.

Three Approaches Worth Knowing

ZeRO Stage 2 and Stage 3 (DeepSpeed). ZeRO-2 shards optimizer states and gradients across data-parallel ranks. ZeRO-3 goes further and shards the parameters themselves. With ZeRO-2, optimizer state memory per GPU drops proportionally to the number of data-parallel ranks; with 64 ranks, your 560 GB of Adam state becomes roughly 8.75 GB per GPU. The communication overhead is real, especially at ZeRO-3, where each forward pass must gather parameters on demand. Profile your all-gather latency before assuming ZeRO-3 is a free lunch.

8-bit Adam (bitsandbytes). Tim Dettmers' bitsandbytes library quantizes the optimizer states to 8-bit using block-wise quantization, cutting optimizer state memory roughly by 4x compared to fp32 moments with minimal accuracy impact on most tasks. As of bitsandbytes 0.41+, this is stable enough to use in production training runs. It won't help with the fp32 master weights if you're keeping them, but it makes a meaningful dent. Speculating here: the impact on loss curves for very long runs and sensitive fine-tuning tasks hasn't been thoroughly characterized in the literature, so monitor loss variance closely when you first adopt it.

Adafactor. Where Adam stores two moment tensors, Adafactor factorizes the second moment into rank-1 approximations, reducing storage from O(n) to roughly O(sqrt(n)) for matrix parameters. It eliminates the first moment entirely in its default configuration. Memory savings are substantial. The trade-off is that Adafactor is more sensitive to learning rate schedules and can require more careful warmup tuning. Use it when memory is the binding constraint and you have time to validate convergence behavior.

Practical Debugging Steps

If you're hitting unexpected OOM during optimizer initialization (not during the forward or backward pass), confirm which component is the actual ceiling. torch.cuda.memory_summary() after model load but before the first optimizer step shows your baseline. Call it again after optimizer.step() to see how much the state tensors contributed.

On multi-node runs, asymmetric OOM (one rank fails while others proceed) often signals that ZeRO sharding is misconfigured or that one rank is holding extra state due to a parameter group boundary issue. Check your param_groups configuration before assuming hardware fault.

For long runs, watch optimizer state memory over time. It should be flat after initialization. Growing optimizer memory mid-run indicates dynamic parameter registration, which happens when code incorrectly creates new parameter tensors inside the training loop rather than reusing registered buffers.

The optimizer is the most patient part of training infrastructure. It doesn't spike, doesn't generate visible throughput degradation, and doesn't produce log warnings. It simply sits there, occupying memory, until the margin runs out. Measure it early, account for it explicitly, and treat it as a first-class resident of the memory budget from day one.