# How do you fix CUDA out of memory errors in PyTorch?

> Source: <https://dev.to/viswaretas_kotra/how-do-you-fix-cuda-out-of-memory-errors-in-pytorch-48ga>
> Published: 2026-11-14 16:41:32+00:00

Fix a CUDA out of memory error by measuring what holds the memory, then cutting the biggest consumer first: leaked tensors, fragmentation, activations, and finally the weights, gradients and optimizer state. When the arithmetic says the job still cannot fit, Nodus runs the same script on a GPU that can. Set a memory floor such as `minMemory: 80Gi` instead of a GPU name, and Nodus picks the accelerator that is cheapest to finish the job.

This post gives one order of fixes, the gigabytes each step buys when fine-tuning Llama 3.1 8B on 4,096-token sequences, and code you can paste. The figures are our own arithmetic from the model's config and published formulas. They are estimates, so check them against your own peak numbers.

The error message already says which kind of problem you have. PyTorch's [caching allocator](https://github.com/pytorch/pytorch/blob/main/c10/cuda/CUDACachingAllocator.cpp) reports the size it tried to allocate, the GPU's total capacity, how much is free, how much is "allocated by PyTorch" and how much "is reserved by PyTorch but unallocated".

Read those numbers as a decision:

`ps -elf | grep python`.
Do not trust `nvidia-smi` for this. The [CUDA semantics notes](https://docs.pytorch.org/docs/stable/notes/cuda.html#memory-management) explain that memory the caching allocator keeps for reuse still shows as used there. Measure the peak inside your script instead, around one full training step:

``` python
import torch

torch.cuda.reset_peak_memory_stats()
loss = model(**batch).loss
loss.backward()
optimizer.step()
optimizer.zero_grad(set_to_none=True)

print(f"peak allocated: {torch.cuda.max_memory_allocated() / 1e9:.1f} GB")
print(f"peak reserved:  {torch.cuda.max_memory_reserved() / 1e9:.1f} GB")
print(torch.cuda.memory_summary(abbreviated=True))
```

If you need to see which line allocated what, the [memory snapshot guide](https://docs.pytorch.org/docs/stable/torch_cuda_memory.html) records history with `torch.cuda.memory._record_memory_history()`, dumps it with `_dump_snapshot("my_snapshot.pickle")`, and opens it in a visualizer that runs locally.

Into four buckets, and the order of fixes follows their size. Llama 3.1 8B has 8,030,261,248 parameters in its [Hugging Face](https://huggingface.co/meta-llama/Llama-3.1-8B) metadata, with 32 layers, a hidden size of 4,096 and a 128,256-token vocabulary.

The [ZeRO paper](https://arxiv.org/abs/1910.02054) puts mixed-precision training with Adam at 16 bytes per parameter: 2 for 16-bit weights, 2 for gradients and 12 for the FP32 master weights plus Adam's two moments. For 8.03 billion parameters that is 128.5 GB before a single activation.

Activations come from the formula in [Korthikanti et al.](https://arxiv.org/abs/2205.05198): s x b x h x (34 + 5as/h) bytes per layer, for sequence length s, micro-batch b, hidden size h and a attention heads. It was derived for a GPT-style block, and Llama's block differs a little, so treat the results as estimates.

| Bucket | Formula | Llama 3.1 8B, one 4,096-token sequence | 
|---|---|---|
| 16-bit weights | 2 bytes x params | 16.1 GB | 
| Gradients + AdamW state | 14 bytes x params | 112.4 GB | 
| Activations, fused attention | 34 x s x b x h x 32 layers | 18.3 GB | 
| Output logits in FP32 | s x 128,256 x 4 bytes | 2.1 GB | 
| Total |  | 148.8 GB | 

Rows are rounded, so they add to 148.9 GB rather than the exact total.

Fused attention matters here. The 5as/h term is the stored attention matrix, and kernels such as FlashAttention, which PyTorch's [`scaled_dot_product_attention`](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) can call, do not keep it. With an eager attention implementation that term stays, and activations for the same sequence grow to 104.2 GB. If your model has an attention implementation setting, make sure it uses SDPA or FlashAttention before anything else.

The cheapest one that targets your biggest bucket. Here is the order we use, with what each step buys on the example above.

*Checkpointing removes most of the activations, but on a full fine-tune the optimizer state is the bucket that decides which GPU you need.*

Look for tensors that keep their autograd history alive. The [PyTorch FAQ](https://docs.pytorch.org/docs/stable/notes/faq.html) gives the classic case: `total_loss += loss` keeps the graph of every step, so memory grows until it runs out. Write `total_loss += float(loss)` instead.

Three more patterns are common. Lists of outputs or attention maps kept for logging hold GPU memory until they are deleted, so `del` them or move them to the CPU. Evaluation run without disabling autograd stores activations nobody will use; wrap it in [`torch.inference_mode()`](https://docs.pytorch.org/docs/stable/generated/torch.autograd.grad_mode.inference_mode.html), which is like `no_grad` with less overhead. And a retry in an `except` block often cannot allocate either, because the exception still references the stack frame that owns the tensors; set a flag and retry outside the `except`.

`torch.cuda.empty_cache()` does not fix any of these. Its [documentation](https://docs.pytorch.org/docs/stable/generated/torch.cuda.memory.empty_cache.html) says it "doesn't increase the amount of GPU memory available for PyTorch". It returns cached blocks to other processes, which only helps when you share the GPU.

Turn on expandable segments before CUDA starts:

```
PYTORCH_ALLOC_CONF=expandable_segments:True python train.py
```

The [CUDA semantics notes](https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-alloc-conf) describe the problem it solves. When tensor sizes change from step to step, as with variable sequence lengths or batch sizes, the allocator leaves unusable slivers at the end of its segments. Expandable segments grow one segment instead of making new ones. The option is marked experimental and is off by default.

`PYTORCH_CUDA_ALLOC_CONF`, the name the error message prints, is now an alias kept for backward compatibility.

Activations scale with micro-batch size, so start there. A micro-batch of 8 sequences needs 8 x 18.3 GB, or 146.0 GB, for activations alone. Run micro-batch 1 and accumulate gradients over 8 steps, and the optimizer sees the same 8-sequence batch for 18.3 GB:

```
accum_steps = 8
for step, batch in enumerate(loader):  # loader yields micro-batches of 1
    with torch.autocast("cuda", dtype=torch.bfloat16):
        loss = model(**batch).loss / accum_steps
    loss.backward()
    if (step + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
```

The `autocast` line is the second fix. The [AMP docs](https://docs.pytorch.org/docs/stable/amp.html) describe autocast running ops such as linear layers in a 16-bit type. The formula above already assumes 16-bit activations; in plain FP32 they double.

Mixed precision does not shrink the 16 bytes per parameter, because the optimizer keeps FP32 copies either way. If you use FP16 instead, the docs pair autocast with `torch.amp.GradScaler`.

The third fix is gradient checkpointing. [`torch.utils.checkpoint`](https://docs.pytorch.org/docs/stable/checkpoint.html) keeps only each layer's input and recomputes the rest during the backward pass. With one checkpoint per layer, Korthikanti et al. put stored activations at 2 x s x b x h x layers bytes, which is 1.1 GB here instead of 18.3 GB:

``` python
from torch.utils.checkpoint import checkpoint

def forward(self, x):
    for layer in self.layers:
        x = checkpoint(layer, x, use_reentrant=False)
    return x
```

It costs compute. The same paper measured 30 to 40% more execution time with full recomputation.

*Micro-batch size and the attention kernel move activations by more than 100 GB before you touch the model.*

Pick by what you are allowed to change. On the example, checkpointing still leaves 131.7 GB, almost all of it model state.

An 8-bit optimizer is a one-line change. [bitsandbytes](https://huggingface.co/docs/bitsandbytes/main/en/optimizers) replaces `torch.optim.Adam` with `bnb.optim.Adam8bit` and stores each moment in 1 byte instead of 4. That takes training from 16 to 10 bytes per parameter, and the 8B example to 83.5 GB. It is still too much for an 80 GB GPU, which is why this step rarely ends the problem alone.

If adapters are enough for your task, LoRA freezes the base weights and trains small matrices. With rank 16 on every linear layer, Llama 3.1 8B gets 41.9 million trainable parameters, about 0.7 GB with their gradients and optimizer state. The run then needs about 19.9 GB.

QLoRA from [Dettmers et al.](https://arxiv.org/abs/2305.14314) stores the frozen base in 4 bits, which brings the example to about 7.9 GB, a little more in practice because embedding and output layers usually stay 16-bit. Our [LoRA vs QLoRA comparison](https://www.nodus-compute.ai/blog/lora-vs-qlora/) covers when the 4-bit base is worth its speed cost.

If you need every weight to train, shard instead. FSDP with full sharding divides the 16 bytes per parameter by the number of GPUs, so 8 GPUs hold about 19.2 GB each, before all-gather buffers. PyTorch's current API is [`fully_shard`](https://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html), and our [multi-GPU training guide](https://www.nodus-compute.ai/blog/multi-gpu-training-pytorch/) has a working script and the launch command.

When the arithmetic still does not fit after the fixes you are willing to make. Swapping full fine-tuning for LoRA is a modelling decision, not a memory trick.

Leave about 10% of the card for the CUDA context and allocator overhead, the same reserve we use in our [70B VRAM guide](https://www.nodus-compute.ai/blog/vram-needed-for-70b-model/). Against the GPUs on the Nodus pricing page, the smallest card that holds each setup is:

| Setup | Estimated need | Smallest listed GPU with 10% headroom | Listed per hour | 
|---|---|---|---|
| QLoRA | 7.9 GB | RTX A4000 16 GB | $0.173 | 
| LoRA | 19.9 GB | RTX 3090 24 GB | $0.180 | 
| Full, 8-bit Adam, checkpointing | 83.5 GB | RTX PRO 6000 96 GB | $1.350 | 
| Full, AdamW, checkpointing | 131.7 GB | B200 180 GB | $5.500 | 

Prices checked November 14, 2026, listed on the [Nodus pricing page](https://www.nodus-compute.ai/pricing/). The H200 at 141 GB would hold the last row with only 7% to spare, which is too tight for us.

*The memory fixes move the hourly price by more than 30 times, before counting how fast each card trains.*

The cheapest card per hour is not always the cheapest per run, because a slower GPU runs longer. On Nodus you do not have to choose. Set `resources.gpu.minMemory` instead of a GPU type in the [Job manifest](https://www.nodus-compute.ai/docs/guides/jobs/), and Nodus picks any accelerator with at least that much memory per GPU, choosing among them by expected cost to finish. For the 83.5 GB row, 88 GiB (94.5 GB) admits the RTX PRO 6000 and everything larger:

```
spec:
  resources:
    gpu: {minMemory: 88Gi}
```

Your training script does not change. If a Job does run out of memory anyway, it ends with the reason `OOMKilled`, and the fix the docs give is more `resources.memory` or a larger GPU.

Usually because one batch is longer than the rest. Activations grow with sequence length, and with eager attention the 5as/h term grows with its square. Run one step on your longest sequence at the start, so the worst case fails in the first minute instead of the fifth hour.

Both target fragmentation. `max_split_size_mb` stops the allocator splitting blocks above a size, which the [CUDA notes](https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-alloc-conf) say can let borderline workloads finish. We try expandable segments first because it is aimed at the changing tensor sizes that usually cause the slivers.

It does if the optimizer sees a smaller batch. Gradient accumulation keeps the effective batch the same, so the update matches the larger batch for most models. Layers that compute batch statistics, such as BatchNorm, still see the micro-batch.

The 16-bit weights take 16.1 GB. The KV cache adds 2 x 32 layers x 8 key-value heads x 128 x 2 bytes, or 0.13 MB per token, so 32 conversations of 8,192 tokens add 34.4 GB.

*Originally published on the [Nodus blog](https://www.nodus-compute.ai/blog/fix-cuda-out-of-memory-pytorch/).*