The Training Memory Budget: Parameters, Optimizer State, Gradients, Activations

A model whose weights occupy 7 GB loads comfortably onto a device with 40 GB of memory and then fails during the first training step. Nothing is wrong with the loader. Training a network holds more state per parameter than inference does, and the extra state is not a small multiplier: for a conventional Adam setup in full precision it is roughly four times the weight footprint before a single activation is stored.

Getting the accounting right changes what one optimises. Most attempts to fit a larger model begin by shrinking the batch, when the batch may be the smallest term in the sum.

Four consumers, one device

Every byte resident during a training step belongs to one of four categories.

Parameters. The weights themselves. In fp32 that is 4 bytes per parameter; in fp16 or bf16, 2 bytes. A model described as having $N$ parameters costs $4N$ bytes as an fp32 checkpoint.

Gradients. One partial derivative per parameter, in the same dtype as the parameter that produced it. With fp32 parameters this doubles the parameter term.

Optimizer state. Adam keeps a first moment and a second moment per parameter, both fp32 in the standard formulation, so 8 bytes per parameter. SGD with momentum keeps one buffer, 4 bytes. SGD without momentum keeps nothing beyond the gradient. This is the term that dominates the sum for any adaptive method: for Adam it is twice the size of the weights.

Activations. Every tensor that backward needs and forward produced, retained until its gradient has been computed. This term is not proportional to parameter count at all. It scales with batch size, sequence length, and depth, and it is the term that moves when the batch size moves.

Summing the first three for fp32 Adam gives $4 + 4 + 8 = 16$ bytes per parameter. A 7-billion-parameter model therefore needs 112 GB for weights, gradients, and optimizer state alone, which is why such a model cannot be trained on one 80 GB accelerator without either sharding that state or changing the optimizer.

Why mixed precision does not halve the budget

Running the forward and backward passes in bf16 or fp16 halves the cost of the stored activations and of the weights used in the matmuls, but the optimiser update needs more precision than a 16-bit format supplies. The update is a small additive step applied to a large value; in fp16 the step is frequently below the representable resolution at that magnitude and rounds to zero. The standard remedy is a full-precision copy of the weights, often called the master weights, to which the update is applied before the result is cast back down for the next forward pass.

The consequence is that mixed precision leaves the parameter, gradient, and optimizer terms close to their fp32 sizes. A typical configuration holds bf16 weights for compute, an fp32 master copy, fp32 gradients accumulated from the bf16 backward pass, and fp32 Adam moments. That is 2 bytes of compute weights plus 4 bytes of master weights plus 4 bytes of gradients plus 8 bytes of optimizer state, or 18 bytes per parameter. Slightly more than the pure fp32 number, not less.

What mixed precision does buy is bandwidth and arithmetic throughput on the matrix multiplies, and a reduction in activation storage. The memory saving is real but it comes from the activation term, not the weight term. The number formats themselves determine how much precision that trade costs.

The activation term, which scales with the batch

Consider a transformer block applied to a batch of $B$ sequences of length $S$ with hidden width $H$. Backward needs the input to each matmul, the normalisation statistics, the attention probabilities, and the intermediate nonlinearity outputs. Per layer, ignoring constants, the dominant stored tensors are on the order of $B \cdot S \cdot H$ elements each, and there are several of them. In bf16 at 2 bytes per element, one such tensor for $B = 8$, $S = 2048$, $H = 4096$ is 134 MB. Multiply by the number of stored tensors per layer and by the number of layers and the term dominates everything else.

Attention adds a term with a different shape. The score matrix before the softmax has $B \cdot A \cdot S \cdot S$ entries for $A$ attention heads, which grows quadratically in sequence length while every other term grows linearly. Doubling $S$ doubles the per-layer hidden-state storage and quadruples the score storage. This is the reason long-context training fails at batch sizes that work fine at short context, and the reason attention kernels that avoid materialising the full score matrix exist at all.

The practical rule is that activation memory is approximately linear in batch size and quadratic in sequence length, with a coefficient set by depth and width. Halving the batch halves it. Halving the sequence length quarters the attention part and halves the rest. Neither operation touches the parameter, gradient, or optimizer terms.

Peak is not the sum of the averages

A step does not hold everything at once for its whole duration. Forward stores activations as it goes, so memory rises monotonically through the forward pass. Backward frees each layer’s stored tensors as soon as that layer’s gradient is computed, so memory falls through the backward pass while the gradient buffers fill. The peak typically occurs at the boundary between the two passes, when the last layer’s activations are still live and the first gradients have already been allocated.

Two consequences follow. First, measuring memory after the step returns a number well below the peak, because the caching allocator retains freed blocks in its pool. Framework memory summaries usually distinguish allocated from reserved bytes for this reason, and the reserved figure is the one that determines whether the next allocation succeeds.

Second, fragmentation matters. The allocator serves a request from a cached block only if a block of sufficient size is free; a pool holding many small free blocks cannot satisfy one large request even when total free bytes are ample. Growing batch sizes and variable sequence lengths produce exactly this pattern, and the failure appears as an out-of-memory error at a peak that arithmetic says should fit.

Working the budget backwards

The useful form of the equation is a residual. Take the device capacity, subtract the framework and kernel workspace overhead, subtract the parameter, gradient, and optimizer terms, and what remains is the activation allowance. Dividing that by the per-sample activation cost gives the largest batch that fits.

def state_bytes(n_params, master_fp32=True, moments=2, grad_bytes=4):
    """Resident bytes for parameters, gradients and Adam moments, per replica."""
    compute = 2 if master_fp32 else 4      # bf16 copy used by the matmuls
    total = compute * n_params
    total += (4 if master_fp32 else 0) * n_params   # fp32 master weights
    total += grad_bytes * n_params
    total += 4 * moments * n_params                  # first and second moment
    return total

def max_batch(capacity, n_params, per_sample_bytes, overhead=2 << 30):
    free = capacity - overhead - state_bytes(n_params)
    return free // per_sample_bytes if free > 0 else 0

The state_bytes result for a model with 7 billion parameters is about 126 GB, which exceeds any single device and says immediately that the batch size is not the variable to adjust.

When the residual is negative, the order of remedies is dictated by which term is large. If optimizer state dominates, shard it across devices, as described in ZeRO stage partitioning; this reduces the per-device state by the data-parallel degree with no change to the model. If activations dominate, either reduce the batch and compensate with gradient accumulation or trade compute for memory with activation checkpointing. If parameters alone exceed capacity, the model must be split across devices, which is a different problem with a different cost structure.

The mistake to avoid is treating “out of memory” as a request to shrink the batch. In a model with a large optimizer state, shrinking the batch from 32 to 16 may recover a few gigabytes while leaving a hundred gigabytes of state untouched. The arithmetic above, done once with the actual dtype choices, identifies which lever has any effect.