Surviving GPU Out-of-Memory — Every Cause, Every Fix

`CUDA out of memory` reports only the allocation that happened to fail last, which is almost never the culprit. We count what actually occupies VRAM in five buckets, derive the sixteen-bytes-per-parameter fixed cost of training, and work through the fixes in order of least damage: gradient checkpointing, optimizer compression, offloading, KV cache limits, and fragmentation.

Training Deep Nets with Sublinear Memory Cost

The truck bed is sized before you drive off

A moving truck's bed is whatever size you booked, and it does not get bigger once you are on the road. If the load does not fit, you repack it, throw things away, or hire a second truck. Those are the only three moves.

GPU memory is that truck bed. And if you run training or inference for long enough, you will meet this:

torch.cuda.OutOfMemoryError: CUDA out of memory.
Tried to allocate 2.00 GiB (GPU 0; 79.15 GiB total capacity;
74.23 GiB already allocated; 1.02 GiB free;
77.31 GiB reserved in total by PyTorch)

The message looks helpful and is not. It reports the 2 GB allocation that happened to fail last, while the culprit is almost certainly hiding in the 74 GB already sitting there. Shaving 2 GB off is usually wasted effort.

So there is exactly one correct first step: produce the breakdown of what is using how much. Once you have it, the fix picks itself.

Five buckets eat your VRAM

Everything resident on the GPU during training falls into five categories.

1. Weights — parameter count times bytes per parameter. Present from first step to last
2. Gradients — an array shaped exactly like the weights, filled during the backward pass
3. Optimizer state — for Adam, first and second moments. Two more arrays shaped like the weights
4. Activations — intermediate forward-pass results, held because the backward pass needs them
5. Temporary buffers, fragmentation, CUDA context — matmul workspaces, communication buffers, and the few hundred MB the CUDA runtime takes just by existing

The decisive fact is that buckets 1 through 3 are fixed the moment you pick a model and an optimizer — dropping the batch size to 1 does not free a single byte of them. Only bucket 4 moves. That asymmetry is why every OOM discussion starts with "lower the batch size" and then stalls at "I set it to 1 and it still dies."

Training's fixed cost: sixteen bytes per parameter

$\Psi$ is the number of parameters, $b_w$ the bytes per weight, $b_g$ the bytes per gradient, $b_{\text{opt}}$ the bytes of optimizer state per parameter. In words: the parameter count, times whatever it costs to feed one parameter.

Plain fp32 with Adam: 4 + 4 + 8 = 16 bytes per parameter. Switch to mixed precision and it halves, right? It does not. bf16 weights 2, bf16 gradients 2, an fp32 master copy 4, two Adam moments 8 — still 16 bytes, the same accounting the ZeRO paper uses.

This trips people constantly. You move to half precision expecting the footprint to halve, and the fixed cost does not budge. What mixed precision shrinks is bucket 4, and its original purpose was never memory at all — it was getting the matmuls onto Tensor Cores (Mixed Precision Training).

For 7 billion parameters, $7\times10^9 \times 16 = 112$ GB. It will not fit on a single 80 GB card before you store a single byte of activations. This is where the intuition that "7B is small" first betrays you.

At inference, weights are fixed cost and the KV cache is variable

Switch to inference and gradients and optimizer state both vanish. Load 7B in bf16 and the weights are 14 GB, leaving 66 GB free. What fills those 66 GB is almost entirely the KV cache.

The leading 2 is for K and V, $L$ is the layer count, $H_{kv}$ the number of KV heads, $d_{\text{head}}$ the per-head dimension, $b$ the bytes per element, $T$ the total number of tokens alive on the server. In words: for every layer and every head, keep one K and one V per token.

With 32 layers, 32 KV heads, head dimension 128, in bf16, that is $2\times32\times32\times128\times2 = 524{,}288$ bytes per token — about 512 KB. A single 4096-token conversation costs 2 GB, so 66 GB of headroom is 33 conversations and request 34 waits. GQA, which cuts the KV head count specifically, became standard because it attacks $H_{kv}$ directly (Understanding the KV Cache from Scratch).

Activations are the only bucket that moves

Activations accumulate as a term proportional to batch × sequence length × hidden size, plus an attention term proportional to batch × heads × sequence length². FlashAttention removed that second, quadratic term from memory, and that is roughly half the reason long-context training became practical.

Activations are also the only one of the five that bends to your wishes. So every OOM fix reduces to one of two families: (A) shrink the activations, or (B) split, compress, or evict the fixed cost. What follows is ordered not by raw effect but by how little damage each does. If the first two suffice, stop there.

Fix 1: Cut the batch, refill with gradient accumulation

Activations scale roughly linearly with batch size, so halving it halves them. Small batches make gradients noisy, but gradient accumulation buys the statistics back: run four micro-batches of 4, letting gradients sum, then step once. Mathematically that is the gradient of a batch of 16.

optimizer.zero_grad()
for i, batch in enumerate(loader):
    loss = model(**batch).loss / ACCUM        # divide: you want the mean
    loss.backward()                            # gradients accumulate into .grad
    if (i + 1) % ACCUM == 0:
        optimizer.step(); optimizer.zero_grad()

Forget the / ACCUM and your effective learning rate jumps by the accumulation factor — entirely silently. And the equivalence is not universal: with BatchNorm statistics, or contrastive objectives that use the rest of the batch as negatives, micro-batching plus accumulation is a different computation from one large batch.

Fix 2: Gradient checkpointing — recompute instead of remember

Activations are held because the backward pass needs them, so throw them away and rebuild them by rerunning the forward pass when the moment comes. Split an $n$-layer network into segments, store only the boundaries, and pick $\sqrt{n}$ segments:

Meaning: what used to grow in proportion to depth now grows only with the square root of depth. At 100 layers you pay for something the size of 10. The bill is one extra forward pass — forward and backward cost roughly 1:2, so 3 units become 4 — but the real slowdown depends on implementation and model shape. Measure it yourself.

model.gradient_checkpointing_enable()   # Hugging Face
model.config.use_cache = False          # ← without this it warns and does nothing

That second line is the trap. use_cache=True (the generation KV cache) is incompatible with checkpointing. You get a warning, execution continues, and you discover it as "I turned it on and memory didn't drop." Calling PyTorch directly, use torch.utils.checkpoint.checkpoint(fn, x, use_reentrant=False); the older reentrant path has subtle problems with RNG reproducibility and hooks.

Fix 3: Attack the optimizer state

Eight of those sixteen bytes — half the fixed cost — are Adam's two moments. Three ways in:

• 8-bit optimizers: quantize the moments block-wise, turning 8 bytes into 2 (Dettmers et al.). Swapping in bitsandbytes' AdamW8bit is a one-line change, so try it first
• Change optimizer: SGD with momentum keeps one moment, plain SGD none. But this changes how training behaves — it is an optimization decision wearing a memory costume
• Shard with ZeRO / FSDP: with $N$ GPUs, split the optimizer state $N$ ways so each card holds $1/N$ (stage 1). Stage 2 shards gradients too, stage 3 the parameters as well

On that last point, kill a very common misconception: plain DDP saves no memory whatsoever. Every GPU holds a complete copy of model, gradients, and optimizer state, and only the gradients are averaged. Going from one card to eight leaves the per-card fixed cost identical. That is the answer to "I added GPUs and it still won't fit," and the entire reason ZeRO and FSDP exist.

Fix 4: Offload — to CPU, then to NVMe

If it still does not fit, move whatever you are not currently touching off the GPU. Optimizer state is only read at update time, so for the rest of the step it can live in host RAM. That is ZeRO-Offload; push it onto an NVMe SSD and it is ZeRO-Infinity.

# DeepSpeed
zero_optimization:
  stage: 2
  offload_optimizer: { device: cpu, pin_memory: true }

The effect is dramatic and so is the bill: your bottleneck moves to PCIe bandwidth. There are orders of magnitude between on-package GPU memory bandwidth and the bus, so shuttling tens of gigabytes across it every step costs exactly what you would expect. This is the move when you must fit on one card and will pay in wall-clock time. It is not a default.

Fix 5: Train fewer weights in the first place

Everything above assumed full fine-tuning. Drop that assumption and the structure of the fixed cost changes rather than shrinking. LoRA freezes the base weights and trains only low-rank update matrices; frozen weights need no gradients and no optimizer state, so 14 of the 16 bytes disappear. Where full fine-tuning of 7B needed 112 GB, LoRA needs 14 GB plus change and fits on one 80 GB card with activations to spare.

The LoRA paper, explained covers why the low-rank approximation holds up. One caveat: LoRA does not reduce activations. The forward pass still traverses every layer, so if you are dying on long sequences you will still die. That is checkpointing's job.

Inference OOM is almost always KV cache management

When a serving process dies, the KV cache did it nine times out of ten. The knobs are well-defined (vLLM's names below).

• --max-model-len — maximum context per request, capping $T$ directly. There is no reason to permit lengths you never serve
• --max-num-seqs — cap on concurrent sequences
• --gpu-memory-utilization — what fraction of VRAM to claim. Lower it if anything else shares the card
• --kv-cache-dtype fp8 — 8-bit KV cache, taking $b$ from 2 to 1 and halving the term

The structural improvement is PagedAttention. Classic implementations reserved a contiguous block sized for the maximum length per request, so most of that space died unused. PagedAttention manages the cache in pages, the way an OS manages virtual memory, allocating only what is actually written. The number of concurrent requests one card can serve changes substantially.

Fragmentation — free space you cannot use

The most unsettling OOM: "10 GB free, and a 2 GB allocation fails."

PyTorch's allocator caches blocks rather than returning them to the driver. Repeated allocation and release leave the free space in slivers, so the total is sufficient but no contiguous 2 GB exists. Variable-length batches make it likely, and the symptom is characteristic: the first few hundred steps are fine, then it dies out of nowhere. Diagnosis is two numbers.

torch.cuda.memory_allocated()   # what tensors actually occupy
torch.cuda.memory_reserved()    # what PyTorch has borrowed from the driver

A wide gap means fragmentation, and the fix is one environment variable: PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. You will be tempted to reach for torch.cuda.empty_cache() — resist. All it does is hand the cache back, which makes the next allocation slower and solves nothing if you are genuinely short. Inside a loop it is almost always a mistake.

Also, do not read nvidia-smi as gospel. It shows what PyTorch has borrowed from the driver (roughly reserved), not what tensors hold.

How this shows up on the job

Who touches it, and when. Platform and MLOps engineers sizing a new model against the GPUs they already have. Research engineers whose run just died at hour three. Serving engineers deciding the concurrency ceiling — and that third one is the infrastructure bill in disguise.

A routine worth fixing in place. (1) Before training, do $\Psi \times 16$ bytes in your head and check it can fit at all. (2) When it dies, get the breakdown from torch.cuda.memory_summary() and max_memory_allocated(). (3) If the fixed cost is the problem, go 8-bit optimizer → ZeRO → LoRA; if activations are, go checkpointing → smaller batch plus accumulation. (4) If reserved and allocated have drifted apart, set expandable_segments:True. Guessing your way down the batch size is how afternoons disappear.

The traps that cost real time.

• Surviving step 1 proves nothing. With variable-length data you die when the longest sample finally comes up. Run one dummy step at your maximum sequence length first and you fail after three seconds instead of three hours
• Forgetting torch.no_grad() in the eval loop. All activations stay alive, so validation eats more than training. model.eval() only switches Dropout and BatchNorm — it does not stop graph retention. You need both
• losses.append(loss) retains the whole graph. Use .item() or .detach(). It grows slowly and kills you hundreds of steps later, which makes it the hardest of these to trace
• DDP saves no memory (above). Reach for ZeRO/FSDP before reaching for more cards
• Check for squatters on the card. A crashed Jupyter kernel or a previous run may still be holding VRAM

How it gets asked in design reviews. "You want to fully fine-tune a 7B model — how many 80 GB GPUs?" Walk the three facts in order: the fixed cost is 112 GB, DDP does not divide it, ZeRO stage 3 does. The number matters less than whether you can decompose the question.

Summary

• The OOM message reports only the last allocation. The culprit is what was already loaded
• Training's fixed cost is 16 bytes per parameter with Adam, and mixed precision does not reduce it
• Only activations move freely, so every fix is either "shrink activations" or "split, compress, evict the fixed cost"
• Inference OOM is the KV cache. Which symbol of $2 \cdot L \cdot H_{kv} \cdot d_{\text{head}} \cdot b \cdot T$ you attack decides the fix
• "Free but unusable" is fragmentation. Spot it in the gap between reserved and allocated

Unlike speed, memory is binary: it fits or it does not. Which is exactly why you can settle it with arithmetic before you start.