Memory Management › OOM Debugging
Language AI HandbookAll chapters Memory Management › OOM DebuggingOne section · Sections in this chapter5 / 10 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Out-of-memory errors are among the most frustrating failures in deep learning, because they often manifest far into a training run, sometimes after hours of successful execution. A systematic debugging approach dramatically reduces the time to resolution.
Anatomy of an OOM Error
A typical CUDA OOM error looks like:
RuntimeError: CUDA out of memory. Tried to allocate 4.50 GiB
(GPU 0; 79.20 GiB total capacity; 71.43 GiB already allocated;
3.32 GiB free; 73.12 GiB reserved by PyTorch memory allocator)
This message tells you several things. The model tried to allocate 4.5 GB in a single operation. The GPU has 79.2 GB total capacity. PyTorch has 71.4 GB actively in use and 73.1 GB reserved (the difference between reserved and allocated is fragmentation). Only 3.3 GB of contiguous usable memory remains, which is not enough for the 4.5 GB allocation.
The allocation that triggered the error is often not the root cause. Memory may have been building up gradually across many operations, and this particular allocation is simply the one that crossed the threshold. Debugging therefore requires the full memory timeline rather than an inspection of only the final allocation.
The fragmentation number here is particularly telling: 73.1 GB reserved but only 71.4 GB allocated means 1.7 GB of fragmented memory that PyTorch's allocator is holding but cannot currently provide as a contiguous block. When you see a large gap between "reserved" and "allocated," fragmentation is likely contributing to your OOM. The solution is often to restart the process and reduce peak allocation spikes rather than just reducing total memory usage.
Systematic Debugging Steps
Step 1: Confirm available GPU memory. Before anything else, verify you know how much GPU memory your target hardware provides, and that no other processes are consuming it:
import subprocess
try:
result = subprocess.run(["nvidia-smi"], capture_output=True, text=True)
if result.returncode == 0:
print(result.stdout)
else:
print("nvidia-smi not available (CPU-only environment)")
except FileNotFoundError:
print("nvidia-smi not found. Run this on a GPU machine to see GPU status.")Other processes may hold GPU memory even when they appear idle. Jupyter kernels are a common culprit: if you have previously run a training experiment in the same kernel, tensors from that run may still occupy memory. Always check nvidia-smi before attributing an OOM to your model's memory requirements.
Step 2: Add memory profiling checkpoints. Insert torch.cuda.memory_summary() calls at key points to see where memory is consumed. On a CUDA device:
torch.cuda.reset_peak_memory_stats()
## ... forward pass ...
print(f"After forward: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
## ... backward pass ...
print(f"After backward: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
print(torch.cuda.memory_summary())
Step 3: Test with minimal batch size. If the model OOMs even at batch size 1, the model itself does not fit in memory. This indicates you need either mixed precision, offloading, or model parallelism. If it runs at batch size 1, you have a scaling issue and the problem is the activation memory.
Step 4: Enable gradient checkpointing. This is typically the first optimization to reach for when activations are the bottleneck:
from torch.utils.checkpoint import checkpoint_sequential
## Or, for HuggingFace models:
model.gradient_checkpointing_enable()
Step 5: Reduce batch size or sequence length. Both linearly reduce activation memory. Gradient accumulation can maintain effective batch size while using smaller micro-batches.
Step 6: Use mixed precision. If training in FP32, switching to FP16 or BF16 halves parameter and gradient memory. Most modern training code defaults to BF16 but it is worth verifying.
Step 7: Check for memory leaks. Memory that is not freed properly can accumulate across iterations. Common culprits include:
- Storing loss values without detaching (
losses.append(loss)instead oflosses.append(loss.item())): keeps the entire computation graph alive - Keeping references to intermediate tensors outside of training loops
- Running validation inside
torch.enable_grad()context instead oftorch.no_grad() - Logging or visualization code that retains tensors
Step 8: Try CPU offloading. If your model does not fit on GPU at all, CPU offloading moves optimizer states or model weights to CPU RAM between uses. This trades compute throughput for memory capacity. Libraries like DeepSpeed ZeRO-3 implement this automatically. The penalty is communication bandwidth between CPU and GPU, which can reduce throughput by 2-5x but makes otherwise impossible configurations feasible.
Memory Profiling Code
Here is a reusable profiling utility that tracks memory at each step of the training loop:
import gc
import time
class MemoryTracker:
"""
Lightweight memory tracker for training loop debugging.
Works on both CPU (RAM) and GPU (VRAM).
"""
def __init__(self, device="cpu"):
self.device = device
self.checkpoints = []
self.use_cuda = device != "cpu" and torch.cuda.is_available()
def checkpoint(self, label):
"""Record current memory usage at a named checkpoint."""
gc.collect()
if self.use_cuda:
torch.cuda.synchronize()
allocated = torch.cuda.memory_allocated() / 1024**2 # MB
reserved = torch.cuda.memory_reserved() / 1024**2
else:
import tracemalloc
current, peak = tracemalloc.get_traced_memory()
allocated = current / 1024**2
reserved = peak / 1024**2
self.checkpoints.append(
{
"label": label,
"allocated_mb": allocated,
"reserved_mb": reserved,
"timestamp": time.time(),
}
)
def report(self):
"""Print a formatted memory report."""
if not self.checkpoints:
print("No checkpoints recorded.")
return
baseline = self.checkpoints[0]["allocated_mb"]
print(
f"\n{'Checkpoint':<30} {'Memory (MB)':>12} {'Delta (MB)':>11} {'Peak (MB)':>10}"
)
print("-" * 65)
prev = baseline
for cp in self.checkpoints:
delta = cp["allocated_mb"] - prev
peak = cp["reserved_mb"]
delta_str = f"+{delta:.1f}" if delta >= 0 else f"{delta:.1f}"
print(
f"{cp['label']:<30} {cp['allocated_mb']:>11.1f} {delta_str:>11} {peak:>10.1f}"
)
prev = cp["allocated_mb"]
# Demonstrate with a simulated training step
import tracemalloc
tracker = MemoryTracker(device="cpu")
tracemalloc.start()
tracker.checkpoint("Initial state")
# Simulate model creation
model_demo = SimpleTransformer(config)
tracker.checkpoint("After model creation")
# Simulate forward pass with dummy data
batch_size, seq_len = 2, 64
dummy_input = torch.randint(0, config["vocab_size"], (batch_size, seq_len))
tracker.checkpoint("After creating input")
logits = model_demo(dummy_input)
tracker.checkpoint("After forward pass")
# Simulate loss and backward
labels = torch.randint(0, config["vocab_size"], (batch_size, seq_len))
loss = nn.CrossEntropyLoss()(
logits.reshape(-1, config["vocab_size"]), labels.reshape(-1)
)
tracker.checkpoint("After loss computation")
loss.backward()
tracker.checkpoint("After backward pass")
# Cleanup
del logits, loss, labels, dummy_input
gc.collect()
tracker.checkpoint("After cleanup")
tracemalloc.stop()Checkpoint Memory (MB) Delta (MB) Peak (MB) ----------------------------------------------------------------- Initial state 0.0 +0.0 0.0 After model creation 0.1 +0.1 0.2 After creating input 0.1 +0.0 0.2 After forward pass 0.1 +0.0 0.2 After loss computation 0.1 +0.0 0.2 After backward pass 0.1 -0.0 0.2 After cleanup 0.1 -0.0 0.2
The delta column shows exactly when memory spikes occur. In a real training run on GPU, you would instrument these checkpoints around the forward pass, loss calculation, and backward pass to pinpoint where unexpected memory consumption appears. The large delta after the backward pass confirms that gradient tensors are being populated, and the cleanup step shows that proper deletion and garbage collection frees this memory.
Common OOM Patterns and Fixes
The following table summarizes the most frequent OOM patterns encountered in practice:
| OOM Pattern | Likely Cause | Fix |
|---|---|---|
| OOM on first batch | Model too large for GPU | Mixed precision, model parallelism, or smaller model |
| OOM after N batches | Memory leak (graphs not freed) | .detach() losses, zero_grad() every step |
| OOM at longer sequences | Quadratic attention memory | Gradient checkpointing, flash attention, shorter sequences |
| OOM during validation | Validation running with gradients | Wrap in torch.no_grad() |
| OOM after increasing batch | Linear activation growth | Reduce batch size, use gradient accumulation |
| OOM after loading optimizer | Optimizer states exceed budget | 8-bit Adam, AdaFactor, or ZeRO sharding |
A surprisingly common source of OOM crashes is running validation without torch.no_grad(). During inference, you do not need to store activations for backpropagation, so wrapping evaluation code in with torch.no_grad(): eliminates all activation memory for that pass. For a typical validation loop, this can reduce memory usage by 30-50%. The analogous pattern in inference pipelines is to use model.eval() alongside torch.no_grad(), which both disables gradient tracking and switches batch normalization and dropout layers to inference mode.