Debugging CUDA Out-of-Memory Errors in PyTorch: Memory Snapshots, Activation Checkpointing and Offload

"CUDA out of memory" is the most common failure in PyTorch training, and the most commonly mis-fixed. Teams reach for a bigger GPU before they know what is using the memory. This tutorial gives a repeatable workflow: measure first, identify which category of memory dominates, then apply the cheapest fix for that category.

1. Know the four buckets

GPU memory during training is made of:

• Parameters: roughly n_params x bytes_per_param.
• Gradients: same size as parameters.
• Optimizer state: Adam keeps two extra fp32 tensors per parameter, so it is often the largest fixed cost.
• Activations: saved tensors for backward. These scale with batch size and sequence length, and are usually what causes the spike.

A 7B-parameter model with bf16 weights and gradients and fp32 Adam states needs about 14 + 14 + 56 GB before a single activation is stored. Do this arithmetic first; it tells you whether the problem is fixed cost or activations.

2. Read the error message

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB.
GPU 0 has a total capacity of 79.15 GiB of which 1.2 GiB is free.
Of the allocated memory 60.1 GiB is allocated by PyTorch,
and 15.3 GiB is reserved by PyTorch but unallocated.

If reserved-but-unallocated is large, you have fragmentation, not a true shortage. Try:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Expandable segments often remove the fragmentation failure with no code change.

3. Capture a memory snapshot

PyTorch can record every allocation with its stack trace:

import torch

torch.cuda.memory._record_memory_history(max_entries=100_000)

try:
    for step, batch in enumerate(loader):
        loss = model(**batch).loss
        loss.backward()
        opt.step(); opt.zero_grad(set_to_none=True)
        if step == 3:
            break
finally:
    torch.cuda.memory._dump_snapshot("snapshot.pickle")
    torch.cuda.memory._record_memory_history(enabled=None)

Load snapshot.pickle at https://pytorch.org/memory_viz. The timeline shows memory climbing through the forward pass and falling during backward. The peak is at the forward/backward boundary; the stacks on the tallest bands tell you which layers own the activations.

For a quick numeric check without the viewer:

torch.cuda.reset_peak_memory_stats()
# ... one training step ...
print(torch.cuda.max_memory_allocated() / 2**30, "GiB peak")

4. Fixes, cheapest first

Free accidental retention

• Log loss.item(), never loss; storing the tensor keeps the whole graph alive.
• Use zero_grad(set_to_none=True).
• Wrap evaluation in torch.inference_mode().

Mixed precision

with torch.autocast("cuda", dtype=torch.bfloat16):
    loss = model(**batch).loss

bf16 halves activation memory for most layers and needs no loss scaling.

Gradient accumulation

Halve the micro-batch and accumulate to keep the effective batch size:

accum = 4
for i, batch in enumerate(loader):
    loss = model(**batch).loss / accum
    loss.backward()
    if (i + 1) % accum == 0:
        opt.step(); opt.zero_grad(set_to_none=True)

Activation checkpointing

Trade compute for memory by recomputing activations in backward. It typically costs 20-35% more step time for a large activation saving:

from torch.utils.checkpoint import checkpoint

class Block(torch.nn.Module):
    def forward(self, x):
        return checkpoint(self.inner, x, use_reentrant=False)

Checkpoint per transformer block, not the whole model. Always pass use_reentrant=False.

Shrink fixed costs

• Use fused or 8-bit optimizers, or shard states with FSDP2 (fully_shard).
• Use LoRA/QLoRA so optimizer state exists only for adapter weights.
• As a last resort, offload optimizer state to CPU; expect a throughput penalty.

Attention and logits

Use F.scaled_dot_product_attention so attention memory stays linear in sequence length. For large vocabularies, the logits tensor (batch x seq x vocab) can be several GB; compute the loss in chunks.

5. A decision table

Symptom	Likely cause	First fix
Large reserved-unallocated gap	Fragmentation	expandable_segments:True
OOM at step 0	Fixed cost too large	Shard or use LoRA
OOM after N steps	Retained graph or longer batch	Check .item(), cap sequence length
OOM only on long sequences	Activations	Checkpointing, SDPA
OOM in eval only	Missing no-grad	inference_mode()

6. Make it a guardrail

Add a CI smoke test that runs three steps at the maximum sequence length and asserts max_memory_allocated stays under a budget. Memory regressions then surface in a pull request rather than at 3 a.m. in a multi-day run.

Takeaways

Measure before you scale up hardware. Classify the memory into parameters, gradients, optimizer state and activations, and apply the matching fix. Most OOMs disappear with bf16, checkpointing and sharded optimizer state, and you keep the GPU budget you planned.

If your team is stuck on memory or throughput in a PyTorch project, IntelliSensei's consultants can profile it and recommend a fix. Get in touch via our contact page.