Fix torch.OutOfMemoryError: PyTorch CUDA OOM Triage Guide

1. Reserved versus Allocated VRAM

The PyTorch caching allocator reserves large blocks of memory from the CUDA driver to avoid expensive kernel synchronization during frequent tensor creations. Total reserved VRAM often far exceeds actively allocated tensor memory.

import torch

# Inspect memory statistics
allocated = torch.cuda.memory_allocated(0) / (1024 ** 3)
reserved = torch.cuda.memory_reserved(0) / (1024 ** 3)
print(f"Allocated: {allocated:.2f} GB | Reserved: {reserved:.2f} GB")

# Print complete breakdown if near capacity
print(torch.cuda.memory_summary(device=0, abbreviated=False))

Note that calling torch.cuda.empty_cache() does not solve training OOMs caused by large forward activations. It only returns unused cached blocks to the driver while introducing severe synchronization stalls.

2. Memory Fragmentation and expandable_segments

Standard caching allocators require contiguous physical memory blocks. Over time, variable tensor lifetimes create fragmentation where total free memory is high, but no single contiguous block satisfies the next allocation.

# Set environment variable before starting Python script
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# Or configure programmatically at entry point
import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
import torch

Enabling expandable segments allows virtual memory segments to grow and merge without requiring contiguous physical underlying pages, virtually eliminating fragmentation crashes.

3. Activation Checkpointing

During the forward pass, intermediate activations are stored for gradient calculation during backpropagation. Activation checkpointing drops these activations and recomputes them on the fly during the backward pass.

import torch
from torch.utils.checkpoint import checkpoint

class CustomModel(torch.nn.Module):
    def __init__(self, layer):
        super().__init__()
        self.layer = layer

    def forward(self, x):
        return checkpoint(self.layer, x, use_reentrant=False)

# For Hugging Face Transformers models
model.gradient_checkpointing_enable(
    gradient_checkpointing_kwargs={"use_reentrant": False}
)

4. Micro-batch Sizing and Gradient Accumulation

When hardware constraints prevent large batch sizes, reduce the per-device micro-batch size and accumulate gradients across multiple steps to preserve effective batch size.

accumulation_steps = 4
optimizer.zero_grad()

for i, batch in enumerate(dataloader):
    outputs = model(batch["input_ids"])
    loss = outputs.loss / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

5. Common Memory Leak Traps

Subtle Python patterns can inadvertently retain computation graphs across iterations, leaking gigabytes of VRAM over time.

Frequently Asked Questions

Related Diagnostic Tools