The error
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.20 GiB is free. Including non-PyTorch memory, this process has 77.80 GiB
memory in use. Of the allocated memory 70.10 GiB is allocated by PyTorch, and 6.90 GiB is reserved
by PyTorch but unallocated.
The message tells you a lot: allocated (live tensors), reserved but unallocated (cached; large values suggest fragmentation), and whether memory is used outside PyTorch.
1. Is someone else on the GPU?
nvidia-smi
nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv
A leftover Jupyter kernel or a previous training run often holds the memory. Stop it. In Kubernetes, make sure two pods aren't sharing one GPU unintentionally.
2. Look inside the process
print(torch.cuda.memory_summary())
print(torch.cuda.max_memory_allocated() / 2**30, "GiB peak")
For hard cases, record a memory snapshot and open it at pytorch.org/memory_viz:
torch.cuda.memory._record_memory_history()
# ... run a few steps ...
torch.cuda.memory._dump_snapshot("snapshot.pickle")
3. Fragmentation
If "reserved but unallocated" is large:
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
4. Training fixes (in order of effort)
| Fix | Saving |
|---|---|
| Smaller batch + gradient accumulation to keep the effective batch size | Linear with batch size |
Mixed precision (torch.autocast(device_type="cuda", dtype=torch.bfloat16)) | Large, plus faster |
Gradient checkpointing (model.gradient_checkpointing_enable() in Hugging Face) | Large activation saving, ~20–30% slower |
| Shorter sequence length or packing | Big for transformers |
| 8-bit optimizers, LoRA/QLoRA instead of full fine-tuning | Large |
| FSDP / DeepSpeed ZeRO across GPUs | Shards weights, gradients and optimizer state |
5. Inference fixes
model.eval()
with torch.inference_mode():
out = model(batch)
Without inference_mode()/no_grad(), PyTorch keeps activations for a backward pass that never happens.
Also: load weights in bf16/fp16 or quantized, limit the batch size and max sequence length, and for LLM serving use an engine with paged KV cache (vLLM, TensorRT-LLM).
6. Memory that grows every step
A slow leak usually means you keep a reference to a tensor with its graph:
total_loss += loss # keeps the graph alive - leak
total_loss += loss.item() # correct
Note
torch.cuda.empty_cache()releases cached blocks to other processes but doesn't free memory still referenced by tensors. It rarely fixes an OOM on its own.