LEARN · DEBUGGING GUIDE

PyTorch CUDA Out of Memory: Diagnosis and Recovery

CUDA OOM errors in PyTorch are often not about total memory limits but fragmentation or hidden allocations. Stop guessing — use the exact diagnostics and fixes here.

AdvancedML / Data7 min read

What this usually means

PyTorch's CUDA allocator caches memory blocks for reuse. An OOM error often means the allocator cannot find a contiguous block of the requested size, even if total free memory is sufficient. This happens due to fragmentation, unreferenced tensors still occupying space, or gradients accumulating across iterations. Non-obvious causes include PyTorch's caching allocator not releasing memory to the system, unused variables in the graph, or forgetting to call .detach() on loss values. The error message's 'Tried to allocate X GiB' is the requested allocation, not necessarily the total memory in use.

( 01 )Fast diagnosis

The first ten minutes — establish facts before touching code.

  • 1Run `nvidia-smi` to check total/free memory. If free memory is low but PyTorch reports little allocated, it's fragmentation.
  • 2Set environment variable `PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128` to reduce fragmentation and rerun.
  • 3Wrap the training loop with `torch.cuda.empty_cache()` after each epoch to release cached allocator memory.
  • 4Check for tensors retained outside the training loop: print `len(torch.cuda.memory_summary())` or use `torch.cuda.memory_snapshot()`.
  • 5Reduce batch size by half and see if OOM disappears. If it does, the issue is memory usage per iteration.
  • 6Enable anomaly detection with `torch.autograd.set_detect_anomaly(True)` to find where gradients explode.
( 02 )Where to look

The specific files, logs, configs, and dashboards that usually own this bug.

  • searchTerminal output: exact error line number and tensor size in allocation message.
  • search`nvidia-smi` output: used/free memory, processes consuming GPU.
  • searchPyTorch memory summary: `torch.cuda.memory_summary()` for detailed allocation breakdown.
  • searchTraining loop code: variables that persist across iterations (e.g., lists of outputs).
  • searchGradient accumulation logic: ensure loss is divided by accumulation steps.
  • searchModel definition: layers with large intermediate tensors (e.g., large linear layers, attention heads).
  • searchData loader: batch size, number of workers, pin_memory settings.
( 03 )Common root causes

Practical causes, not theory. These are the things you will actually find.

  • warningMemory fragmentation from variable-sized tensors (e.g., dynamic batch sizes, attention masks).
  • warningGradient accumulation without dividing loss, causing gradients to grow unbounded.
  • warningUnreleased references to tensors in lists or dictionaries across iterations.
  • warningPyTorch caching allocator holding memory that nvidia-smi reports as used but is reclaimable.
  • warningDataLoader pin_memory=True causing pinned CPU memory to exhaust GPU memory.
  • warningUsing too many workers in DataLoader, each holding GPU copies.
  • warningModel too large for GPU but fits in CPU memory — OOM on .to('cuda').
( 04 )Fix patterns

Concrete fix directions. Pick the one that matches your root cause.

  • buildSet `PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:32` or `128` to reduce fragmentation.
  • buildUse `torch.cuda.empty_cache()` at strategic points (e.g., after validation) to free cached blocks.
  • buildConvert model to half precision: `model.half()` or use AMP with `torch.cuda.amp`.
  • buildReduce batch size or use gradient accumulation with proper scaling (loss = loss / accumulation_steps).
  • buildProfile memory with `torch.cuda.memory_summary()` and refactor code to release tensors early with `del`.
  • buildMove evaluation to CPU or use `with torch.no_grad():` to avoid graph retention.
  • buildSwitch to `torch.utils.checkpoint` (gradient checkpointing) to trade compute for memory.
( 05 )How to verify

A fix you cannot prove is a guess. Close the loop.

  • verifiedRun the same training step with `torch.cuda.memory_summary()` before and after to see allocated vs cached.
  • verifiedAfter fix, run `nvidia-smi` during peak iteration; used memory should be stable and below 90%.
  • verifiedTrain for at least 10% of total iterations to ensure no late-stage OOM.
  • verifiedIf using gradient accumulation, verify loss curve matches non-accumulated run.
  • verifiedCheck that `torch.cuda.memory_allocated()` does not increase across epochs without corresponding increases in batch size.
  • verifiedUse `torch.cuda.set_per_process_memory_fraction(0.9)` to fail fast if memory exceeds 90%.
( 06 )Mistakes to avoid

Things that make this bug worse or harder to find.

  • warningCalling `torch.cuda.empty_cache()` after every iteration — it slows training and doesn't fix fragmentation.
  • warningSetting `PYTORCH_CUDA_ALLOC_CONF` incorrectly (e.g., negative values or non-existent keys).
  • warningBlindly reducing batch size without checking if it's actually the problem.
  • warningUsing `torch.no_grad()` incorrectly outside inference — prevents gradient computation.
  • warningIgnoring warnings from PyTorch about memory usage (e.g., 'Memory allocated: X' messages).
  • warningAssuming OOM means you need a bigger GPU — often it's fixable with software changes.
( 07 )War story

The Fragmenting Transformer: A 16GB GPU Dies Mid-Epoch

ML EngineerPyTorch 1.12, NVIDIA A4000 16GB, Transformers 4.21

Timeline

  1. 09:00Start training BERT-large on custom dataset, batch size 32, sequence length 512.
  2. 09:05OOM error: 'Tried to allocate 512.00 MiB. GPU has 16.00 GiB total capacity.'
  3. 09:07Check nvidia-smi: 14.5 GiB used, 1.5 GiB free.
  4. 09:10Reduce batch size to 16 — still OOM after a few iterations.
  5. 09:15Call torch.cuda.memory_summary(); see large number of small allocations (fragmentation).
  6. 09:20Set PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:64, retry.
  7. 09:25Training runs for 2 epochs then OOM again — fragmentation returns.
  8. 09:30Add torch.cuda.empty_cache() after each validation step.
  9. 09:35Training completes 5 epochs without OOM.

I was training a BERT-large model on a 16GB A4000. The dataset had variable-length sequences, so I padded to 512. First attempt with batch size 32 crashed immediately — nvidia-smi showed 14.5GB used but PyTorch said only 10GB allocated. Classic fragmentation.

I dropped batch size to 16, hoping to reduce pressure. It ran for a few hundred iterations then died with the same error. I printed memory_summary() and saw thousands of tiny allocations from attention masks of different sizes. The allocator couldn't find a contiguous 512MB block.

Setting max_split_size_mb:64 helped initially, but after two epochs fragmentation returned. I added empty_cache() after validation, which cleared the cached allocator blocks. That stable train ran for 5 epochs without a single OOM. Lesson: don't just reduce batch size — fix fragmentation and evict cached blocks periodically.

Root cause

Memory fragmentation caused by variable-length attention masks creating many non-contiguous allocations. PyTorch's caching allocator held onto blocks, preventing coalescing.

The fix

Set max_split_size_mb:64 in environment variable and call torch.cuda.empty_cache() after each validation epoch.

The lesson

CUDA OOM is often about fragmentation, not total memory. Use memory_summary() to diagnose, then tune the allocator and strategically empty cache.

( 08 )Understanding PyTorch's CUDA Caching Allocator

PyTorch uses a caching memory allocator to avoid costly cudaMalloc calls. When you free a tensor, the allocator keeps the block in a cache for reuse. Over time, this cache becomes fragmented with blocks of varying sizes. When a new tensor requires a contiguous block larger than any cached block, PyTorch calls cudaMalloc, which may fail if system memory is exhausted — even if total cached memory is sufficient.

The environment variable `PYTORCH_CUDA_ALLOC_CONF` controls allocator behavior. Key settings: `max_split_size_mb` (prevents splitting large blocks beyond this size), `roundup_power2_divisions` (rounds allocation sizes to reduce fragmentation), and `garbage_collection_threshold` (triggers cache cleaning when memory usage exceeds a threshold). Setting `max_split_size_mb:128` is a common first step to reduce fragmentation.

( 09 )Diagnosing with torch.cuda.memory_summary()

Call `torch.cuda.memory_summary()` to get a detailed report of allocated, cached, and active memory. Look for the 'Allocated memory' vs 'Cached memory' gap. A large gap indicates fragmentation. Also check 'Number of allocated blocks' — a high number suggests many small tensors.

For a live view, use `torch.cuda.memory_snapshot()` which returns a dictionary of all allocated blocks. You can visualize this with `torch.cuda.memory._dump_snapshot()` (PyTorch >=1.10) to generate a timeline of allocations. This is invaluable for finding which operation causes a spike.

( 10 )Gradient Accumulation and Memory Bloat

Gradient accumulation can silently increase memory usage if the loss is not properly scaled. Each backward pass accumulates gradients into .grad attributes, which persist until optimizer.step(). If the loss is not divided by the number of accumulation steps, gradients become large, consuming more memory. More subtly, intermediate activations from each forward pass are retained until backward is called — with accumulation, they pile up.

The fix: divide loss by accumulation steps: `loss = loss / accumulation_steps`. Also ensure you call `optimizer.zero_grad()` at the start of each accumulation cycle, not after every backward. Use `with torch.no_grad():` for the forward passes that are only for accumulation? No — you need gradients. Instead, use checkpointing or reduce batch size per accumulation step.

( 11 )DataLoader and Pin Memory Pitfalls

Setting `pin_memory=True` in DataLoader speeds up CPU-to-GPU transfer by pinning host memory. However, it can also cause GPU memory exhaustion if the pinned memory is not released promptly. The DataLoader prefetches batches in pinned memory, and if the GPU is busy, pinned memory accumulates. This is especially problematic with many workers.

Monitor pinned memory with `nvidia-smi` — look for 'Pinned' memory usage. Solutions: reduce `num_workers`, set `pin_memory=False` (slower but safer), or use `prefetch_factor=2` to limit prefetching. Also, call `del batch` or use context managers to release references.

( 12 )Advanced: Gradient Checkpointing and Mixed Precision

When all else fails, trade compute for memory. Gradient checkpointing (`torch.utils.checkpoint`) recomputes intermediate activations during backward, reducing peak memory — works well for large models like transformers. Trade-off: ~20% slower training.

Mixed precision training with `torch.cuda.amp` reduces memory by storing tensors in float16. Combined with gradient scaling, it often halves memory usage. Enable with `with torch.cuda.amp.autocast():` and `scaler.scale(loss).backward()`. This is the most effective single fix for OOMs due to model size.

Frequently asked questions

Why does nvidia-smi show high memory usage but PyTorch reports low allocated memory?

PyTorch's caching allocator holds onto freed memory blocks for reuse. nvidia-smi reports all memory held by the process, including cached but not currently allocated blocks. This is normal — the cached memory can be reclaimed by PyTorch when needed, but if fragmentation prevents reuse, it appears as 'used' in nvidia-smi. Use `torch.cuda.memory_summary()` to see the split between allocated and cached.

What does the 'Tried to allocate X GiB' error actually mean?

It means PyTorch attempted to allocate a contiguous block of X GiB for a new tensor. The total GPU memory may have enough free space, but not in a single contiguous chunk. This often indicates fragmentation. The error does not necessarily mean total memory is exhausted; check `torch.cuda.memory_summary()` for fragmentation metrics.

How do I clear PyTorch's CUDA cache without restarting the kernel?

Call `torch.cuda.empty_cache()`. This releases all cached memory blocks back to the system. However, it does not free memory currently allocated to tensors. Use it sparingly, as it can slow down training by forcing re-allocation. Effective after validation or evaluation when you have a break in gradient computation.

Can I limit PyTorch's GPU memory usage to avoid OOM?

Yes. Use `torch.cuda.set_per_process_memory_fraction(0.9)` to limit PyTorch to 90% of total GPU memory. This causes PyTorch to fail early with an OOM if it exceeds the limit, rather than crashing later. Also, set environment variable `PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.8` to trigger cache cleanup when memory usage hits 80%.

Why does OOM happen after several iterations, not immediately?

Gradual memory increase can be due to unreleased tensors in lists or dictionaries that grow each iteration, gradient accumulation without proper scaling, or caching allocator fragmentation building up. Use `torch.cuda.memory_allocated()` to track allocation over time and identify the growth. Common culprits: storing loss values in a list for logging, or not detaching tensors before appending.