GPU Memory Anatomy

Understanding where GPU memory goes is the foundation of all training infrastructure decisions. This chapter gives you the complete breakdown — every byte accounted for — so you can predict memory usage before launching a single training run.

---

1. The Four Categories of Training Memory

Training memory splits into static (constant regardless of input) and dynamic (varies with batch size and sequence length):

Category	Static/Dynamic	Typical share	Depends on
Model parameters	Static	12-15%	Parameter count Ψ
Gradients	Static	12-15%	Parameter count Ψ
Optimizer states	Static	55-75%	Parameter count Ψ, optimizer type
Activations	Dynamic	10-30%	Batch size, sequence length, model width
Framework overhead	Static	2-5%	CUDA context, NCCL buffers, fragmentation

Key insight: Optimizer states dominate training memory. For Adam/AdamW, they consume 3× the memory of the parameters themselves.

---

2. Model Parameters

The working copy of the model used in forward and backward passes.

$$M_{\text{params}} = \Psi \times \text{bytes\_per\_param}$$

Precision	Bytes	7B	13B	70B	405B
FP32	4	28 GB	52 GB	280 GB	1,620 GB
BF16	2	14 GB	26 GB	140 GB	810 GB
FP8	1	7 GB	13 GB	70 GB	405 GB

Modern training uses BF16 (or FP8 on Hopper+) for the working copy, with FP32 reserved for the optimizer’s master weights.

Counting Parameters Accurately

For a Transformer with $L$ layers, hidden dim $h$, vocab $V$, FFN dim $d_{ff}$:

$$\Psi \approx V \times h + L \times (4h^2 + 2h \times d_{ff} + \text{norms})$$

Component	Parameters per layer	70B (h=8192)
QKV projection	$3h^2$ (or $h^2 + 2h_{kv} \times h$ for GQA)	~201M
Output projection	$h^2$	67M
FFN up + gate + down	$3h \times d_{ff}$ (SwiGLU)	707M
LayerNorm (×2)	$4h$	65K
Per-layer total		~975M
80 layers		~78B
Embedding + LM head	$2 \times V \times h$	~2.6B (V=128K)

Counting Parameters: Worked Example (Llama-3.1 70B)

Architecture: $h=8192$, $L=80$, $V=128{,}256$, GQA with $a=64$ query heads, $a_{kv}=8$ KV heads, $d_{ff}=28{,}672$ (SwiGLU), no bias terms.

Component	Formula	Count
Embedding	$V \times h$	128,256 × 8,192 = 1.050B
Per-layer QKV projection	$h \times (a \times d_h + 2 \times a_{kv} \times d_h)$	8,192 × (64×128 + 2×8×128) = 75.5M
Per-layer output projection	$a \times d_h \times h$	64 × 128 × 8,192 = 67.1M
Per-layer FFN (SwiGLU: gate + up + down)	$3 \times h \times d_{ff}$	3 × 8,192 × 28,672 = 704.6M
Per-layer RMSNorm (×2)	$2 \times h$	2 × 8,192 = 16.4K
Per-layer total		847.2M
80 layers		67.78B
Final RMSNorm	$h$	8.2K
LM head (tied with embedding)	0 (shared)	0
Grand total		~68.8B

The “70B” label is approximate — actual parameter counts vary by implementation details (bias terms, head dim rounding, tied/untied LM head).

---

3. Gradients

Same shape as parameters — one gradient value per parameter:

$$M_{\text{gradients}} = \Psi \times \text{bytes\_per\_grad}$$

With BF16 training: 2 bytes per gradient. With FP32 gradient accumulation: 4 bytes.

Gradient accumulation doesn’t change peak gradient memory — you still hold one full set of gradients, just accumulated over multiple micro-batches before the optimizer step.

---

4. Optimizer States (The Dominant Cost)

Different optimizers require different additional state per parameter:

Optimizer	States per param	Bytes/param (FP32)	7B overhead	Notes
SGD	None	0	0 GB	Rarely used for LLMs
SGD + Momentum	velocity $v$	4	28 GB	One FP32 buffer
Adam/AdamW	$m$ (1st moment) + $v$ (2nd moment)	8	56 GB	Two FP32 buffers
Adam + FP32 master	master weights + $m$ + $v$	12	84 GB	Standard mixed-precision

The Adam 12Ψ Rule

In standard BF16 mixed-precision training with AdamW, the optimizer maintains:

1. FP32 master weights (4Ψ bytes) — high-precision copy for accurate updates
2. FP32 first moment $m$ (4Ψ bytes) — exponential moving average of gradients
3. FP32 second moment $v$ (4Ψ bytes) — EMA of squared gradients

Total optimizer overhead: 12Ψ bytes.

Why FP32 for optimizer states? Small weight updates ($\text{lr} \times \text{gradient} \approx 10^{-4} \times 10^{-3} = 10^{-7}$) fall below BF16’s precision floor ($\approx 10^{-3}$ relative). FP32 is necessary to accumulate these small changes accurately. Without it, weights simply don’t update — training stalls.

The Full Static Ledger

Component	Bytes/param	Formula	7B	70B
BF16 parameters	2	$2\Psi$	14 GB	140 GB
BF16 gradients	2	$2\Psi$	14 GB	140 GB
FP32 master weights	4	$4\Psi$	28 GB	280 GB
FP32 Adam $m$	4	$4\Psi$	28 GB	280 GB
FP32 Adam $v$	4	$4\Psi$	28 GB	280 GB
Total	16	16Ψ	112 GB	1,120 GB

For a 7B model: 112 GB of static memory — already exceeds a single 80 GB GPU. This is why even “small” models need distributed training.

---

5. Activations (The Dynamic Cost)

Activations are intermediate tensors saved during forward pass for the backward pass. They scale with batch size and sequence length.

For a Transformer layer:

Saved tensor	Shape	Purpose
Input to attention	$b \times s \times h$	Backward through QKV projection
Q, K, V tensors	$3 \times b \times s \times h$	Backward through attention scores
Attention weights	$b \times a \times s \times s$	Backward through softmax
Attention output	$b \times s \times h$	Backward through output projection
FFN input	$b \times s \times h$	Backward through FFN
FFN intermediate	$b \times s \times d_{ff}$	Backward through gate/activation
LayerNorm inputs	$2 \times b \times s \times h$	Backward through normalization

Approximate total per layer (BF16):

$$M_{\text{act/layer}} \approx 2 \times (22bsh + 2bas^2) \text{ bytes}$$

For 70B ($h=8192$, $a=64$, $s=4096$, $b_\mu=1$):

• Linear terms: $14 \times 4096 \times 8192 + 6 \times 4096 \times 28672 \approx 1.2$ GB
• Attention terms: $3 \times 64 \times 4096^2 \approx 3.2$ GB
• Per layer: ~4.5 GB
• 80 layers: ~357 GB without activation checkpointing

Activations scale as $O(s^2)$ for standard attention. This is one reason long-context training is so memory-intensive. FlashAttention removes this quadratic term. See Activation Memory Analysis for detailed formulas and mitigation strategies.

---

6. Framework Overhead

Source	Typical size	Notes
CUDA context	0.5-1.5 GB	Per GPU, allocated on first CUDA call
NCCL buffers	0.5-2 GB	Communication scratch space
PyTorch allocator fragmentation	1-5 GB	Wasted space in memory pools
Temporary buffers	0.5-2 GB	Intermediate computation tensors

Total overhead: typically 3-8 GB per GPU — significant when memory is tight.

# Check actual memory usage vs theoretical
import torch
print(f"Allocated: {torch.cuda.memory_allocated()/1e9:.1f} GB")
print(f"Reserved:  {torch.cuda.memory_reserved()/1e9:.1f} GB")
print(f"Max allocated: {torch.cuda.max_memory_allocated()/1e9:.1f} GB")
# Reserved - Allocated = fragmentation + cached blocks

---

7. Complete Memory Picture

For a 7B model on a single 80 GB H100 (BF16 + AdamW, $b_\mu=4$, $s=2048$):

Component	Memory
Parameters (BF16)	14 GB
Gradients (BF16)	14 GB
Optimizer (FP32 master + m + v)	84 GB
Static subtotal	112 GB — already OOM!

This is why distributed training exists. Even a 7B model can’t train on a single GPU with Adam. Solutions:

1. FSDP/ZeRO: Shard static memory across GPUs → 112/8 = 14 GB per GPU
2. CPU offloading: Move optimizer to system RAM → only 28 GB on GPU
3. 8-bit Adam: Use INT8 optimizer states → 56 GB instead of 84 GB
4. SGD instead of Adam: 0 optimizer overhead → 28 GB total (but worse convergence)

For sharding strategies, see ZeRO Stage 1/2/3.

---

Self-Check

• Name the four categories of training memory and their approximate shares
• Calculate the 16Ψ static memory for a 13B model with Adam
• Explain why optimizer states consume 75% of static training memory
• Explain why Adam needs FP32 (precision floor argument)
• Calculate per-layer activation memory for a given architecture
• Use torch.cuda.memory_allocated() to verify theoretical predictions

---

References

• ZeRO: Memory Optimizations (Rajbhandari et al., 2019)
• Reducing Activation Recomputation in Large Transformer Models (Korthikanti et al., 2022)
• Adam: A Method for Stochastic Optimization (Kingma & Ba, 2014)