GPU Out of Memory During Inference: Why It's Different From Training OOM and How to Fix It

TL;DR

A GPU out of memory error during inference is a different problem from the same error during training. Training OOM is usually caused by activations, gradients, and optimizer states. Inference OOM is caused by model weight loading, the KV cache growing during generation, batch size at inference time, and model precision. The fixes are completely different. This post explains what is consuming your GPU memory during inference, how to diagnose which one pushed you over the limit, and the specific fixes for each cause.

Your training run finished cleanly. The model downloaded. You load it up for inference and get this:

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 23.69 GiB total capacity; 21.23 GiB already allocated)

Your first instinct might be to apply the same fixes you use for training OOM. Cut the batch size. Enable gradient checkpointing.

But gradient checkpointing does nothing during inference. And the batch size at inference is probably already 1.

Inference OOM has its own causes. Here is what they are and how to fix them.

What Is Actually in GPU Memory During Inference

During inference your GPU is holding several things at once, and they behave differently than during training.

Model weights are the largest fixed cost. A 7B model in 16-bit precision takes roughly 14 GB. In 8-bit it takes about 7 GB. In 4-bit it takes about 4 GB. These numbers do not change once the model is loaded.

The KV cache is where inference gets tricky. Every token the model generates requires storing key and value matrices for all previous tokens in the context. This is called the KV cache. It grows with every token generated and with every position in your input prompt.

For a long prompt or a long generation, the KV cache can consume several gigabytes on top of the model weights. A 7B model with a 4096 token context at 16-bit precision can have a KV cache that adds 2 to 4 GB to your memory footprint mid-generation.

Activations during the forward pass are smaller during inference than training because you are not computing gradients. But they still exist and they scale with context length.

Multiple concurrent requests each have their own KV cache. If you are serving the model and batching requests, the memory adds up fast.

Cause One: Model Precision Is Too High for Your GPU

This is the most common cause of inference OOM and the easiest to fix.

If you loaded your model in full 16-bit or 32-bit precision and your GPU does not have enough VRAM for the full model, you hit OOM before the first token generates.

Diagnose it:

nvidia-smi

If your GPU is already near capacity right after model loading, before you generate anything, this is your problem.

Fix: load the model in lower precision

With Hugging Face Transformers:

from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained(
    "your-model-path",
    torch_dtype=torch.float16,  # or torch.bfloat16
    device_map="auto"
)

For even lower memory use 4-bit quantization with bitsandbytes:

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_4bit=True)

model = AutoModelForCausalLM.from_pretrained(
    "your-model-path",
    quantization_config=quantization_config,
    device_map="auto"
)

4-bit quantization cuts memory use by roughly 75 percent compared to 16-bit. For inference on a fine-tuned model, the quality difference is small for most tasks.

Cause Two: The KV Cache Is Filling Your GPU Mid-Generation

This one is harder to spot because the model loads fine but crashes partway through generating a long response.

The KV cache grows linearly with your context length. If you are generating a 2000 token response to a 1000 token prompt, you are accumulating 3000 tokens worth of KV cache by the end. On a 7B model at 16-bit precision, that can be 3 to 5 GB on top of your base model memory.

Diagnose it:

Watch your GPU memory during inference:

watch -n 0.5 nvidia-smi

If memory climbs steadily during generation and then crashes, the KV cache is the cause.

Fix option one: reduce max new tokens

Cap how many tokens the model can generate per request:

outputs = model.generate(
    inputs,
    max_new_tokens=512  # lower this
)

Shorter generations mean a smaller maximum KV cache size.

Fix option two: reduce max context length

If your inputs are long, truncate them before passing to the model:

inputs = tokenizer(
    prompt,
    max_length=1024,
    truncation=True,
    return_tensors="pt"
)

Fix option three: use a more memory-efficient inference framework

vLLM uses PagedAttention to manage the KV cache much more efficiently than the default Transformers generate loop. For serving a model with variable-length inputs and outputs, switching to vLLM can dramatically reduce the peak KV cache memory footprint.

pip install vllm

from vllm import LLM

llm = LLM(model="your-model-path", gpu_memory_utilization=0.85)

The gpu_memory_utilization parameter tells vLLM what fraction of your GPU memory it can use for the KV cache. Lowering it from the default 0.90 gives you headroom if you are hitting OOM.

Cause Three: device_map Is Not Splitting the Model Correctly

If you are using device_map="auto" and have multiple GPUs, Transformers is supposed to split the model across them automatically. Sometimes it does not balance correctly and one GPU gets overloaded.

Diagnose it:

nvidia-smi

If one GPU is nearly full and another is nearly empty, the split is wrong.

Fix: specify a manual device map

device_map = {
    "model.embed_tokens": 0,
    "model.layers.0": 0,
    # ... first half of layers on GPU 0
    "model.layers.16": 1,
    # ... second half on GPU 1
    "lm_head": 1
}

model = AutoModelForCausalLM.from_pretrained(
    "your-model-path",
    device_map=device_map
)

This gives you explicit control over which layers land on which GPU.

Cause Four: Memory Fragmentation From Previous Runs

If you loaded and unloaded a model in the same Python session and then tried to load again, PyTorch may have fragmented GPU memory that it cannot reclaim cleanly.

Fix:

import torch
import gc

del model
gc.collect()
torch.cuda.empty_cache()

Run this before reloading. If you are in a notebook, restart the kernel entirely. Memory fragmentation in notebooks is common and the clean fix is a kernel restart.

The Inference OOM Checklist

Work through these in order before reaching for a bigger GPU.

One: check your model precision. Are you loading in 32-bit when 16-bit or 4-bit would work?

Two: watch memory during generation. Is it climbing to OOM or hitting it immediately on load?

Three: reduce max new tokens and max input length if the KV cache is growing too large.

Four: check GPU balance if you have multiple GPUs. Is one overloaded?

Five: clear memory between loads if you are in a notebook or iterating in the same session.

If you have worked through all five and you are still hitting OOM, your model genuinely requires more VRAM than your GPU has. The options at that point are a GPU with more memory, a smaller model, or more aggressive quantization.

The Key Difference From Training OOM

Training OOM is about activations, gradients, and optimizer states. The fixes involve gradient checkpointing, batch size, and LoRA rank.

Inference OOM is about model weight precision and the KV cache. The fixes involve quantization, context length limits, and more efficient attention implementations.

Knowing which problem you have saves you from applying the wrong fix and wondering why nothing works.