Memory Budget Verification

After choosing (TP, PP, DP), you must verify that the per-GPU memory actually fits. This chapter provides the exact formulas and a systematic method to add up every memory component.

1. The Complete Training Memory Ledger

Training memory has two parts: static (independent of batch size) and dynamic (scales with batch/sequence).

Static Memory: 16Ψ per Parameter

Memory Ledger

With BF16 mixed precision + AdamW:

Component	Precision	Bytes/param	7B model	70B model
Parameters (working copy)	BF16	2	14 GB	140 GB
Gradients	BF16	2	14 GB	140 GB
FP32 master weights	FP32	4	28 GB	280 GB
Adam 1st moment ( $m$ )	FP32	4	28 GB	280 GB
Adam 2nd moment ( $v$ )	FP32	4	28 GB	280 GB
Static total		16	112 GB	1,120 GB

Why 16 and not 18? Some implementations keep an additional FP32 gradient accumulator (extra 4Ψ → 20Ψ). The canonical BF16 mixed-precision setup uses 16Ψ (BF16 gradients, FP32 optimizer states). Our Envelope Estimation chapter uses 18Ψ for a more conservative estimate. Both are valid — just be consistent.

How Parallelism Reduces Static Memory

Strategy	Parameters per GPU	Optimizer per GPU	Gradients per GPU
No parallelism	Ψ	Ψ	Ψ
TP=T	Ψ/T	Ψ/T	Ψ/T
PP=P	Ψ/P (layers assigned)	Ψ/P	Ψ/P
FSDP (ZeRO-3)	Ψ/D (gathered on demand)	Ψ/D	Ψ/D
TP=T + PP=P + FSDP(D)	Ψ/(T×P×D)	Ψ/(T×P×D)	Ψ/(T×P×D)

FSDP is the most memory-efficient for static memory because it shards everything across all DP ranks, gathering parameters on-the-fly during forward/backward.

2. Dynamic Memory: Activations

Activations are the intermediate tensors saved during forward pass for use in backward. This is the “variable cost” that depends on batch size and sequence length.

Per-Layer Activation Memory (Transformer)

For one Transformer layer with hidden dimension $h$ , sequence length $s$ , batch size $b$ , attention heads $a$ :

Component	Shape	Bytes (BF16)
Input to attention (Q, K, V)	$3 \times b \times s \times h$	$6bsh$
Attention scores	$b \times a \times s \times s$	$2bas^2$
Attention output	$b \times s \times h$	$2bsh$
FFN intermediate	$b \times s \times 4h$ (or $b \times s \times d_{ff}$ )	$8bsh$
LayerNorm inputs (×2)	$2 \times b \times s \times h$	$4bsh$
Dropout masks	(various)	~ $2bsh$
Per-layer total		$\approx 22bsh + 2bas^2$

For $L$ layers:

M_{\text{activations}} = L \times (22bsh + 2bas^2)

Worked Example: 70B Model

Architecture: $h=8192$ , $L=80$ , $a=64$ (GQA with 8 KV heads), $s=4096$ , $b=1$ (micro-batch):

\begin{aligned} \text{Per-layer (linear terms)} &= (14 \times 8192 + 6 \times 28672) \times 4096 \times 1 = 1.2 \text{ GB} \\ \text{Per-layer (attention)} &= 3 \times 64 \times 4096^2 = 3.2 \text{ GB} \\ \text{Total (80 layers)} &= 80 \times 4.5 = 357 \text{ GB (!!)} \end{aligned}

357 GB of activations for a single micro-batch! This is clearly impossible without mitigation.

Activation Checkpointing

Instead of saving all activations, recompute them during backward:

Strategy	Memory	Compute overhead
No checkpointing	$L \times (22bsh + 2bas^2)$	0%
Selective checkpointing	Save only cheap tensors, recompute attention	~30% memory reduction, ~5% compute
Full checkpointing	Save only layer inputs, recompute everything	~ $L \times 2bsh$ (just inputs)

With full activation checkpointing for the 70B example:

M_{\text{activations}} \approx 80 \times 67\text{ MB} = 5.4 \text{ GB}

That’s 357 GB → 5.4 GB — a 66× reduction at the cost of 33% extra forward compute.

How Parallelism Reduces Activation Memory

Strategy	Effect on activations
TP=T	Reduces some tensors by T (attention, FFN), but not all (LayerNorm is replicated)
PP=P	Only stores activations for assigned layers: $L/P$ layers
SP (Sequence Parallel)	Reduces $s$ by the SP degree for LayerNorm/Dropout regions
CP (Context Parallel)	Reduces $s$ by CP degree for attention computation
Gradient accumulation	Doesn’t reduce peak activation memory (still processes full micro-batch)

3. The Complete Memory Budget Formula

M_{\text{per-GPU}} = \underbrace{\frac{16\Psi}{\text{TP} \times \text{PP} \times \text{ZeRO-stage}}}_{\text{static}} + \underbrace{\frac{L}{\text{PP}} \times f(b_\mu, s, h, a, \text{ckpt})}_{\text{activations}} + \underbrace{M_{\text{buffers}}}_{\text{misc}}

Where:

$b_\mu$ = micro-batch size (not global batch)
ZeRO-stage: 1 (optimizer only) shards by D, 2 adds gradients, 3 adds params
$M_{\text{buffers}}$ ≈ 1-4 GB for communication buffers, CUDA context, fragmentation

4. Memory Verification Checklist

For a given configuration, fill in this table:

Component	Formula	Value	Notes
Parameters	$2\Psi / (\text{TP} \times \text{PP})$	__ GB	BF16 working copy
Gradients	$2\Psi / (\text{TP} \times \text{PP})$	__ GB	Sharded further with FSDP
Optimizer	$12\Psi / (\text{TP} \times \text{PP} \times D)$	__ GB	Always sharded across DP
Activations	see formula	__ GB	With/without checkpointing
Comm buffers	~2 GB	2 GB	NCCL buffers
CUDA context	~1 GB	1 GB	Per-GPU fixed overhead
Total	sum	__ GB	Must be < 80 GB

Worked Example: 70B, TP=8, PP=1, DP=32

Component	Calculation	Value
Parameters	140 GB / 8	17.5 GB
Gradients	140 GB / 8	17.5 GB
Optimizer	840 GB / (8 × 32)	3.28 GB
Activations (full ckpt, b_μ=1)	5.4 GB / 8 (rough TP reduction)	~0.7 GB
Buffers + CUDA		3 GB
Total		42.6 GB ✅

Plenty of headroom on 80 GB. Can increase micro-batch size to 2-4.

Worked Example: 405B, TP=8, PP=4, DP=8

Component	Calculation	Value
Parameters	810 GB / (8 × 4)	25.3 GB
Gradients	810 GB / (8 × 4)	25.3 GB
Optimizer	4,860 GB / (8 × 4 × 8)	18.98 GB
Activations (full ckpt, b_μ=1)	~2 GB (126/4 layers, TP=8)	~2 GB
Buffers + CUDA		3 GB
Total		74.6 GB ⚠️ tight

Tight fit — micro-batch size must stay at 1, or use gradient accumulation.

5. Common Pitfalls

Pitfall	What happens	Fix
Forgetting optimizer states	“70B = 140 GB, TP=8 = 17.5 GB, plenty of room!”	Include 12Ψ optimizer overhead
Ignoring activations	OOM at forward pass, not at parameter loading	Add activation checkpointing
Wrong micro-batch size	Calculated with $b_\mu=1$ but training uses $b_\mu=4$	Verify with actual micro-batch
Fragmentation	Theoretical fits, practical OOM	Leave 10-15% headroom
PP stage imbalance	First/last stages have embedding layers	Account for embedding table size

Self-Check

Calculate the 16Ψ memory breakdown for a 13B model
Explain why optimizer states are 12 out of 16 bytes per parameter
Calculate activation memory for one Transformer layer with and without checkpointing
Fill in the memory verification table for 70B on TP=8, DP=8
Identify when a configuration won’t fit and what to adjust