RuntimeError: CUDA out of memory is the most-hated line in machine learning. You are mid-training, the epoch is crawling along, and the run dies because the GPU ran out of VRAM by a few hundred megabytes. The reflex is to rent a bigger GPU. You usually don't need to — three changes fix the large majority of cases, and none of them require new hardware.

🎬 The 60-second version is on Instagram — watch the reel. This is the full version with the mechanism behind each fix and the silent memory leak that quietly causes OOMs.

TL;DR

  • Gradient accumulation — shrink the batch, accumulate gradients over several steps, keep the same effective batch size.
  • Mixed precision — wrap the forward pass in torch.autocast("cuda", dtype=torch.bfloat16) to roughly halve activation memory.
  • Gradient checkpointing — recompute activations in the backward pass instead of storing them: trade compute for memory.
  • Avoid the silent leak — accumulate loss.item(), never the live loss tensor.

What "CUDA out of memory" actually means

GPU memory during training holds four things: the model weights, the gradients, the optimizer state (Adam keeps two extra tensors per parameter), and — usually the biggest and most variable — the activations saved from the forward pass so the backward pass can compute gradients. Weights, gradients, and optimizer state are fixed by your model size. Activation memory is not: it scales with batch size and sequence length. That is why the error so often appears when you bump the batch size, and why the batch is the first place to look.

One important nuance: OOM is an allocation failure, not necessarily a leak. PyTorch's caching allocator may also be holding fragmented free blocks it can't hand back for a large contiguous request. So the fixes fall into two buckets — genuinely reduce peak memory, and reduce fragmentation and accidental retention.

Fix 1: Gradient accumulation

Large effective batch sizes stabilize training, but they blow up activation memory. Gradient accumulation decouples the two. You run several small "micro-batches," call backward() on each so gradients add up in place, and only step the optimizer once every N micro-batches. Peak activation memory is set by the small micro-batch; the gradient update reflects the large effective batch.

python
ACCUM = 4  # effective batch = micro-batch size * 4

for i, batch in enumerate(loader):
    loss = model(batch) / ACCUM      # scale so accumulated grads average correctly
    loss.backward()
    if (i + 1) % ACCUM == 0:
        optimizer.step()
        optimizer.zero_grad()

The one easy-to-miss detail is dividing the loss by ACCUM. Because gradients sum across the accumulation window, you scale each contribution down so the result matches a true large-batch average rather than being N times too large.

Fix 2: Mixed precision

By default many training scripts keep activations in 32-bit float. Most of that precision is wasted. Mixed precision runs the forward pass in a 16-bit type — bfloat16 is the safe default on modern GPUs because it keeps float32's exponent range and rarely needs loss scaling — while keeping a master copy of the weights in float32. Halving the width of the activation tensors roughly halves the memory they consume, and it usually speeds training up too, since 16-bit matmuls run faster on tensor cores.

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

If you use float16 instead of bfloat16, add a GradScaler to prevent gradients from underflowing to zero. With bfloat16 you generally don't need it.

Fix 3: Gradient checkpointing

When a model is simply too deep to hold all its activations, gradient checkpointing is the escape hatch. Instead of storing every intermediate activation for the backward pass, it stores only a few checkpoints and recomputes the rest on the fly during backprop. You pay for it with extra compute — an additional partial forward pass — but the reduction in activation memory can be dramatic, and it is often what makes a large model fit at all.

python
# Hugging Face Transformers models expose it directly:
model.gradient_checkpointing_enable()

# For a custom nn.Module, wrap the expensive blocks:
from torch.utils.checkpoint import checkpoint
out = checkpoint(expensive_block, x)

The tradeoff is explicit: more time per step in exchange for a lower memory ceiling. On a memory-bound run that would otherwise not start, that is a trade worth making.

The silent leak that causes mystery OOMs

Sometimes memory climbs steadily across iterations until it dies — a genuine leak, and almost always the same bug. If you track the running loss by accumulating the tensor, you keep the entire autograd graph for every iteration alive, because the tensor still references it:

python
running_loss += loss          # ❌ retains the graph for every step — memory grows
running_loss += loss.item()   # ✅ detaches to a plain Python float

Two more knobs for stubborn cases:

  • torch.cuda.empty_cache() returns cached-but-unused blocks to the driver. It does not fix a true shortage, but it helps when you switch phases (say, training to evaluation) and want the cache released.
  • PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True reduces fragmentation, which is exactly the case where you have enough total free memory but not one contiguous block big enough for the next allocation.

Key Takeaways

  • Activation memory scales with batch size — so the batch is the first thing to cut.
  • Gradient accumulation keeps a large effective batch at a small memory footprint (remember to divide the loss by the accumulation steps).
  • bfloat16 mixed precision roughly halves activation memory and usually trains faster.
  • Gradient checkpointing trades extra compute for a much lower memory ceiling when a model won't otherwise fit.
  • Accumulate loss.item(), not the live tensor, or you will leak the autograd graph and OOM out of nowhere.

This is the deep-dive companion to a @teqnersinsights reel. Follow @teqnersinsights for the 60-second fixes, and browse the full reel library for more ML engineering deep-dives.

PyTorchCUDADeep LearningGPUMachine Learning