Training memory: model states, activations and checkpointing

Before this lesson: lesson 5 (bytes per number format, fp32 master weights), lesson 4 (FlashAttention never stores the $s \times s$ score matrix), backpropagation, the transformer block, and Adam (two moment buffers per parameter). Time: about 70 minutes, plus the notebook.

---

Why this matters #

"Out of memory" is the first wall you hit when a model or a sequence grows. A 7B-parameter model trained with Adam needs 112 GB for its model states alone. That is more than one 80 GB H100 SXM (2022) holds, before a single activation is stored. Activations grow with the number of tokens in a micro-batch, and the attention part grows with $s^2$.

Every parallelism choice in the next lessons is first a memory decision: data parallelism with ZeRO and FSDP shards the model states, and tensor and pipeline parallelism also split the activations. In ML for science the sequences are long. A global weather field on a 0.25° grid has $721 \times 1440 \approx 1.04$ million points. Even with $8 \times 8$ patches this gives about 16,000 tokens per field, and then activations, not parameters, decide what fits.

1. Predict first #

2. The picture #

One linear layer #

Start with the simplest case: one linear layer $y = xW$. The input $x$ has $n$ rows (tokens) and $d_\text{in}$ columns; $W$ is $d_\text{in} \times d_\text{out}$. We train it with mixed-precision Adam.

• Model states belong to $W$: the weights, their gradients and the optimizer states. With mixed-precision Adam this is 16 bytes per entry of $W$, so $16\,d_\text{in}d_\text{out}$ bytes. The number of tokens does not appear.
• Activations belong to the data. The backward pass needs the weight gradient $\partial L/\partial W = x^\top\,\partial L/\partial y$. So autograd keeps $x$: $2\,n\,d_\text{in}$ bytes in bf16. It does not need to keep $y$.

Small numbers. Take $d_\text{in} = d_\text{out} = 1024$ and $n = 4096$ tokens. Model states: $16 \cdot 1024^2 \approx 16.8$ MB. Activation: $2 \cdot 4096 \cdot 1024 \approx 8.4$ MB. Double the tokens and the activation doubles to 16.8 MB; the model states stay at 16.8 MB.

The rule: every operation keeps what its backward pass needs #

Each operation's backward formula decides what it keeps. Here are the operations of a transformer layer:

operation	keeps	why the backward pass needs it
linear $y = xW$	its input $x$	$\partial L/\partial W = x^\top \partial L/\partial y$
GeLU $y = g(x)$	its input $x$	$\partial L/\partial x = g'(x) \odot \partial L/\partial y$, and $g'$ needs $x$
softmax $p = \mathrm{softmax}(z)$	its output $p$	$\partial L/\partial z = p \odot (\bar{p} - \langle p, \bar{p} \rangle)$ with $\bar{p} = \partial L/\partial p$
dropout	its mask (1 byte per element)	the gradient passes only where the mask is 1
LayerNorm	its input (and a mean and a variance per row)	the gradient depends on the normalised input
product of two activations, $QK^\top$ or $PV$	both factors	each factor's gradient needs the other factor
residual add $x + f(x)$	nothing	the gradient is copied to both branches

Memory over one training step #

Now put these tensors on a time axis. The model states are allocated once and stay. During the forward pass, each layer adds its activations on top. At the end of the forward pass all $L$ layers' activations are in memory. The backward pass then goes through the layers in reverse order. Each layer uses its activations once and frees them. So memory rises during the forward pass, peaks when the backward pass starts, and falls during the backward pass.

Interactive figure: loading…

So the peak is roughly: model states $+$ the activations of all layers $+$ the logits $+$ temporary buffers. Each term has its own fix. ZeRO and FSDP shard the model states (lesson 8). A smaller micro-batch, FlashAttention, recomputation and tensor or sequence parallelism shrink the activations.

3. Definitions and main results #

Notation (the same in the whole track): $s$ sequence length, $b$ micro-batch size in sequences, $h$ hidden size, $a$ attention heads, $L$ layers, $V$ vocabulary size, $\Psi$ number of parameters. Memory is in bytes; 1 GB $= 10^9$ bytes.

Parameters: $12h^2 + 13h$ per layer #

A GPT-style layer has biases, two LayerNorms and an MLP with hidden width $4h$.

1. Attention. $W_Q, W_K, W_V, W_O$ are each $h \times h$ with a bias of size $h$: $4h^2 + 4h$.
2. MLP. The first matrix is $h \times 4h$ with a bias of size $4h$; the second is $4h \times h$ with a bias of size $h$: $8h^2 + 5h$.
3. Two LayerNorms, each with a scale and a shift of size $h$: $4h$.

Total: $12h^2 + 13h$. The model adds the token embedding $Vh$ (often shared with the output layer), a final LayerNorm $2h$, and learned positions if there are any. For large $h$, $\Psi \approx 12Lh^2$.

A Llama-style layer has no biases, uses RMSNorm (a scale only, $h$ each) and a SwiGLU MLP with three matrices of size $h \times f$. With $f = \frac{8}{3}h$ the MLP has $3 \cdot \frac{8}{3}h^2 = 8h^2$ parameters, so the layer again has about $12h^2$. Grouped-query attention (GQA) with $g$ query heads per key/value head makes $W_K$ and $W_V$ of size $h \times h/g$. Llama 2 7B ($h = 4096$, $f = 11008$, $L = 32$, $V = 32000$, untied embeddings) gives $32 \cdot (4h^2 + 3hf + 2h) + 2Vh + h = 6{,}738{,}415{,}616$ parameters: the "7B" of its name.

Interactive figure: loading…

Model states: 16 bytes per parameter #

Mixed-precision Adam keeps five numbers per parameter (ZeRO paper, section 3.1). In the track notation the optimizer part is $K = 12$ bytes per parameter.

item	format	bytes per parameter
weights used in forward and backward	bf16	2
gradients	bf16	2
master weights	fp32	4
Adam first moment $m$	fp32	4
Adam second moment $v$	fp32	4
total		16 $= 4 + K$

So the model states take $16\Psi$ bytes: 112 GB for 7B parameters. Lesson 5 explains why the master weights and the moments are fp32. Common variants:

• PyTorch autocast with fp32 parameters also gives 16 bytes (4 weights + 4 gradients + 8 Adam), plus short-lived bf16 copies of the weights.
• Megatron-LM with fp32 gradient buffers gives 18 bytes.
• Inference needs only the bf16 weights: 2 bytes per parameter.

Activations: $sbh(34 + 5as/h)$ bytes per layer #

We use the pre-LN GPT layer of Korthikanti et al. (2022): LayerNorm, attention, dropout, residual add; LayerNorm, MLP, dropout, residual add. Inside attention, $Q$, $K$ and $V$ are split into $a$ heads. The scores $QK^\top$ of one head are an $s \times s$ matrix.

Units. A tensor with one value per token and per hidden unit has $s \cdot b \cdot h$ elements; in bf16 it takes $2sbh$ bytes. A dropout mask of the same shape takes $sbh$ bytes (1 byte per element). One score matrix per head and per sequence gives $a \cdot s^2 \cdot b$ elements in total.

Theorem (activation memory of one transformer layer). Assume a pre-LN layer with a $4h$ GeLU MLP, 16-bit activations, dropout after the softmax, after the attention projection and after the MLP, with 1-byte masks, and no FlashAttention. Then one layer keeps

Proof. List every tensor that the backward pass needs, in the order of the forward pass.

1. LayerNorm 1 keeps its input, the layer input: $2sbh$.
2. The QKV projection keeps its input, the LayerNorm output (one tensor for all three matrices): $2sbh$.
3. $QK^\top$ keeps $Q$ and $K$: $2sbh + 2sbh = 4sbh$.
4. The softmax keeps its output: $a$ score matrices of size $s \times s$ per sequence, in bf16: $2as^2b$.
5. The dropout after the softmax keeps its mask: $as^2b$.
6. The product with $V$ keeps both factors: the dropout output $2as^2b$ and $V$ $2sbh$.
7. The output projection keeps its input $2sbh$; the dropout after it keeps a mask $sbh$. The attention block so far: $11sbh + 5as^2b$.
8. The MLP: the first linear ($h \to 4h$) keeps its input $2sbh$; GeLU keeps its input, which is $4h$ wide: $8sbh$; the second linear ($4h \to h$) keeps its input, the GeLU output: $8sbh$; the dropout keeps a mask $sbh$. Together $19sbh$.
9. LayerNorm 2 keeps its input: $2sbh$. With LayerNorm 1 (step 1), the two norms give $4sbh$.
10. Sum: $(11 + 19 + 4)\,sbh = 34sbh$, plus the core $5as^2b$. Together $sbh(34 + 5as/h)$. $\blacksquare$

Not counted: the mean and variance of each LayerNorm row ($O(sb)$, very small) and the residual adds (they keep nothing).

part	items	bytes
attention, linear in $s$	LN1 input 2, QKV input 2, $Q$ 2, $K$ 2, $V$ 2, projection input 2, mask 1	$13sbh$
attention core, grows with $s^2$	softmax output 2, dropout mask 1, dropout output 2	$5as^2b$
MLP and LN2	LN2 input 2, fc1 input 2, GeLU input 8, fc2 input 8, mask 1	$21sbh$
total		$34sbh + 5as^2b$

Why the core grows with $s^2$. Everything else has one row per token, so it grows like $s \cdot b$. The score matrices have one entry per pair of tokens, so they grow like $s^2 \cdot b$. The ratio of the two parts is $5as/(34h)$. It passes 1 when $s > 34h/(5a)$: for $h = 4096$ and $a = 32$, at $s \approx 870$ tokens.

Interactive figure: loading…

What changes the count #

setting	kept per layer	cost
baseline	$sbh(34 + 5as/h)$	—
FlashAttention	$\approx 34sbh$	none; usually faster
selective recomputation	$34sbh$	recompute the attention core: about 1% more FLOPs for GPT-3
full checkpointing	$2sbh$ (the layer input)	one extra forward pass: about 33% more FLOPs
tensor parallelism, degree $t$	$sbh(10 + 24/t + 5as/(ht))$	all-reduces in every layer
tensor + sequence parallelism	$sbh(34 + 5as/h)/t$	the same communication volume as tensor parallelism

FlashAttention (lesson 4) never writes the $s \times s$ matrices. It keeps the output and one fp32 log-sum-exp value per row and per head ($4asb$ bytes, small), and rebuilds the scores in its backward pass. It also regenerates the dropout mask from the saved random-number state instead of storing it.

Selective activation recomputation (Korthikanti et al., section 5) keeps everything except the attention core, and recomputes the core from $Q$, $K$ and $V$ in the backward pass. It is the right part to recompute: it holds the most memory ($5as^2b$) but needs few FLOPs. Recomputing it costs the two matmuls $QK^\top$ and $PV$ again, $4bs^2h$ FLOPs per layer, plus elementwise work. A training step costs about $3 \cdot (24bsh^2 + 4bs^2h)$ FLOPs per layer, so the overhead is about $s/(18h + 3s)$. For GPT-3 ($a = 96$, $s = 2048$, $h = 12288$) $5as/h = 80$: the core is $80/114 = 70\%$ of the activations, and recomputing it costs about 0.9% more FLOPs.

Full activation checkpointing (also called gradient checkpointing or rematerialization) keeps only the input of each layer, $2sbh$. In the backward pass it reruns the layer's forward pass to rebuild the activations, then runs the layer's backward pass.

• Memory: $2sbhL$ for the layer inputs, plus the full activations of the one layer that is being rebuilt.
• Compute: one extra forward pass. A training step costs about $6\Psi$ FLOPs per token ($2\Psi$ forward, $4\Psi$ backward); with checkpointing it costs $8\Psi$, which is 33% more.

Interactive figure: loading…

Checkpoint every $k$ layers. Split the $L$ layers into segments of $k$ layers and keep only the input of each segment. In the backward pass, rebuild one segment at a time. Write $c$ for the bytes of one checkpoint and $A$ for the full activations of one layer. The peak is at the end of the first rebuild, when all checkpoints and one rebuilt segment are in memory:

The two terms pull in opposite directions. Setting the derivative $-Lc/k^2 + A$ to zero gives $k^* = \sqrt{Lc/A}$ and $M(k^*) = 2\sqrt{LcA}$.

• When a checkpoint costs as much as a layer's activations ($c = A$), $k^* = \sqrt{L}$ and the memory is $O(\sqrt{L})$ layers instead of $L$. This is the classic result of Chen et al. (2016) for general networks, for the price of one extra forward pass.
• In a transformer the layer input is small: $c = 2sbh$ against $A = 34sbh$ (with FlashAttention). Then $k^* = \sqrt{L/17}$, a small number: $k = 1$ gives the lowest peak for $L < 34$, and $k = 2$ for $L = 80$ ($M(2) = 148\,sbh$ against $M(1) = 194\,sbh$). So transformers checkpoint every layer or every few layers.
• The extra compute does not depend on $k$: every layer is rerun once, so the cost is always one extra forward pass.

Tensor and sequence parallelism, pipeline and data parallelism

• Tensor parallelism with degree $t$ (lesson 9) splits $Q$, $K$, $V$, the projection input and both MLP intermediates ($24sbh$), and the attention core, over $t$ GPUs. The LayerNorm inputs ($4sbh$), the inputs of the QKV and first MLP linears ($4sbh$) and the two dropout masks after the attention projection and the MLP ($2sbh$) stay on every GPU. This gives $sbh(10 + 24/t + 5as/(ht))$.
• Sequence parallelism splits those $10sbh$ along the sequence dimension $s$. Then every term is divided by $t$: $sbh(34 + 5as/h)/t$. It replaces each all-reduce by a reduce-scatter and an all-gather, so the communication volume does not change.
• Pipeline parallelism with $p$ stages and the 1F1B schedule: the first stage holds $p$ micro-batches of $L/p$ layers, which is still $L$ layers' worth of activations (Korthikanti et al., section 4.2.3).
• Data parallelism and ZeRO do not change the activations. Every GPU still runs whole layers on its own micro-batch. ZeRO divides only the model states (lesson 8).

The logits #

The output layer turns $sb$ hidden vectors into $sb$ rows of $V$ scores, the logits. The loss is computed in fp32, so the logits take $4sbV$ bytes (Korthikanti et al., section 4.3). For a large vocabulary this is large: with $s = 2048$, $b = 8$ and $V = 50257$ it is 3.3 GB. During the backward pass of the loss, the gradient of the logits has the same size for a short time.

Does it fit, and what is the cheapest fix? #

Keep 10–20% of the GPU memory free for temporary buffers, communication buffers and fragmentation (free memory split into pieces too small to use). The ZeRO paper (section 3.2) reports out-of-memory errors with more than 30% of memory still free in extreme cases.

Interactive figure: loading…

If it does not fit, estimate the time cost of each fix and take the cheapest one:

1. FlashAttention (or selective recomputation): almost free.
2. ZeRO-1 or ZeRO-2, if there is data parallelism: the same communication volume as plain data parallelism (lesson 8).
3. A smaller micro-batch with gradient accumulation: free until the matmuls get too few tokens and become memory-bound (lesson 2).
4. Full checkpointing (+33% FLOPs) or ZeRO-3 / FSDP (1.5× the communication of plain data parallelism): choose by whether the step is limited by compute or by communication.
5. Tensor + sequence parallelism inside a node: divides model states and activations by $t$, but adds communication in every layer; then pipeline parallelism across nodes.
6. Offload to CPU memory: the last resort, because the CPU link is slow.

Answers to section 1. (1) 21 GB: 16 bytes per parameter, $16 \cdot 1.3 \times 10^9 \approx 21$ GB. 2.6 GB is the bf16 weights only. (2) It grows. The $34sbh$ part depends only on the product $s \cdot b$, so it stays the same. The core $5as^2b = 5as \cdot (sb)$ doubles, because each token now attends to twice as many tokens. (3) At the start of the backward pass. At that moment all layers' activations (and the logits) are in memory; the backward pass then frees them layer by layer.

4. Worked example #

A 1.3B GPT-style model: $L = 24$, $h = 2048$, $a = 16$, $V = 50257$, $s = 2048$, $b = 8$. Training in bf16 with mixed-precision Adam, rotary positions (no learned position table). One GPU with 80 GB.

Step 1: parameters.

• Per layer: $12h^2 + 13h = 12 \cdot 4{,}194{,}304 + 13 \cdot 2048 = 50{,}331{,}648 + 26{,}624 = 50{,}358{,}272$.
• 24 layers: $24 \cdot 50{,}358{,}272 = 1{,}208{,}598{,}528$.
• Embedding (shared with the output layer): $Vh = 50257 \cdot 2048 = 102{,}926{,}336$. Final LayerNorm: $2h = 4{,}096$.
• Total: $\Psi = 1{,}311{,}528{,}960 \approx 1.31$B. Here $12Lh^2$ is only 92% of $\Psi$: at this size the embedding matters.

Step 2: model states. $16 \cdot 1.3115 \times 10^9 = 21.0$ GB.

Step 3: activations.

• $sbh = 2048 \cdot 8 \cdot 2048 = 33{,}554{,}432$, and $5as/h = 5 \cdot 16 \cdot 2048/2048 = 80$.
• Per layer: $(34 + 80) \cdot sbh = 114 \cdot 33{,}554{,}432 = 3.83$ GB.
• 24 layers: $24 \cdot 3.83 = 91.8$ GB. The attention core is $80/114 = 70\%$ of it.

Step 4: logits. $4sbV = 4 \cdot 2048 \cdot 8 \cdot 50257 = 3.29$ GB.

Step 5: total. $21.0 + 91.8 + 3.3 = 116$ GB. It does not fit on an 80 GB GPU.

Step 6: fix 1, FlashAttention. Activations $34 \cdot sbh \cdot 24 = 27.4$ GB. Total $21.0 + 27.4 + 3.3 = 51.7$ GB. It fits, with 28 GB free (35%).

Step 7: full checkpointing on top. Layer inputs $2sbh \cdot 24 = 1.6$ GB, plus one rebuilt layer $34sbh = 1.1$ GB. Total $21.0 + 1.6 + 1.1 + 3.3 = 27.0$ GB, for 33% more FLOPs. It is not needed at $b = 8$, but it allows a larger micro-batch.

Step 8: the largest micro-batch. Keep 15% free, so the budget is $0.85 \cdot 80 = 68$ GB, and $68 - 21.0 = 47$ GB is left after the model states. Activations and logits grow with $b$. Per sequence ($b = 1$):

• FlashAttention only: $34 \cdot s h \cdot L + 4sV = 3.42 + 0.41 = 3.83$ GB, so $b \le 12$.
• FlashAttention and full checkpointing: $(2L + 34) \cdot sh + 4sV = 0.34 + 0.41 = 0.76$ GB, so $b \le 62$. Now the logits take more memory than the activations.

Step 9: what if we had 8 GPUs of 40 GB with data parallelism? Plain data parallelism still needs 51.7 GB per GPU: it does not fit. ZeRO-2 (lesson 8) keeps $2\Psi + 14\Psi/8 = 2.62 + 2.30 = 4.92$ GB of model states per GPU, so $4.9 + 27.4 + 3.3 = 35.6$ GB. It fits, with only 4.4 GB (11%) free. Halving the micro-batch to $b = 4$ (two accumulation steps, same global batch) gives $4.9 + 13.7 + 1.6 = 20.3$ GB.

5. In code #

The first block checks the worked example with the formulas of section 3. NumPy only.

import numpy as np

def params_gpt(L, h, V):
    return L * (12 * h * h + 13 * h) + V * h + 2 * h        # shared embedding, final LayerNorm

def act_per_layer(s, b, h, a, flash=False):
    return s * b * h * 34 + (0 if flash else 5 * a * s * s * b)

L, h, a, V, s, b = 24, 2048, 16, 50257, 2048, 8
psi = params_gpt(L, h, V)
states, logits = 16 * psi, 4 * s * b * V
for flash in (False, True):
    act = L * act_per_layer(s, b, h, a, flash)
    print(f"flash={flash}: states {states/1e9:.1f} + act {act/1e9:.1f} + logits {logits/1e9:.1f}"
          f" = {(states + act + logits)/1e9:.1f} GB")
print(psi, round(12 * L * h * h / psi, 3))
# flash=False: states 21.0 + act 91.8 + logits 3.3 = 116.1 GB
# flash=True: states 21.0 + act 27.4 + logits 3.3 = 51.7 GB
# 1311528960 0.921

The second block lets PyTorch report what autograd keeps. A saved-tensors hook sees every tensor that is saved for the backward pass; we add up their bytes (without the parameters) for one layer in bf16, on the CPU. You do not need to run it; it shows that the theorem matches a real framework.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

s, b, h, a, p = 256, 2, 256, 4, 0.1      # sequence, micro-batch, hidden, heads, dropout

class GPTLayer(nn.Module):               # pre-LN, 4h GeLU MLP, three dropouts
    def __init__(self, flash=False):
        super().__init__()
        self.flash = flash
        self.ln1, self.ln2 = nn.LayerNorm(h), nn.LayerNorm(h)
        self.qkv, self.proj = nn.Linear(h, 3 * h), nn.Linear(h, h)
        self.fc1, self.fc2 = nn.Linear(h, 4 * h), nn.Linear(4 * h, h)

    def forward(self, x):                # x: (s, b, h)
        drop = lambda t: torch.native_dropout(t, p, True)[0]   # 1-byte mask, as on a GPU
        q, k, v = self.qkv(self.ln1(x)).split(h, dim=-1)
        q, k, v = [t.reshape(s, b * a, h // a).transpose(0, 1) for t in (q, k, v)]
        if self.flash:                   # fused attention: no s x s tensor is kept
            o = F.scaled_dot_product_attention(q, k, v)   # (attention dropout off here)
        else:
            o = drop(torch.softmax(q @ k.transpose(1, 2) / (h // a) ** 0.5, dim=-1)) @ v
        x = x + drop(self.proj(o.transpose(0, 1).reshape(s, b, h)))
        return x + drop(self.fc2(F.gelu(self.fc1(self.ln2(x)))))

def saved_bytes(layer, x, ckpt=False):
    """Bytes of the tensors autograd saves for backward (parameters not counted)."""
    params = {w.untyped_storage().data_ptr() for w in layer.parameters()}
    seen = {}
    def pack(t):
        st = t.untyped_storage()
        if st.data_ptr() not in params:
            seen[st.data_ptr()] = st.nbytes()
        return t
    with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t):
        checkpoint(layer, x, use_reentrant=False) if ckpt else layer(x)
    return sum(seen.values())

torch.manual_seed(0)
x = torch.randn(s, b, h, dtype=torch.bfloat16, requires_grad=True)
for flash, ckpt, predicted in [(False, False, 54), (True, False, 34), (False, True, 2)]:
    layer = GPTLayer(flash).to(torch.bfloat16)
    print(f"flash={flash}, checkpoint={ckpt}: {saved_bytes(layer, x, ckpt) / (s * b * h):.2f} sbh"
          f" (predicted {predicted})")
print(sum(w.numel() for w in layer.parameters()), 12 * h * h + 13 * h)
# flash=False, checkpoint=False: 54.03 sbh (predicted 54)
# flash=True, checkpoint=False: 34.09 sbh (predicted 34)
# flash=False, checkpoint=True: 2.00 sbh (predicted 2)
# 789760 789760

Here $5as/h = 5 \cdot 4 \cdot 256/256 = 20$, so the theorem predicts $34 + 20 = 54$. The small extras are the LayerNorm statistics and, with fused attention, the fp32 log-sum-exp per row. torch.native_dropout keeps a 1-byte mask, like the fused GPU dropout; F.dropout on the CPU keeps a mask in the activation's format and would add $2sbh + as^2b$ more.

6. Where it shows up #

• LLM training. Before a run, teams do this sum to choose the micro-batch, the recomputation policy (torch.utils.checkpoint, Megatron-LM's --recompute-granularity selective or full) and the ZeRO or FSDP stage. After the run starts, they measure with torch.cuda.max_memory_allocated() or a memory snapshot.
• Long context. At $s = 32{,}768$, $a = 32$, $h = 4096$, $b = 1$: $5as/h = 1280$, far above 34, so FlashAttention is required. Even $34sbh$ is 4.6 GB per layer, 146 GB for 32 layers. This needs checkpointing or context parallelism (lesson 10).
• ML for science. Weather, climate and seismic models have long token sequences and large activations but a moderate number of parameters. For them FlashAttention, checkpointing and sequence sharding matter more than ZeRO. Physics-informed losses that call autograd with create_graph=True also keep the graph of the first backward pass, so the activation memory grows further.
• Fine-tuning. With LoRA the base weights are frozen: they need no gradients and no Adam states. The model states drop to about 2 bytes per frozen parameter, but the activations stay the same, so checkpointing is still common.

7. Common mistakes #

• 2 bytes per parameter for training. That is inference. Training with mixed-precision Adam needs 16.
• Expecting data parallelism or ZeRO to cut activation memory. They only divide model states.
• Using $34 + 5as/h$ for another kind of layer. A Llama-style layer (no dropout, SwiGLU, RMSNorm) needs its own count, tensor by tensor.
• Forgetting the logits ($4sbV$, large for big vocabularies), the layer that is being rebuilt during checkpointing, and fragmentation.
• Thinking checkpointing is free. It costs about a third more FLOPs. It raises the hardware FLOPs utilization, not the model FLOPs utilization (lesson 13).
• Mixing GB ($10^9$) and GiB ($2^{30}$). They differ by 7%. GPU tools often print GiB.

8. Check yourself #

1. How many parameters does one GPT-style layer with $h = 1024$ have ($4h$ MLP, biases)?
   Answer

   $12 \cdot 1024^2 + 13 \cdot 1024 = 12{,}582{,}912 + 13{,}312 = 12{,}596{,}224$.
2. At least how many 80 GB GPUs are needed just to hold the model states of a 70B model trained with mixed-precision Adam, fully sharded?
   Answer

   $16 \cdot 70 \times 10^9 = 1.12$ TB, and $1120/80 = 14$ GPUs. In practice you need more, for activations and buffers.
3. Which $10sbh$ of the layer does tensor parallelism not split, and why?
   Answer

   The two LayerNorm inputs ($4sbh$), the inputs of the QKV and the first MLP linear ($4sbh$), and the two dropout masks after the attention projection and the MLP ($2sbh$). They lie outside the regions that tensor parallelism splits, so every GPU holds the full tensor. Sequence parallelism splits them along $s$.
4. The same layer without any dropout: how many bytes per layer?
   Answer

   $32sbh + 2as^2b$. The three masks go ($-2sbh$ and $-as^2b$). The softmax output now feeds both the softmax backward and the product with $V$, so it is stored once and the separate dropout output ($-2as^2b$) goes too.
5. Why is selective recomputation almost as good as FlashAttention for memory, and why is it cheap?
   Answer

   Both remove the $5as^2b$ core and keep $34sbh$. Rebuilding the core needs only $Q$, $K$, $V$ and few FLOPs: two batched matmuls ($4bs^2h$) and elementwise work, against $72bsh^2$ for the large linear layers over a training step. For GPT-3 this is about 1%.
6. You checkpoint a 24-layer transformer every $k = 6$ layers, with FlashAttention. Using $c = 2sbh$ and $A = 34sbh$, what is the activation peak in units of $sbh$, and is $k = 1$ better?
   Answer

   $M(6) = (24/6) \cdot 2 + 6 \cdot 34 = 8 + 204 = 212\,sbh$. $M(1) = 24 \cdot 2 + 34 = 82\,sbh$. So $k = 1$ is much better: $k^* = \sqrt{24 \cdot 2/34} \approx 1.2$.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. How many bytes per parameter does mixed-precision Adam need, and how do they split?

   Show answer

2. How many parameters does a GPT-style transformer layer with a $4h$ MLP and biases have?

   Show answer

3. What is the activation memory of one transformer layer without FlashAttention (Korthikanti et al.)?

   Show answer

4. Why does the attention core of the activations grow with $s^2$?

   Show answer

5. What does FlashAttention remove from the activation memory?

   Show answer

6. What does full activation checkpointing cost and save?

   Show answer

7. When during a training step is memory highest, and why?

   Show answer

8. With checkpoints every $k$ layers, what is the peak and the best $k$?

   Show answer

9. How much memory do the logits of the loss take?

   Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "How much memory does it take to train a 7B model with Adam?" — 16 bytes per parameter: 112 GB of model states before activations. So you need ZeRO, FSDP or tensor parallelism even on 80 GB GPUs.
• "What does activation checkpointing cost?" — One extra forward pass, about 33% more FLOPs. It keeps $2sbh$ per layer instead of $34sbh$ or more.
• "Why does FlashAttention save memory, not only time?" — It never stores the $s \times s$ matrices; it rebuilds them in the backward pass.
• "Out of memory at step 1, or after hours?" — At step 1: the estimate is wrong (logits, the micro-batch, recomputation off). After hours: fragmentation, a longer batch than usual, or a leak, such as keeping loss tensors with their graph in a Python list.

10. Practice #

The notebook for this lesson runs in Colab, with NumPy only. It has five exercises, each with a check cell and a folded solution:

1. params_per_layer(h, mlp_ratio, n_kv_ratio, style) and the model total, checked against GPT-2, the 1.3B model and Llama 2 7B and 70B.
2. model_state_bytes(psi, K), and the GPU count for a few model sizes.
3. activation_bytes(s, b, h, a, mode) for the baseline, FlashAttention, selective recomputation and full checkpointing, checked against a tensor-by-tensor list.
4. A memory-timeline simulator for one step with checkpoints every $k$ layers; find the $k$ that minimises the peak and compare it with $\sqrt{Lc/A}$.
5. cheapest_fix(...): try the fixes of section 3 in order and return the first one that fits.

11. Watch and read #

Lightning Talk: New Activation Checkpointing APIs in PyTorch - Jeffrey Wan & Horace He, MetaPyTorch Foundation · 12 min

Why checkpointing trades compute for memory, and how selective checkpointing chooses what to recompute.

EfficientML.ai Lecture 19: On-Device Training and Transfer Learning (MIT 6.5940, Fall 2023)MIT HAN Lab · Song Han · 76 min

Watch the first part, on why activations, not parameters, are the memory bottleneck of training.

Ultimate Guide To Scaling ML Models - Megatron-LM | ZeRO | DeepSpeed | Mixed PrecisionAleksa Gordić - The AI Epiphany · 82 min

A paper walk-through that connects mixed precision, the 16 bytes per parameter, and the ways to split memory over GPUs.

• Korthikanti et al. (2022), Reducing Activation Recomputation in Large Transformer Models: section 4.1 (the per-layer count), 4.2 (tensor, sequence and pipeline parallelism), 4.3 (logits), 5 (selective recomputation).
• Chen et al. (2016), Training Deep Nets with Sublinear Memory Cost: checkpointing with $O(\sqrt{n})$ memory.
• Rajbhandari et al. (2019), ZeRO: Memory Optimizations Toward Training Trillion Parameter Models: section 3.1 (model states, 16 bytes per parameter), 3.2 (activations, buffers, fragmentation).
• Austin et al., How to Scale Your Model, All the Transformer Math You Need to Know: parameter and FLOP counts, and gradient checkpointing.
• EleutherAI, Transformer Math 101: memory and compute rules of thumb for training.