TopolearnOut-of-memory errors are the most common blocker when scaling models; fixing them needs a model of where memory goes.
Training memory has four parts: parameters, gradients, optimizer state and activations saved for backward. Activations scale with batch size and sequence length and usually dominate.
PyTorch uses a caching allocator: freed blocks stay reserved for reuse instead of returning to the driver, because cudaMalloc is slow. So torch.cuda.memory_allocated() (live tensors) is lower than memory_reserved() (what the process holds and what nvidia-smi shows). torch.cuda.empty_cache() releases unused cached blocks to other processes; it does not free live tensors and rarely fixes an OOM.
Fragmentation happens when free memory is split into blocks too small for a request. Setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True lets segments grow and often reduces it.
To see what holds memory, record a memory snapshot with torch.cuda.memory._record_memory_history() and _dump_snapshot(), then open it in the visualiser at pytorch.org/memory_viz, which shows every allocation with its stack trace.
Fixes, roughly in order: smaller batch with gradient accumulation, mixed precision, activation checkpointing (torch.utils.checkpoint.checkpoint(block, x, use_reentrant=False), which drops a block's activations and recomputes them during backward, about 30% more compute), fused optimizers, then sharding with FSDP.
Example
import torch
from torch import nn
from torch.utils.checkpoint import checkpoint
assert torch.cuda.is_available(), "needs a CUDA GPU"
torch.cuda.memory._record_memory_history(max_entries=100_000)
block = nn.Sequential(nn.Linear(4096, 4096), nn.GELU(), nn.Linear(4096, 4096)).cuda()
x = torch.randn(64, 4096, device="cuda", requires_grad=True)
torch.cuda.reset_peak_memory_stats()
y = checkpoint(block, x, use_reentrant=False) # keep only the block's input
y.sum().backward() # recompute activations here
mib = 2 ** 20
print(f"peak {torch.cuda.max_memory_allocated() / mib:.0f} MiB")
print(f"allocated {torch.cuda.memory_allocated() / mib:.0f} MiB")
print(f"reserved {torch.cuda.memory_reserved() / mib:.0f} MiB")
torch.cuda.memory._dump_snapshot("mem.pickle") # drag into pytorch.org/memory_viz
torch.cuda.memory._record_memory_history(enabled=None)Remember
- Allocated = live tensors; reserved = cached by the allocator
- Snapshots show exactly which allocations fill the GPU
- Activation checkpointing trades ~30% compute for activation memory
Watch out for
- Calling
torch.cuda.empty_cache()in the training loop, which slows training and does not free live tensors.