Diagnosing PyTorch CUDA Out-of-Memory Errors on a Single GPU

When a Training Batch Triggers `RuntimeError: CUDA out of memory`

The problem is a PyTorch training script halting because the CUDA allocator cannot satisfy a memory request on a single GPU. The useful takeaway is that the error often reflects how memory is reserved and fragmented, not just how much is actively used. Distinguishing allocated memory from reserved memory determines whether you should shrink the batch, change the allocator, or restructure the model.

Condition, Cause, and First Check

What you see	Likely cause	First check
Allocation fails during forward or backward pass	Active tensors and activations exceed free device memory	Log `torch.cuda.memory_allocated()` against total memory
Reserved memory is high while allocated memory is low	Allocator holds free blocks that cannot be reused for the requested shape	Run `torch.cuda.memory_summary()`
Error persists at `batch_size=1`	Model parameters and optimizer states alone occupy most memory	Inspect parameter memory in `memory_summary()`

Allocated memory is memory currently holding tensors. Reserved memory is memory the PyTorch CUDA allocator has obtained from the driver but has not necessarily handed to a tensor. Fragmentation happens when reserved blocks are free but too small or shaped incorrectly for a new allocation.

Ordered Diagnostic Checks

1. Confirm the environment. Run in the shell where you launch training:

python -c 'import torch; print(torch.__version__, torch.version.cuda, torch.cuda.is_available())'

Required permissions: none beyond running Python as your user. Expected check: CUDA is available and the printed CUDA version matches your driver. Risk: a mismatch between PyTorch build and driver can produce allocator errors unrelated to memory size.

2. Log allocation before and after the failing batch. Insert this near the training step:

import torch

device = torch.device('cuda:0')
print(f'Allocated before: {torch.cuda.memory_allocated(device) / 1024**2:.2f} MiB')
print(f'Max allocated: {torch.cuda.max_memory_allocated(device) / 1024**2:.2f} MiB')
# run one batch here
print(f'Allocated after: {torch.cuda.memory_allocated(device) / 1024**2:.2f} MiB')

Expected check: if allocated after approaches total GPU memory, the model and batch are too large for the device.

3. Inspect reserved blocks. Run:

print(torch.cuda.memory_summary(device))

Look for a large gap between "Allocated" and "Reserved" in the summary. A wide gap suggests fragmentation or tensors that were freed but not returned to the operating system.

Fixes Tied to Findings

Allocation exhaustion

• Reduce `batch_size` or `sequence_length` if the failure occurs during allocation of activations.
• Enable gradient checkpointing. This trades compute for memory by discarding intermediate activations during the forward pass and recomputing them during backward. In PyTorch, wrap the model or use `torch.utils.checkpoint` on selected layers.
• Use mixed precision with `torch.autocast`. This stores activations in lower precision such as float16 or bfloat16. Verify loss convergence after enabling it, because numerics can change.

Fragmentation

• Set the allocator configuration before launching Python:

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:64 python train.py

Where to run: the shell command that starts your script. Permissions: user running the script. Placeholder: `train.py` is your entry point. Expected check: rerun `memory_summary()` and confirm reserved blocks better match allocation requests. Risk: this does not free memory held by active tensors, and allocator behavior varies by PyTorch and CUDA version.

• Call `torch.cuda.empty_cache()` only at phase boundaries, such as between epochs. Do not call it inside the inner batch loop; it releases unused reserved blocks but cannot free tensors still referenced by the computation graph, and repeated calls slow training.

Unreleased tensors

• Log scalar values with `loss.item()` instead of appending the full `loss` tensor to a list, which keeps the graph alive.
• Wrap inference or validation code in `torch.no_grad()` so PyTorch does not store activations for backward.

Concrete Configuration Example

If `memory_summary()` shows high reserved memory and the failure is a large intermediate allocation, launch with:

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:128 python -m torch.distributed.launch --nproc_per_node=1 train.py --batch-size 4

This keeps the allocator from splitting blocks too aggressively. Monitor `max_memory_allocated()` across several batches to confirm the peak stays below total memory.

Escalation Criteria

Escalate when `batch_size=1` still produces OOM and `memory_summary()` shows model parameters and optimizer states consuming most reserved memory. At that point, code-level fixes are insufficient. Move to DistributedDataParallel across multiple GPUs, or offload optimizer states and parameters to CPU. File a bug report only with a minimal, self-contained reproduction script and the output of the environment check above.

Verification and Limitations

Verification: After applying a fix, run the same batch size that previously failed. Confirm success when training completes and `torch.cuda.max_memory_allocated()` remains below the GPU's total memory for the full run. Re-run the environment check on the target machine to record versions.

Limitation: `empty_cache()` does not free active tensors. Mixed precision can alter convergence. Allocator environment variables behave differently across PyTorch releases; test on the exact stack you plan to deploy.

Rollback

Remove `PYTORCH_CUDA_ALLOC_CONF` from the launch command or unset it in the shell to return to the default allocator. Revert gradient checkpointing or mixed-precision code changes to restore previous memory and compute behavior.