Optimiser State and Memory

A 7-billion-parameter model stored in fp32 is about 28 GB. Training it with plain Adam and fp32 master weights needs far more than that: the parameters themselves (4 bytes each), the gradients (4 bytes each), and Adam's two moving-average tensors, m and v (4 bytes each), the same shape as the model, held per parameter. That is roughly 16 bytes of state per parameter beyond the weights, close to 112 GB for the same 7B model, before a single activation from a forward pass has been stored. "Can I fit this model" is, in practice, almost always "can I fit this model's optimiser state," and that distinction shapes real hardware decisions long before training starts.

Where each byte goes

Under mixed-precision training, the forward and backward passes run in a lower-precision format like bf16, but an fp32 master copy of the weights is kept and updated by the optimiser, then cast back down to bf16 for the next round of matmuls. Micikevicius et al.'s mixed-precision training recipe established this pattern: bf16 has too little mantissa precision to accumulate small parameter updates reliably on its own, so the update happens in fp32 even while the bulk of the compute runs in a cheaper format (see mixed-precision-bf16-fp8 for the full precision picture). Adam's m and v tensors are conventionally kept in fp32 as well, for the same stability reason, adding two more full-precision tensors per parameter on top of the master weights and gradients.

Cutting it down, three real levers

Sharding. ZeRO and FSDP split the optimiser state, and optionally the gradients and parameters themselves, across data-parallel ranks, so no single GPU ever holds the full state for the whole model, only its own shard. Rajbhandari et al.'s ZeRO paper turns what was a fixed per-GPU memory ceiling into a bandwidth problem instead, see zero-fsdp-sharding for the full mechanism.

Cheaper optimisers. Adafactor (Shazeer and Stern) factorises the second-moment tensor into row and column vectors instead of storing a full matrix, cutting that piece of memory sharply at the cost of a coarser approximation. Lion drops the second moment entirely and tracks only momentum, halving optimiser memory outright (see the tradeoffs in optimisers-adam-adamw-lion).

Quantised optimiser states. Dettmers et al.'s 8-bit optimizers keep Adam's m and v in int8 with per-block scaling factors rather than fp32, shrinking roughly 8 bytes per parameter of moment storage down to around 2, with the scaling handled transparently enough that the update math still behaves like full-precision Adam most of the time.

Why this is infrastructure, not a micro-optimisation

Optimiser state, not weights, is frequently the binding constraint on how many GPUs a given model requires, and by extension on cost and on which sharding or offload strategy is worth the engineering effort. A team that only budgets for weight memory will discover the real constraint the first time they try to actually start training. Memory-frugal optimisers and state sharding are decisions made at the same level as choosing a parallelism strategy, not afterthoughts bolted on once a run is already slow to fit.

When it falls down

• Sharding trades memory for communication. ZeRO stage 2 and 3 remove the fixed per-GPU memory ceiling but add all-gather and reduce-scatter traffic on every step; on a slower interconnect this can make a run bandwidth-bound instead of memory-bound.
• Quantised states are not a strict free lunch. 8-bit moments introduce quantisation noise into the momentum and variance estimates; convergence is usually close to full precision but not always identical, and the difference is worth checking rather than assuming away.
• Factorisation is an approximation with real failure cases. Adafactor's row/column factorisation of the second moment assumes it is well approximated by a rank-1 outer product, which fits some parameter shapes far better than others.
• The tricks are rarely stress-tested in combination. Running 8-bit states, sharding, and a memory-cheap optimiser together at once is a much less common configuration than any one of them alone, and papers demonstrating each individually do not guarantee they interact cleanly together.

Further reading

• Micikevicius et al., 2017, Mixed Precision Training, arXiv:1710.03740 - the fp32 master-weight pattern.
• Rajbhandari et al., 2019, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, arXiv:1910.02054 - sharded optimiser state.
• Shazeer and Stern, 2018, Adafactor: Adaptive Learning Rates with Sublinear Memory Cost, arXiv:1804.04235 - factorised second moments.
• Dettmers et al., 2022, 8-bit Optimizers via Block-wise Quantization, arXiv:2110.02861 - quantised Adam moments.