Activation Memory Math

What counts as an activation, and why autograd keeps it

An activation is any intermediate tensor the backward pass still needs, and the rule comes straight from the chain rule: for each op, ask what the gradient formula references. A matmul Y = X · W has ∂L/∂W = X^T · ∂L/∂Y, so X must survive until backward reaches that op. A nonlinearity y = σ(x) needs x (or y, if the derivative can be written in terms of the output). Dropout needs its Boolean mask. Softmax needs its output, since ∂L/∂x = (∂L/∂y − (∂L/∂y · y)) · y.

So the saved set is not “everything,” nor one tensor per layer — it is precisely the operands the gradient formulas name. Two consequences matter. First, the count is a property of the ops, not of the parameter count: a model with few parameters but wide intermediates can be activation-heavy. Second, everything saved stays alive from the moment it is produced until the backward sweep reaches it, so peak memory holds all layers at once.

Bytes: element counts times element width

Every estimate is the same two steps: count elements from the shapes, multiply by bytes per element. Notation for a decoder layer:

b = batch size        s = sequence length
h = hidden size       a = attention heads
L = layers            V = vocabulary size
d_head = h / a

hidden-state tensor : [b, s, h]      → b·s·h elements
attention scores    : [b, a, s, s]   → b·a·s² elements
MLP intermediate    : [b, s, 4h]     → 4·b·s·h elements

Element width is the quiet multiplier. In mixed-precision training the saved tensors are 2 bytes (fp16/bf16), dropout masks are 1 byte per element in most implementations (a Boolean stored as a byte, not a bit), and anything that runs in fp32 — typically the loss head — is 4. The budget is linear in this factor, so bf16 versus fp32 storage is exactly 2× on the largest line item in your run. That is why the canonical formula below is quoted directly in bytes, with fp16 storage already folded in.

Walking one transformer layer, tensor by tensor

Take a standard pre-norm block with GeLU and dropout, and list what each op hands to the backward pass:

two LayerNorms, inputs saved            4·sbh
QKV matmul input                        2·sbh
Q and K, kept for QK^T                  4·sbh
V, kept for the AV matmul               2·sbh
output-projection input                 2·sbh
attention dropout mask (1 byte)         1·sbh
softmax out + mask + dropout out        5·a·s²·b
MLP first-matmul input                  2·sbh
GeLU input        [b, s, 4h]            8·sbh
MLP second-matmul input                 8·sbh
MLP dropout mask                        1·sbh
-----------------------------------------------
A_layer ≈ 34·sbh + 5·a·s²·b bytes

This is the Korthikanti et al. (2022) accounting, and the breakdown beats the total: attention contributes 11sbh + 5as²b, the MLP 19sbh, the norms 4sbh. The MLP — usually thought of as the cheap half — is the larger linear contributor, purely because its intermediate is 4× wide and gets saved twice.

Linear versus quadratic: where the crossover sits

The two terms differ in character: the linear one grows as s, the score term as s². Setting them equal gives a crossover you can carry in your head:

5·a·s²·b = 34·s·b·h
⇒ s* = (34/5) · (h/a) = 6.8 · d_head

The batch size cancels, and so does everything except the head dimension. With the near-universal d_head = 128, the crossover is s* ≈ 870 tokens. Below that, activation memory is essentially linear in sequence length and the MLP dominates; above it, the attention score tensors take over and the curve bends upward. At s = 4096 the ratio is 5as/(34h) = s/s* ≈ 4.7, so scores are already 82% of the per-layer bill.

That single number explains why a model trains comfortably at 512 tokens and blows up at 4k for the same batch, and why long-context work is a memory problem before it is a compute problem.

The tensors outside the layer stack

Per-layer accounting misses two items at the ends of the network, and one is frequently the single largest tensor in the step. The embedding lookup is cheap — an index tensor plus the resulting [b, s, h] hidden state. The output head is not: it produces logits of shape [b, s, V], and because cross-entropy is numerically fragile, frameworks usually compute it in fp32.

Compare the shapes: a hidden state is b·s·h elements, the logits are b·s·V. With h = 2048 and V = 32000 the logits are 15× wider than a hidden state, at double the element width, and softmax typically keeps both the logits and their gradient alive at once. Hence the loss-head spikes so common in small models with large vocabularies: the model is small, but V ≫ h makes the last tensor enormous. Chunked or fused cross-entropy exists precisely to never materialize it in full.

Worked example: a 1.3B small language model

Take L = 24, h = 2048, a = 16 (d_head = 128), V = 32000, trained at s = 2048, b = 4 in bf16:

sbh   = 2048 · 4 · 2048 = 16.78e6
linear term  34·sbh              = 570 MB / layer
score  term  5·16·2048²·4       = 1342 MB / layer
A_layer      ≈ 1.91 GB
all layers   24 · 1.91 GB      ≈ 45.9 GB

logits [4, 2048, 32000] fp32   ≈ 1.05 GB (+ grad)

Sanity-check against the crossover: s/s* = 2048/870 = 2.35, and indeed 1342/570 = 2.35. Better still, normalize per token: divide by s·b to get 34h + 5as = 233 KB per token per layer, or 5.6 MB per token across the stack. Multiply by tokens-per-microbatch and you have the answer — which also shows why gradient accumulation helps (fewer tokens live at once) and a bigger microbatch does not.

Where activations sit in the total budget

The training footprint has a fixed part and a variable part. With Adam in mixed precision the fixed part is about 16 bytes per parameter: 2 fp16 weights, 2 fp16 gradient, 4 fp32 master copy, 4 each for the two moments.

fixed:  1.28e9 params · 16 B  ≈ 20.5 GB   (independent of b, s)
variable: activations           ≈ 45.9 GB   (∝ b, and ∝ s to s²)

Activations are more than twice the entire model-state budget here — and they are the part that moves. That is the structural reason they get the engineering attention: you cannot shrink the fixed part without changing or sharding the optimizer, but the variable part responds immediately to batch size, sequence length, and a few implementation choices. It is also why “it is only 1.3B parameters, it should fit” is such a reliable way to be wrong.

What FlashAttention deletes

The 5as²b term exists only because a naive implementation materializes the score matrix. FlashAttention tiles the computation, keeping a running maximum and sum (online softmax) so no [b, a, s, s] tensor is ever written to memory. Backward saves only the output and a per-row log-sum-exp statistic — O(b·a·s) instead of O(b·a·s²) — and recomputes the score tiles on the fly.

with FlashAttention:  A_layer ≈ 34·sbh   (quadratic term gone)
worked example:       24 · 570 MB ≈ 13.7 GB  (was 45.9 GB)

A second, cheaper deletion: modern LLM training usually sets dropout to zero. That drops both masks from the linear term (34 → 32) and, more importantly, cuts the quadratic term from 5as²b to about 2as²b, since only the softmax output is still needed for the value matmul.

Inference is a completely different regime

Do not carry training intuitions into serving. With no backward pass there is no autograd graph, so nothing must be saved: activation memory collapses to the working set of the executing layer, reused layer after layer.

prefill:  O(b · s · h) live, one layer at a time
decode:   O(b · h) live — one token wide
KV cache: 2 · L · b · s · h_kv · bytes
        = 2 · 24 · 4 · 2048 · 2048 · 2 B ≈ 1.61 GB

At decode time transient activations are megabytes while the KV cache is gigabytes, so the cache — not the activations — is the binding constraint, and it is what grouped-query attention and KV quantization target. Training intuition survives in exactly one place: prefill, where a long prompt processed in one shot recreates the same s² pressure inside the attention kernel. That is why chunked prefill exists.

Measuring it, and why the estimate is always low

Formulas give a lower bound; allocators report reality, and reality is consistently higher. Three reasons. Fragmentation: caching allocators reserve in blocks, so reserved exceeds allocated, and a workload with variable sequence lengths fragments badly. Transients: kernels need scratch space, and a matmul may hold input, output, and workspace at the exact instant of peak. Backward overlap: gradients start accumulating while activations are still being freed, so the true peak is usually early in the backward sweep, not at the end of the forward.

Measure the peak, not the steady state, at your real sequence length. Compare peak allocated against peak reserved: a large gap means fragmentation, which bucketing or padding fixes, whereas a high allocated peak needs a structural lever — FlashAttention, a fused loss head, gradient accumulation, or recomputation.

Pitfalls that break the accounting

Assuming the 34sbh formula transfers. It is derived for a GeLU MLP with a 4h intermediate and dropout on. A SwiGLU block with d_ff = (8/3)h saves three [b, s, d_ff] tensors instead of two [b, s, 4h] ones: 6 · (8/3) = 16sbh plus the input, landing near the same 19 — a coincidence, not a law.

Forgetting that peak holds all layers at once. Multiply the per-layer number by L.

Counting elements when the formula counts bytes. The 34 and 5 already assume 2-byte storage.

Ignoring the loss head. With a big vocabulary it can rival several transformer layers.

Benchmarking at a short sequence. Below s* = 6.8·d_head you are on the linear branch, extrapolating a curve that does not exist.