Where the memory actually goes

Five consumers, in order of surprise

When a run reports CUDA out of memory, people reach for a smaller batch size. Sometimes that is right. Often the batch was not the problem, and knowing which of five consumers is eating the card saves an afternoon.

1. Weights. Parameters times bytes per parameter. A 7B model is 14 GB in bf16, 28 GB in fp32, about 3.9 GB in 4-bit NF4. Fixed for the whole run, easy to compute, rarely the surprise.

2. Gradients. Same size as the trainable weights, in the gradient dtype. Full fine-tuning a 7B in bf16: another 14 GB. LoRA with 20M trainable parameters: 80 MB. This single line is why parameter-efficient methods exist.

3. Optimiser state. 8 bytes per trainable parameter for AdamW. 56 GB for a full 7B; 160 MB for the LoRA.

4. Activations. The intermediate tensors kept from the forward pass so the backward pass can use them. This is the one people underestimate, and it scales with things they change casually:

activations ≈ batch × seq_len × hidden × layers × bytes × c

where c is a small constant: roughly 10 to 20 depending on architecture and what is fused. For a 7B model (32 layers, hidden 4,096) at batch 4 and sequence 2,048 in bf16, that is on the order of 10 to 20 GB. It grows linearly in both batch and sequence length, so doubling your context doubles it.

5. Attention, if it is not flash attention. A naive attention implementation materialises the full score matrix: batch × heads × seq² × bytes. At sequence 8,192 with 32 heads in bf16 that is about 4.3 GB per layer, transiently. This is the quadratic term everybody has heard of and few have costed. FlashAttention computes attention in tiles without ever writing the matrix, turning that term from quadratic to linear in memory. Switch it on (attn_implementation="flash_attention_2", or the PyTorch SDPA backend) before you touch anything else at long context.

Gradient checkpointing, and what it actually trades

Gradient checkpointing throws away most activations after the forward pass and recomputes them during the backward pass from a small number of saved boundaries. Memory for activations drops from roughly linear in depth to roughly the square root of it; compute rises by about 30%.

It is nearly always worth it on a memory-bound run, because the memory you free usually buys a larger batch, and a larger batch recovers much of the lost throughput.

One trap: with LoRA and gradient checkpointing together, you may need model.enable_input_require_grads(). Without it the checkpointed segments see no input requiring grad, the graph is not built, and you get a run where the loss never changes and nothing obviously errors. This has cost a great many people an evening.

The logits, again

From the first lesson of this module: batch × seq × vocab can exceed the weights. At batch 4, sequence 2,048 and a 128k vocabulary that is 1 billion values: 2 GB in bf16, 4 GB in fp32, and typically both exist at once during the loss computation.

Fused cross-entropy kernels avoid materialising it. Liger Kernel and cut cross-entropy are the common options, both free and both a one-line change in TRL. On large-vocabulary models this is frequently the single largest saving available.

The diagnostic sequence

When you are out of memory, work down this list rather than guessing:

1. Print torch.cuda.max_memory_allocated() and compare against the card. If allocated is far below capacity, you have fragmentation, not a shortage: set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True and try again.
2. Is flash attention on? At sequence length over 4,000 this is usually the answer.
3. Is gradient checkpointing on?
4. Is the cross-entropy fused?
5. Only now reduce batch size, raising gradient accumulation to keep the effective batch constant.
6. Then reduce sequence length, and measure how many of your examples get truncated by the new limit, because a max length that cuts 30% of your completions is a data problem, not a memory fix.

The estimate before you rent

Do the arithmetic on paper first. For QLoRA on a 7B at sequence 2,048, batch 2: 3.9 GB weights, negligible gradients and optimiser state, perhaps 5 to 8 GB of activations with checkpointing off, 1 GB of logits. Around 11 GB, which fits a 16 GB T4 with room and comfortably fits a 24 GB card.

If the arithmetic says 70 GB and you were planning to use a 24 GB card, you learned that for free.