Counting the bytes

A 7B model in bf16 is 14 GB. An 80 GB card should fit it five times over. It does not fit once, and the first time that happened to me I reached for gradient checkpointing without knowing what I was trading away.

The arithmetic takes ten minutes and it makes every subsequent decision obvious.

Four occupants

$$\underbrace{2}_{\text{bf16 weights}} + \underbrace{2}_{\text{grads}} + \underbrace{4}_{\text{fp32 master}} + \underbrace{4+4}_{m,\;v} = 16 \;\text{bytes}$$

For 7B that is 112 GB. Before a single activation, and the card has 80.

The piece that surprised me is the fp32 master copy. Mixed precision is advertised as halving memory and it increases weight memory — the bf16 copy is additional, not a replacement. The reason is that bf16 has seven mantissa bits, so adding a $10^{-7}$ update to a $10^{-2}$ parameter returns the parameter unchanged. The update is not rounded; it is discarded, and the parameter stops moving. The master copy exists so the optimiser has somewhere the arithmetic works.

The fifth occupant, which is the variable one

Activations. Everything the forward pass computed that the backward pass will need, and it scales with batch and sequence rather than with parameters.

$$M_{\text{layer}} = s\,b\,h\left(34 + 5\,\frac{a \cdot s}{h}\right)$$

Two terms. The 34 is about a dozen tensors of size $sbh$ — layer norm inputs, the QKV projections, the FFN hidden state and its activation input, dropout masks. None individually large, and there are a lot of them across $L$ layers.

The second term carries $s^2$, because the attention score matrix is $[b, a, s, s]$ and is stored per layer. At $s = 4096$, $a = 32$, $b = 4$ that is 2.7 GB per layer against roughly 0.6 GB for everything else.

So for a long-sequence run the answer to “where is my memory” is: the score matrices, by a factor of four or five, and everything else is noise.

Which is why the first move is free

FlashAttention never writes the score matrix to memory. It tiles the computation, keeps a running softmax normaliser, and stores only the log-sum-exp statistics needed to recompute scores during the backward pass.

The quadratic term goes to zero. Not reduced — gone, exactly, with no approximation and while being faster than the version that materialised it.

The ladder, in order of what it costs

Having the two equations turns the technique list into an ordering.

Move	Attacks	Cost
FlashAttention	the $s^2$ term	none, negative
smaller micro-batch + accumulation	all activations	wall-clock below ~8
selective recomputation	most of the 34	~4% FLOPs
8-bit optimiser states	the $4{+}4$	none measured
full recomputation	the 34	~36% FLOPs
ZeRO / FSDP	the $16$, across ranks	collectives
offload to host	anything	PCIe bandwidth

Work down it. Most runs stop after three rows.

The distinction that took me longest

ZeRO and gradient checkpointing sound like alternatives and address different occupants entirely. ZeRO shards parameters, gradients and optimiser state across data-parallel ranks — the 16 in eq. 1 — and does nothing for activations. Checkpointing addresses activations and does nothing for the 16.

A run that is out of memory because of a long sequence will not be helped by ZeRO stage 3, however much it shards. A run that is out of memory because the optimiser state does not fit will not be helped by checkpointing. Knowing which number is large tells you which of the two you need, and the numbers are two equations away.

Part 2 is about the largest single lever on that list — recomputation — and the argument for why it should never be applied uniformly.