GPU memory is a hard, often limiting constraint in both training and serving — this note covers practical memory optimization techniques for fitting larger models and batch sizes within available memory.
What Consumes GPU Memory During Training
| Component | Notes |
|---|---|
| Model parameters | Scales directly with model size |
| Gradients | One gradient value per parameter — roughly doubles the parameter memory footprint |
| Optimizer state | Optimizers like Adam store additional per-parameter state (momentum, variance estimates) — can be 2x the parameter memory on top of gradients |
| Activations | Intermediate layer outputs saved for the backward pass — scales with batch size and network depth, often the largest and most variable consumer |
Code — Gradient Checkpointing (Trading Compute for Memory)
import torch
from torch.utils.checkpoint import checkpoint
class MemoryEfficientBlock(torch.nn.Module):
def forward(self, x):
# Instead of storing this block's activations for the backward pass,
# checkpoint recomputes them during backprop -- saving memory at the
# cost of extra computation
return checkpoint(self.block, x, use_reentrant=False)Gradient checkpointing deliberately avoids storing every intermediate activation, instead recomputing them during the backward pass when needed — a direct, deliberate tradeoff of extra compute time for significantly reduced memory usage, valuable when memory (not compute) is the binding constraint.
Code — Mixed Precision Training for Memory Savings
import torch
scaler = torch.cuda.amp.GradScaler()
for x_batch, y_batch in train_loader:
optimizer.zero_grad()
with torch.cuda.amp.autocast(): # uses FP16 for most operations, reducing memory
output = model(x_batch)
loss = loss_fn(output, y_batch)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()Using lower-precision (FP16/BF16) representations for most operations roughly halves the memory footprint of activations and gradients compared to full FP32 precision, with GradScaler managing numerical stability — one of the most broadly effective and commonly used memory optimization techniques.
Other Practical Memory-Saving Techniques
- Gradient accumulation — simulating a larger effective batch size by accumulating gradients over several smaller batches before an optimizer step, avoiding the memory cost of one large batch.
- Reducing batch size — the simplest lever, though it directly affects training dynamics and GPU utilization, so it's often a last resort after other techniques are exhausted.
- Model parallelism — splitting a model across multiple GPUs when it's too large to fit on a single device at all (see Distributed Training).
Common Mistakes
- Immediately reducing batch size as the first response to an out-of-memory error, without first trying mixed precision or gradient checkpointing — these often recover substantial memory with less impact on training dynamics or GPU utilization.
- Using gradient checkpointing indiscriminately across an entire model when memory isn't actually the binding constraint — the added recomputation cost is a real tradeoff, not a free optimization, and should be applied where memory savings are genuinely needed.
Interview Relevance
Q: "You're training a large model and hit an out-of-memory error. What would you try before simply reducing the batch size?" Mixed precision training (FP16/BF16) typically cuts activation and gradient memory roughly in half with minimal accuracy impact, and gradient checkpointing trades extra recomputation for significantly reduced activation memory — both often recover enough memory to keep the original batch size, preserving training dynamics and GPU utilization better than simply shrinking the batch. Reducing batch size (or using gradient accumulation to compensate) remains a valid fallback if these techniques aren't sufficient on their own.
Practice Question
Why does gradient checkpointing reduce memory usage at the cost of increased training time, rather than being a purely "free" optimization?
