# Diagnose and fix PyTorch CUDA OOM and allocator fragmentation
Use this when a PyTorch training or inference job dies with `CUDA out of memory`,
or when `torch.cuda.memory_reserved()` climbs far above what the tensors seem to
need. Most CUDA OOMs are one of two distinct problems with different fixes, and
this procedure tells them apart before you change anything.
## 1. Read the error message carefully
The allocator's error tells you what it wanted and what it saw, e.g.:
```
CUDA out of memory. Tried to allocate 1.24 GiB (GPU 0; 15.78 GiB total capacity;
10.34 GiB already allocated; 435.50 MiB free; 14.21 GiB reserved in total by PyTorch)
```
- **Tried to allocate** — size of the single contiguous request that failed.
- **already allocated** — same as `torch.cuda.memory_allocated()`: memory held by live tensors.
- **reserved in total by PyTorch** — same as `torch.cuda.memory_reserved()`: everything the
caching allocator holds, including inactive cached blocks it could reuse.
- **free** — what the CUDA driver still reports free on the device.
## 2. Classify the failure: genuine OOM vs fragmentation
Compare the three numbers at failure time:
- `memory_allocated` ~ `memory_reserved` ~ total capacity, and "free" is small:
**genuine OOM**. The model, batch, or activations truly do not fit. Fix with
smaller batch size, gradient accumulation, gradient checkpointing
(`torch.utils.checkpoint`), mixed precision (`torch.amp`), or optimizer-state
offload.
- `memory_reserved` is much larger than `memory_allocated` and "Tried to allocate"
is much smaller than `reserved - allocated`: **fragmentation**. The allocator has
enough cached memory in total, but no single free block is big enough for the
requested contiguous chunk. This is where allocator tuning (step 4) helps.
- `memory_reserved` is small but the driver reports little free memory: something
outside PyTorch's allocator holds the memory (another process, a second model on
the same GPU, `cudaMalloc` outside the caching allocator). Check `nvidia-smi`.
Quick one-liner during a run:
```python
import torch
a = torch.cuda.memory_allocated() / 1024**3
r = torch.cuda.memory_reserved() / 1024**3
ma = torch.cuda.max_memory_allocated() / 1024**3
mr = torch.cuda.max_memory_reserved() / 1024**3
print(f"allocated={a:.2f} GiB reserved={r:.2f} GiB "
f"peak_allocated={ma:.2f} GiB peak_reserved={mr:.2f} GiB")
```
Call `torch.cuda.reset_peak_memory_stats()` before the region of interest and read
`max_memory_allocated()` / `max_memory_reserved()` after it to get the true peak of
that region — the counters record the peak since the last reset (or since CUDA
init), not the peak of the loop you care about. If you only want the model's
steady-state peak, reset right before the training loop starts.
## 3. Do not sprinkle `torch.cuda.empty_cache()` around
`empty_cache()` releases only segments that are entirely inactive back to the CUDA
driver. It does not reduce the peak of live tensors, it does not defragment split
blocks, and it does not make the next allocation succeed that would otherwise fail:
the allocator already reuses cached memory before asking the driver for more.
Calling it in a training loop typically makes things slower (each call forces the
allocator to re-acquire memory from the driver on the next step) without fixing OOMs.
The one legitimate use: a one-time call before a known-large allocation in a
fresh phase of the program (e.g. after teardown, before loading a large model),
or to release memory so a *different* process on the same GPU can use it.
## 4. Fix fragmentation with the allocator config
Tuning happens through the `PYTORCH_CUDA_ALLOC_CONF` environment variable, read
when the CUDA allocator initializes — it must be set **before the first CUDA
call in the process**:
```
PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 python train.py
```
The format is `OPTION:VALUE,OPTION2:VALUE2` (comma-separated pairs). Useful options:
- `max_split_size_mb` — prevents the allocator from splitting blocks larger than
this size (in MB). Helps when the OOM shows a large amount of inactive split
blocks. Default is unlimited. Per the docs, treat this as a last resort for a
workload that aborts with OOM and shows lots of inactive split blocks — it can
cost anywhere from zero to substantial performance depending on allocation
patterns. Common starting values: 128 or 512; retune with `torch.cuda.memory_stats()`
and `torch.cuda.memory_summary()`. Ignored when using `backend:cudaMallocAsync`.
- `roundup_power2_divisions` — rounds requested allocation sizes to the nearest
power-of-2 division instead of coarse buckets, improving block reuse for large,
similarly-sized allocations and wasting less capacity on near-miss sizes.
- `backend:cudaMallocAsync` — switch the allocator to CUDA's built-in async
allocator (requires CUDA 11.4 or newer) instead of PyTorch's native caching
allocator. Changes fragmentation behavior wholesale; `max_split_size_mb` is
ignored under this backend.
- `expandable_segments` — if your PyTorch build supports it, lets the allocator
request virtual address space and grow segments on demand, which can eliminate
some fragmentation patterns entirely.
Because the config is per-process, A/B test the failing command with and without
the option and compare the true peak (`max_memory_reserved`) — don't eyeball
`nvidia-smi` "reserved", which lags the allocator's internal state.
## 5. Find the true peak with a memory snapshot
When numbers disagree and you need ground truth, capture an allocator snapshot:
```python
torch.cuda.memory._record_memory_history()
run_your_code() # the loop or step that OOMs
torch.cuda.memory._dump_snapshot("my_snapshot.pickle")
```
Open the pickled file in the interactive viewer at `pytorch.org/memory_viz`
(it runs locally; nothing is uploaded). The viewer shows:
- **Active Memory Timeline** — every live tensor over time; hover a block for its
allocating stack trace. This answers "what is actually alive at the peak".
- **Allocator State History** — every allocator event; OOMs are recorded as `oom`
events with the full segment/block state at failure. This answers "why did the
request fail when reserved memory still existed" — e.g. all free blocks are too
small or in the wrong pool.
Each segment in the snapshot has `total_size` (what `cudaMalloc` returned),
`allocated_size`, and per-block states: `active_allocated`, `active_awaiting_free`
(waiting on another stream), and `inactive` (free for reuse). A failure with many
large `inactive` blocks is fragmentation; a failure where everything is
`active_allocated` is genuine OOM. `torch.cuda.memory_summary()` prints a
condensed version of the same data when you don't need the full trace.
## 6. Checklist for the failing command
1. Reproduce and capture: error message + allocated/reserved/free numbers.
2. Classify: fragmentation (reserved >> allocated) vs genuine OOM vs external holder.
3. If genuine: batch size / gradient accumulation / activation checkpointing /
mixed precision / optimizer offload — fix the math, not the allocator.
4. If fragmentation: set `PYTORCH_CUDA_ALLOC_CONF` before process start
(`max_split_size_mb:128` first), or try `backend:cudaMallocAsync`;
confirm with a snapshot that inactive split blocks were the problem.
5. If external holder: `nvidia-smi` to find the other process.
6. Never add `empty_cache()` to the loop as a fix.