Fix RuntimeError: CUDA out of memory (PyTorch)

The error that stops every GPU training run eventually:

RuntimeError: CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0;
15.78 GiB total capacity; 13.52 GiB already allocated; 1.23 GiB free;
13.90 GiB reserved in total by PyTorch)

Most guides jump straight to "reduce your batch size" or "call torch.cuda.empty_cache()." The first is sometimes right; the second almost never helps. Before you change anything, the numbers in the error tell you which problem you actually have — and there are two very different ones, with different fixes.

This guide reads the error properly, separates fragmentation from genuine memory pressure, then walks the memory-reduction levers in order of impact.

Quick Answer

Read the numbers. If reserved ≫ allocated, it's fragmentation — set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. If allocated fills the card, you genuinely need less memory: lower batch size, use gradient accumulation, enable mixed precision (AMP), add gradient checkpointing, or quantize. torch.cuda.empty_cache() does not lower peak usage, so it rarely fixes OOM.

TL;DR

• Read allocated vs reserved vs free before touching code.
• reserved ≫ allocated → fragmentation → PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True.
• allocated ≈ total → genuinely out → reduce memory.
• Levers by impact: batch size ↓ → gradient accumulation → AMP → checkpointing → quantize → shard.
• empty_cache() is a myth for fixing OOM in your own run.
• Wrap eval in torch.no_grad(); use loss.item(), not total += loss.

Read the Error First

Four numbers matter:

The decisive comparison is reserved vs allocated:

• reserved ≈ allocated (like 13.90 vs 13.52 above) → the card is genuinely full. You must use less memory.
• reserved ≫ allocated (e.g. 13.90 reserved but only 6 allocated) → fragmentation. PyTorch holds lots of memory in small cached blocks, none big enough for the new allocation. Fix the allocator, not the model.

The error even hints at it: "If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation."

The Fix Decision

Step 0: Rule out another process

Before assuming your code is the problem, check who's on the GPU:

nvidia-smi

If another PID holds most of the memory (a zombie run, a notebook you forgot, someone else's job), that's your fix — kill it or move to a free GPU with CUDA_VISIBLE_DEVICES=1.

Step 1: Fragmentation → tune the allocator

If reserved ≫ allocated, let the allocator manage memory more flexibly. Set this before the process starts:

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python train.py

• expandable_segments:True → segments grow instead of fragmenting; the modern, usually-best option.
• Older alternative: PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 caps block splitting, which also reduces fragmentation.
• This can't create memory — it only helps when memory is reserved-but-unusable.

Step 2: Genuinely out → reduce memory, in order of impact

When the card is truly full, apply these from cheapest to most involved.

Lower the batch size (biggest, simplest lever)

Activation memory scales roughly linearly with batch size. Halving the batch roughly halves activation memory:

# 32 → 8 if you're OOMing
train_loader = DataLoader(dataset, batch_size=8, shuffle=True)

Gradient accumulation (keep the effective batch)

Don't want to lose the large effective batch? Run small micro-batches and step the optimizer every N of them:

accum_steps = 4                      # effective batch = micro_batch * 4
optimizer.zero_grad()
for i, (x, y) in enumerate(loader):
    loss = model(x, y) / accum_steps # scale so grads average correctly
    loss.backward()
    if (i + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

• Peak memory tracks the micro-batch, not the effective batch.
• loss / accum_steps keeps the gradient magnitude equivalent to a full batch.

Mixed precision (AMP) — ~half the activation memory

Run the forward/backward in float16/bfloat16 where safe:

from torch.cuda.amp import autocast, GradScaler
 
scaler = GradScaler()
for x, y in loader:
    optimizer.zero_grad()
    with autocast():                 # fp16/bf16 ops
        loss = model(x, y)
    scaler.scale(loss).backward()    # scaled to avoid underflow
    scaler.step(optimizer)
    scaler.update()

• autocast() casts eligible ops to half precision → smaller activations.
• GradScaler prevents fp16 gradient underflow.

Gradient checkpointing (trade compute for memory)

Recompute activations during backward instead of storing them — big memory savings for deep models, at ~20–30% more compute:

# Hugging Face models
model.gradient_checkpointing_enable()
 
# Raw PyTorch
from torch.utils.checkpoint import checkpoint
out = checkpoint(expensive_block, x)

Quantize or use a smaller model

For inference or QLoRA-style training, load in 8-bit/4-bit to cut weight memory dramatically:

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
bnb = BitsAndBytesConfig(load_in_4bit=True)
model = AutoModelForCausalLM.from_pretrained("model", quantization_config=bnb)

If quantization fails to initialize, see fixing bitsandbytes CUDA Setup failed.

Shard across GPUs

When one model won't fit any single card, split it: FSDP (torch.distributed.fsdp) or DeepSpeed ZeRO shard parameters, gradients, and optimizer state across devices.

The Silent Leaks Everyone Hits

Two coding mistakes cause OOM that no batch-size change fixes.

Retaining the graph during eval

Inference and validation don't need gradients. Without no_grad, PyTorch stores the whole graph:

# ❌ builds and keeps the autograd graph for every batch
for x, y in val_loader:
    preds = model(x)
 
# ✅ no graph, far less memory
with torch.no_grad():                # or torch.inference_mode()
    for x, y in val_loader:
        preds = model(x)

Accumulating tensors with grad history

total_loss += loss keeps every batch's graph alive because loss still carries history. Detach the scalar:

# ❌ leaks — loss keeps the graph
total_loss += loss
 
# ✅ store the plain number
total_loss += loss.item()            # or loss.detach()

The same applies to appending tensors to a list for logging — call .item() or .detach().cpu() first.

Debugging: See What's Using Memory

import torch
print(torch.cuda.memory_allocated() / 1e9, "GiB allocated")
print(torch.cuda.max_memory_allocated() / 1e9, "GiB peak")
print(torch.cuda.memory_summary())   # full breakdown by size class

• max_memory_allocated() is the number that must fit under your card's capacity.
• Reset the peak between phases with torch.cuda.reset_peak_memory_stats().
• On GPU, an OOM traceback can point at the wrong line due to async execution — see the device-side assert guide for the same CUDA_LAUNCH_BLOCKING=1 technique.

Verification Steps

1. torch.cuda.max_memory_allocated() stays comfortably under total capacity.
2. One full epoch (train + validation) completes without OOM.
3. Memory doesn't creep upward across iterations — a rising memory_allocated() over steps means a leak (retained graph or tensors).
4. With expandable_segments:True, reserved tracks allocated closely.
5. nvidia-smi shows only your process on the target GPU.

Prevention

• Wrap all eval/inference in torch.no_grad() / inference_mode().
• Log losses with .item(), never accumulate live tensors.
• Set a batch size that leaves headroom; use accumulation for large effective batches.
• Enable AMP by default for training on modern GPUs.
• Track max_memory_allocated() and alert if it nears capacity.
• Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True for long/large jobs.
• Pin one job per GPU, or partition with CUDA_VISIBLE_DEVICES.

Troubleshooting Matrix

Related Guides

• Fix "bitsandbytes CUDA Setup failed"
• Resolve CUDA error: device-side assert triggered
• Fix torch.cuda.is_available() returns False
• Fix "Token indices sequence length is longer than the specified maximum"
• A practical guide to 4-bit QLoRA fine-tuning

External References

• PyTorch — CUDA memory management
• PyTorch — Automatic mixed precision (AMP)
• PyTorch — torch.utils.checkpoint
• PyTorch — FSDP

FAQs

Is CUDA out of memory a bug in PyTorch? No — it means your workload's peak memory exceeded the GPU's capacity, or memory fragmented. It's a resource/config issue you resolve by using less memory or tuning the allocator, not a framework bug.

Why does the first iteration OOM but not later ones? Peak memory is usually hit on the first forward/backward when activations and the graph are largest, or on the first batch with the longest sequence. If it survives the first iteration, it typically survives the rest — unless you have a leak.

Does a smaller learning rate help OOM? No. Learning rate doesn't affect memory. Only the sizes of tensors (batch, sequence, model, activations) and whether you keep the graph do.

Should I use bfloat16 or float16 for AMP? On Ampere and newer GPUs, bfloat16 is more numerically stable and often needs no GradScaler. On older cards, float16 with GradScaler is the path. Both roughly halve activation memory.

Can I catch the OOM and retry with a smaller batch? You can, but after an OOM the CUDA context may be in a poor state; it's cleaner to size the batch correctly up front (or use accumulation). If you do retry, call torch.cuda.empty_cache() first — this is the one place it genuinely helps, by releasing the failed allocation's cache before the smaller retry.