Activation Memory Accounting
Backpropagation needs the forward pass’s intermediates. There are a lot of them, they scale with batch and sequence rather than with parameters, and one term is quadratic.
Standing
Current standardPart of the default recipe. Present in most frontier models trained today, and the thing a new design departs from rather than argues for.
The arithmetic every memory technique on this branch is trying to reduce. Worth doing on paper before reaching for any of them.
judged as of 2026-09 · what the labels mean
Theory
Training memory has four occupants. Parameters, gradients and optimiser state scale with and are fixed for a given model. Activations scale with and are the term you control.
They exist because backpropagation is the chain rule: computing for a layer requires that layer’s input, so every input consumed on the way forward must survive until the backward pass reaches it.
with sequence length, batch, hidden size, heads. Two terms, and the second is the one that matters: it carries , because the attention score matrix is and is stored per layer.
Where the 34 comes from
Worth unpacking once, because it explains which techniques pay.
| Kept | Bytes per token |
|---|---|
| layernorm input, ×2 | 4h |
| QKV input | 2h |
| Q, K, V | 6h |
| attention output | 2h |
| attention projection input | 2h |
| FFN input | 2h |
| FFN hidden, ×4 width | 8h |
| GELU input | 8h |
| dropout masks | 2h |
Nothing is individually large. There are simply about ten of them, and then layers of it.
The quadratic term
per layer. At , , that is 2.7 GB per layer, against roughly 0.6 GB for everything else — so the score matrices are most of the activation memory and all of the scaling problem.
FlashAttention deletes the term. It computes attention in tiles and never writes the matrix to HBM, storing only the log-sum-exp statistics needed to recompute it on the backward pass. Memory becomes linear in .
This is why the entry begins here rather than with recomputation: everything else on this branch is trading arithmetic for memory at some cost, and this one is free.
Implementation
def activation_memory(
layers: int, hidden: int, heads: int, seq: int, batch: int,
flash: bool = True, checkpointing: str = "none",
) -> dict[str, float]:
"""Bytes of stored activations. Korthikanti et al. eq. 2, in GB."""
linear = 34 * seq * batch * hidden
quadratic = 0 if flash else 5 * heads * seq**2 * batch
per_layer = linear + quadratic
total = {
"none": layers * per_layer,
# Full: only the block input survives; 2·s·b·h per layer.
"full": layers * 2 * seq * batch * hidden,
# Selective: recompute the attention block, keep the rest.
"selective": layers * (linear - 11 * seq * batch * hidden) + layers * quadratic * 0,
}[checkpointing]
return {"gb": total / 1e9, "per_layer_mb": per_layer / 1e6,
"quadratic_share": quadratic / per_layer if per_layer else 0}What each lever does
Having the formula makes the options comparable rather than a list of tricks.
Batch size divides everything, and micro-batching with gradient accumulation buys the reduction without changing the optimiser’s view of the batch.
Sequence length divides the linear term and quarters the quadratic one, if you still have it.
Recomputation trades the 34 for a repeated forward pass; selective recomputation targets the largest contributors and pays about a tenth of the compute for most of the saving.
Offloading moves them off the device entirely, bounded by PCIe rather than by arithmetic.
Notably absent: ZeRO, which shards parameters, gradients and optimiser state and does nothing at all for activations. The two address different occupants and are routinely confused.