Quick Fix (TL;DR)

The most common cause of torch.cuda.OutOfMemoryError on modern Linux systems is memory fragmentation in the caching allocator rather than physical VRAM exhaustion. Set the expandable segments allocator flag in your shell or Python script before launching your workload:

# Export in your shell or .bashrc
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# Launch your training or inference script
python train.py

In Python code, set this environment variable before importing PyTorch:

import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

import torch
# PyTorch now allocates dynamic virtual memory chunks instead of fixed blocks

---

1. Why PyTorch Throws CUDA OOM

PyTorch does not allocate VRAM directly from the NVIDIA driver for every tensor. It uses a caching allocator. When a tensor is deleted, the memory returns to the PyTorch memory pool, not to the GPU hardware pool.

This caching mechanism leads to three distinct failure modes:

Failure Mode	Error Signature	Root Cause	Immediate Fix
Memory Fragmentation	Reserved > Allocated	Small allocations block contiguous free space	expandable_segments:True
Physical VRAM Exceeded	Reserved == Total VRAM	Batch size or model weights exceed hardware limit	Enable Gradient Checkpointing
Zombie CUDA Process	nvidia-smi shows memory in use	Dead background process holding GPU context	fuser -v /dev/nvidia* and kill PID
Accumulated Graph History	Memory climbs steadily every step	Loss tensors appended to list without .item()	Detach loss tensors

---

2. Diagnosing Active GPU Memory Allocation

Before modifying batch sizes or quantization settings, inspect the exact split between allocated memory, cached memory, and driver reserved memory.

Run this diagnostic snippet inside your training loop:

import torch

def print_cuda_memory_stats(device: int = 0):
    allocated = torch.cuda.memory_allocated(device) / (1024 ** 3)
    reserved = torch.cuda.memory_reserved(device) / (1024 ** 3)
    max_allocated = torch.cuda.max_memory_allocated(device) / (1024 ** 3)
    
    print(f"Device [{device}] Memory Stats:")
    print(f"  Allocated (Active Tensors): {allocated:.2f} GB")
    print(f"  Reserved (PyTorch Cache):   {reserved:.2f} GB")
    print(f"  Peak Allocated:             {max_allocated:.2f} GB")

# Test on GPU 0
print_cuda_memory_stats(0)

If Reserved is close to your GPU limit (e.g., 23.5 GB on an RTX 4090) while Allocated is significantly lower (e.g., 14.0 GB), your script failed due to fragmentation. The allocator could not find a single contiguous block large enough for the incoming tensor.

---

3. Four Methods to Eliminate PyTorch CUDA OOM

Apply these solutions based on your memory diagnostic output:

Method A: Expandable Segments (PyTorch 2.1+)

PyTorch 2.1 introduced virtual memory management via expandable_segments:True. Instead of reserving large fixed segments that cannot be split, PyTorch maps non-contiguous physical pages into contiguous virtual GPU memory spaces.

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:512

This prevents fragmentation crashes during generation tasks with variable input sequence lengths.

Method B: Activation Checkpointing (Gradient Checkpointing)

During model training, storing intermediate forward activations for backward pass backpropagation consumes up to 70% of total VRAM. Gradient checkpointing discards forward activations and recomputes them on the fly during the backward pass.

Enable checkpointing in Hugging Face Transformers:

model.gradient_checkpointing_enable()

# For custom PyTorch modules:
from torch.utils.checkpoint import checkpoint
output = checkpoint(custom_layer, input_tensor, use_reentrant=False)

This trades 20% slower step speed for a 60% reduction in activation memory.

Method C: Finding and Killing Zombie GPU Processes

Often, a previous PyTorch run crashed or was killed with Ctrl+Z rather than Ctrl+C, leaving background CUDA workers holding VRAM that nvidia-smi may report as unknown processes.

Find all processes holding CUDA device file descriptors:

sudo fuser -v /dev/nvidia*

Kill the orphaned processes cleanly:

sudo fuser -vk -9 /dev/nvidia0

Method D: Detaching History in Metrics and Logging

A common subtle bug is calculating running loss averages by accumulating raw tensors:

# BROKEN: Keeps the entire computation graph in memory indefinitely
total_loss += loss

# FIXED: Extracts raw scalar float, freeing the computation graph
total_loss += loss.item()

---

4. Multi-GPU Distributed Memory Balancing: FSDP & DeepSpeed

In distributed multi-GPU training (DDP), each GPU stores a full copy of model weights, optimizer states, and gradients. When model size scales past 14B parameters, DDP exhausts memory even on 24GB GPUs.

Switch from DDP to Fully Sharded Data Parallel (FSDP):

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.fully_sharded_data_parallel import (
    CPUOffload,
    BackwardPrefetch,
)

model = FSDP(
    model,
    cpu_offload=CPUOffload(offload_params=True),
    backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
    device_id=torch.cuda.current_device(),
)

FSDP shards model parameters, gradients, and optimizer states across all available GPUs, reducing per-card memory footprints proportionally to worker counts.

---

5. Verification: Safe Memory Stress Test

Verify your memory configuration using an allocation loop:

import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
import torch

device = torch.device("cuda:0")
tensors = []

try:
    for i in range(10):
        # Allocate 2GB chunks dynamically
        t = torch.empty((500, 1024, 1024), dtype=torch.float32, device=device)
        tensors.append(t)
        print(f"Step {i+1}: Allocated {(i+1) * 2} GB VRAM successfully.")
except torch.cuda.OutOfMemoryError:
    print("Caught clean OOM. Testing cache recovery...")
    del tensors
    torch.cuda.empty_cache()
    print(f"Cache cleared. Current VRAM: {torch.cuda.memory_allocated() / (1024**3):.2f} GB")

Setting expandable_segments:True paired with explicit cache clearing ensures your training loops run stably without random allocator panics.