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.