Where the memory goes: parameters, optimizer state and activations

Running out of accelerator memory is the most common obstacle in training large models. An out-of-memory error stops a job immediately, and the techniques used to avoid it, such as recomputation, smaller micro-batches, offloading and sharding, all cost time or communication. Planning a training run therefore starts with a memory budget: how many bytes each component needs, how that scales with model size, sequence length and batch size, and which techniques trade memory for compute most cheaply. This lesson breaks training memory into its parts, derives the size of activations in a transformer, explains activation checkpointing, gradient accumulation and offloading, covers memory management inside the framework, and finishes with the key-value cache that dominates memory during inference.

The components of training memory

During training, accelerator memory holds:

• Parameters: the model weights, usually in BF16 (2 bytes each) for computation.
• Gradients: one per parameter, usually BF16 or FP32.
• Optimizer state: for Adam or AdamW, two moments per parameter in FP32 (8 bytes), plus the FP32 master weights (4 bytes) in mixed precision.
• Activations: intermediate results of the forward pass, saved because the backward pass needs them to compute gradients.
• Temporary buffers: workspace for matrix multiplication libraries, communication buffers and memory lost to fragmentation.

The first three, the model state, depend only on the parameter count: about 16 bytes per parameter with Adam in mixed precision. Activations depend on the architecture, sequence length and batch size, and for long sequences they can exceed the model state.

Optimizer choice and memory

Adam's two moments double the memory of the parameters themselves in FP32, which is a heavy price at scale. Alternatives reduce it: plain stochastic gradient descent with momentum keeps one moment; Adafactor factorizes the second moment of each weight matrix into row and column statistics; 8-bit optimizers, such as those in the bitsandbytes library, store moments in quantized form with block-wise scaling. Memory-saving optimizers can affect training dynamics, so large language models are still mostly trained with AdamW, and the optimizer state is instead divided across devices by sharding (Lesson 7).

Activation memory in a transformer

Korthikanti and colleagues at NVIDIA (2022) derived the activation memory of a standard transformer layer stored in 16-bit precision. For sequence length \(s\), micro-batch size \(b\), hidden size \(h\) and \(a\) attention heads, each layer stores about:

The first term, \(34\,sbh\), covers the inputs saved for the layer normalizations, the query, key and value projections, the attention output projection, the two MLP matrices, the activation function and dropout masks. The second term, \(5as^2b\), covers the attention score and probability matrices, which grow with the square of the sequence length. FlashAttention (Lesson 2) never stores those matrices, recomputing them block by block in the backward pass, which removes the quadratic term entirely. Even so, activations grow linearly with sequence length and batch size, and are multiplied by the number of layers.

import numpy as np
GB = 1e9
# A 7B-class decoder: hidden size h, layers L, heads a, sequence s, micro-batch b
h, L, a, s, b, vocab = 4096, 32, 32, 4096, 1, 32000
params = 12 * L * h * h + 2 * vocab * h      # attention + MLP weights, embeddings and output layer
state = 16 * params                          # BF16 weights and grads, FP32 master, Adam m and v
print(f"parameters: {params/1e9:.2f}B  training state: {state/GB:.0f} GB")
# Activation memory per layer in bytes (Korthikanti et al. 2022, 16-bit activations)
per_layer_naive = s * b * h * (34 + 5 * a * s / h)
per_layer_flash = s * b * h * 34                   # attention scores never stored
per_layer_recompute = s * b * h * 2                # keep only each layer input
for name, per in [("store everything", per_layer_naive), ("FlashAttention", per_layer_flash), ("full recomputation", per_layer_recompute)]:
    print(f"activations, {name:19s}: {L*per/GB:7.1f} GB per sequence of {s} tokens")
# Gradient accumulation: keep the global batch, shrink the micro-batch
global_tokens = 4e6
for micro in (8, 4, 1):
    steps = global_tokens / (micro * s)
    print(f"micro-batch {micro}: activations {micro*L*per_layer_flash/GB:6.1f} GB, accumulation steps {steps:.0f}")
# Inference: key-value cache for a 70B-class model with grouped-query attention
L70, kv_heads, head_dim = 80, 8, 128
def kv_cache(seq, batch, bytes_per=2):
    return 2 * L70 * kv_heads * head_dim * seq * batch * bytes_per
for seq, batch in [(4096, 1), (32768, 1), (8192, 32)]:
    print(f"KV cache, context {seq:6d}, batch {batch:3d}: {kv_cache(seq, batch)/GB:6.1f} GB")

parameters: 6.70B  training state: 107 GB
activations, store everything   :   104.2 GB per sequence of 4096 tokens
activations, FlashAttention     :    18.3 GB per sequence of 4096 tokens
activations, full recomputation :     1.1 GB per sequence of 4096 tokens
micro-batch 8: activations  146.0 GB, accumulation steps 122
micro-batch 4: activations   73.0 GB, accumulation steps 244
micro-batch 1: activations   18.3 GB, accumulation steps 977
KV cache, context   4096, batch   1:    1.3 GB
KV cache, context  32768, batch   1:   10.7 GB
KV cache, context   8192, batch  32:   85.9 GB

For a 4,096-token sequence, storing everything needs over 100 GB of activations, more than an H100 holds even before the 107 GB of model state. FlashAttention cuts activations to about 18 GB, and full recomputation to about 1 GB, at the cost of extra compute. Model state still does not fit on one GPU, which is why even a 7B model is normally trained with sharding across several devices.

Activation checkpointing

Activation checkpointing, also called gradient checkpointing or rematerialization, saves only some activations during the forward pass and recomputes the rest during the backward pass. Chen and colleagues (2016) showed that storing activations at about \(\sqrt{L}\) evenly spaced points in an \(L\)-layer network cuts activation memory from \(O(L)\) to \(O(\sqrt{L})\) for the cost of one extra forward pass. In practice, transformer training usually applies one of two policies:

• Full recomputation: save only the input to each transformer layer (about \(2sbh\) bytes) and rerun the whole layer's forward pass during backpropagation. This adds roughly one third to the compute, since the forward pass is about a third of the cost of a training step.
• Selective recomputation: save the activations that are expensive to recompute, such as matrix multiplication outputs, and recompute only the cheap but memory-hungry ones, such as attention scores and elementwise results. Korthikanti and colleagues found that this removed most of the activation memory for only a few percent of extra compute.

Because recomputation adds FLOPs that are not part of the model's useful work, it lowers MFU while raising HFU. Choosing what to recompute is a classic memory–compute trade-off, and frameworks increasingly automate the choice.

Gradient accumulation

The global batch size, the number of tokens contributing to each optimizer step, is a training hyperparameter chosen for convergence, often several million tokens for large language models. It need not be processed at once. With gradient accumulation, each device runs several micro-batches forward and backward in sequence, adding their gradients together, and only then performs the optimizer step. Activation memory depends on the micro-batch size, not the global batch, so accumulation lets a device train with a large effective batch in limited memory. The cost is less parallel work per kernel, since a smaller micro-batch produces smaller matrices, which can reduce efficiency. With data parallelism across many devices (Lesson 6), the global batch is split across devices first and accumulation covers only the remainder.

Offloading

Offloading moves data that is not needed immediately from accelerator memory to host CPU memory or even NVMe solid-state drives, which are far larger and cheaper. ZeRO-Offload and ZeRO-Infinity in Microsoft's DeepSpeed library offload optimizer state, and optionally parameters, and run the optimizer step on the CPU. Activation offloading stores saved activations in host memory during the forward pass and prefetches them for the backward pass. Offloading is limited by the link between the accelerator and the host, about 64 GB/s each way over PCIe 5.0 compared with 3.35 TB/s for HBM, so it works best when transfers overlap with computation. It is especially useful for fine-tuning large models on a small number of GPUs.

Memory management in the framework

Frameworks such as PyTorch allocate accelerator memory through a caching allocator that keeps freed blocks for reuse instead of returning them to the device, because device allocation is slow. Tensors of different sizes can fragment the cache, leaving memory that is free but unusable for a large allocation, so a job can fail with an out-of-memory error while reporting gigabytes free. Remedies include allocator settings such as expandable segments, keeping tensor shapes consistent between steps, and avoiding unnecessary copies. Memory profilers record allocations over a training step, showing peaks and their causes; the peak typically occurs at the start of the backward pass, when all saved activations are still held and gradients begin to accumulate.

The KV cache in inference

During autoregressive generation, a transformer stores the keys and values of every previous token in every layer, so that each new token attends to them without recomputing. This key-value (KV) cache grows linearly with context length and batch size:

For a 70B-class model with 80 layers, 8 key-value heads of dimension 128 and BF16 values, each token takes about 330 KB, so 32 sequences of 8,192 tokens need about 86 GB, more than the 70 GB of 8-bit weights. Architecture choices reduce it: grouped-query attention (GQA) shares key-value heads among several query heads, multi-query attention uses one, and multi-head latent attention, used in DeepSeek models, compresses keys and values into a smaller latent vector. Serving systems such as vLLM manage the cache in fixed-size pages, a technique called PagedAttention, to avoid fragmentation and share common prefixes between requests.

Memory planning in practice

A practical memory plan is built bottom up and checked by measurement. Engineers compute model state from the parameter count and optimizer, estimate activations from the formula for the planned sequence length and micro-batch, add a margin of 10 to 20 % for buffers and fragmentation, and then run a few steps while recording peak memory with the framework's memory statistics. If the measured peak differs from the estimate by more than the margin, the profile usually reveals why: an operation that keeps an unexpected copy, a layer that is not being recomputed as intended, or communication buffers larger than assumed.