Activation Memory Estimation

Activations are the intermediate outputs of each forward-pass layer; backprop reads them to compute gradients, so training must store or recompute them. Unlike the static param budget $16\Psi$, which depends only on parameter count, activation memory scales with batch size $b$ × sequence length $s$ × hidden dim $h_l$. Storing each layer's input activation (FP16, 2 bytes/element) needs $2 \cdot s \cdot b \cdot h_l$ bytes, so total activation memory ≈ $\sum_{l=1}^{L} 2 \cdot s \cdot b \cdot h_l \approx 2 L s b h$ bytes; rule of thumb: ~16 bytes per token per layer (a coarse ~8 FP16 elements per token per layer, hidden-size dependent), total ≈ $16 L s b$ bytes, with the exact value from $\sum_l 2 s b h_l$. Worked example: Llama-2 70B ($L=80$, $h=8192$), $s=4096$, $b=2$ → per layer $2 \times 4096 \times 2 \times 8192 \approx 134$ MB, 80 layers ≈ 10.7 GB — about 1% of the static param budget (1120 GB), but it grows linearly with $s \times b$ and takes over under long context and large batches. Gradient checkpointing (activation recomputation): store a checkpoint every $k$ layers instead of every layer, and rerun the forward pass during backward to regenerate intermediates; activation memory drops from $O(L)$ to $O(\sqrt{L})$ (with $k \approx \sqrt{L}$) at ~1/3 extra FLOPs — trading memory for compute; selective recomputation reruns only the biggest tensors (the $s \times s$ attention score matrix, MLP intermediates) to minimize the overhead.

a complete memory budget for 70B training must add the activation term; interview favorites are deriving the activation formula, the cost of gradient checkpointing, and activation memory vs KV cache.

counting only $16\Psi$ badly underestimates long-sequence/large-batch training (activations can exceed param memory); checkpointing cuts activation memory from $O(L)$ to $O(\sqrt{L})$ for ~33% extra compute, complementing ZeRO's param sharding — the total budget is $16\Psi/N_d$ (params) + activations (after recomputation), and together they decide whether a model fits.

Derive the activation memory formula $\sum_l 2 \cdot s \cdot b \cdot h_l$: explain each factor (2 bytes/element FP16 × layer input $h_l$ elements × $s$ positions × $b$ samples), why it is independent of parameter count, and why it grows linearly with batch×seq.

Worked example: Llama-2 70B ($L=80$, $h=8192$), $s=4096$, $b=2$ → ≈134 MB/layer, ≈10.7 GB total; put it beside the static $16\Psi = 1120$ GB and state what each budget scales with (params vs batch×seq) — the complementary relationship.

Gradient checkpointing: keep a checkpoint every $k$ layers and recompute the forward pass in backward — activation memory drops from $O(L)$ to $O(\sqrt{L})$ at ~1/3 extra FLOPs; explain the memory-for-compute trade.

Which tensors dominate activation memory (the $s \times s$ attention score matrix, $4h$ MLP intermediates, dropout masks), and how selective recomputation reruns only the big ones to minimize overhead.

Distinguish activation memory (training) from KV cache (inference): activations are forward intermediates read back by backprop, growing with $L \times s \times b$; the KV cache stores generated K/V read sequentially once, formula $2 \cdot 2 \cdot L \cdot H \cdot d \cdot s \cdot b$; why the two budgets must not be conflated.