Troubleshooting CUDA Out of Memory During Model Loading

One of the most frustrating errors a machine learning engineer can encounter is the dreaded RuntimeError: CUDA out of memory. While this error can occur at any point during training or inference, it is particularly insidious when it happens during model loading — before you've even begun computing anything. This tutorial walks you through the root causes, diagnostic techniques, and concrete fixes for CUDA Out of Memory (OOM) errors that surface while loading models into GPU memory.

What Is CUDA Out of Memory During Model Loading?

CUDA Out of Memory is an error raised by PyTorch (and other CUDA-backed frameworks) when the GPU's available VRAM is insufficient to satisfy an allocation request. When this happens during model loading, it typically means the framework is attempting to materialize model weights, intermediate buffers, or optimizer state onto the GPU, and the requested block of memory exceeds what is currently free.

A typical traceback looks like this:

RuntimeError: CUDA out of memory. Tried to allocate 2.00 GiB.
GPU 0 has a total capacity of 15.78 GiB of which 1.42 GiB is free.
Process 12345 has 14.36 GiB memory in use. Of the allocated memory
13.80 GiB is allocated by PyTorch, and 432.00 MiB is reserved by
PyTorch but unallocated.

The key insight is that the error is not always about the model itself being too large — it is about the gap between what the GPU has free and what the load operation requests. That gap can be caused by leftover allocations, fragmentation, framework overhead, or simply a model that genuinely exceeds the hardware.

Why It Matters

Model loading is the gateway to every downstream workflow: fine-tuning, inference, evaluation, and deployment. If you cannot load the model, nothing else proceeds. Worse, OOM during loading often surfaces in production scenarios — when you are trying to serve a larger model on constrained hardware, or when you are migrating from a development GPU to a smaller inference GPU. Understanding how to diagnose and resolve these errors is therefore essential for:

Shipping large language models and diffusion models on cost-effective hardware.
Avoiding silent failures in CI/CD pipelines that load checkpoints.
Maximizing utilization of multi-GPU setups without over-provisioning.
Debugging shared environments where multiple processes compete for VRAM.

Common Causes

Before jumping to fixes, it helps to understand the typical culprits. Most loading-time OOM errors fall into one of the following categories:

Leftover allocations: A previous process or notebook kernel still holds GPU memory.
Loading to GPU by default: The checkpoint or loader pushes tensors to cuda before you have a chance to control placement.
Full-precision weights: Loading FP32 weights when FP16 or BF16 would suffice doubles or quadruples memory usage.
Fragmentation: The GPU has enough total free memory, but it is split into small non-contiguous blocks.
Optimizer and gradient state: Loading a training checkpoint that includes optimizer state can triple effective memory usage.
Multiple models on one GPU: Loading a second model without releasing the first.

Diagnosing the Problem

The first step is always to measure. Guessing at OOM causes wastes time. Use the following commands and snippets to inspect the actual state of your GPU and process.

From a shell, check overall GPU usage:

nvidia-smi

For a continuously updating view:

watch -n 1 nvidia-smi

Look for processes holding memory that you did not expect. Often a stale Jupyter kernel or a crashed training job will linger. To find Python processes consuming GPU memory:

fuser -v /dev/nvidia*

Inside your Python script, inspect PyTorch's view of memory:

import torch

print("Allocated:", torch.cuda.memory_allocated() / 1e9, "GiB")
print("Reserved: ", torch.cuda.memory_reserved() / 1e9, "GiB")
print("Max alloc:", torch.cuda.max_memory_allocated() / 1e9, "GiB")

If memory_allocated() is high before you load anything, you have leftover tensors. If memory_reserved() is high but memory_allocated() is low, you have fragmentation or a caching issue.

Fix 1: Clear Leftover State

The simplest and most common fix is to release stale memory before loading. In a fresh script this is unnecessary, but in notebooks and long-running processes it is critical.

import gc
import torch

# Force garbage collection of Python objects referencing tensors
gc.collect()

# Empty PyTorch's caching allocator
torch.cuda.empty_cache()

# Verify memory is actually free
print("Free after cleanup:", torch.cuda.mem_get_info()[0] / 1e9, "GiB")

Note that empty_cache() only returns cached memory to the CUDA driver — it does not free tensors that are still referenced by live Python objects. You must ensure those references are gone first, which is why gc.collect() precedes it.

Fix 2: Load on CPU First

Many OOM errors occur because the loader defaults to device="cuda" or because torch.load restores tensors to their original device. By loading to CPU first, you gain control over how and when memory moves to the GPU.

import torch

# Load checkpoint with explicit CPU mapping
checkpoint = torch.load("model.pt", map_location="cpu", weights_only=True)

# Build the model on CPU
model = MyModel()
model.load_state_dict(checkpoint["model_state_dict"])

# Move to GPU explicitly, only when ready
device = torch.device("cuda")
model = model.to(device)

The map_location="cpu" argument is the key. Without it, PyTorch attempts to restore each tensor to the device it was saved on, which may be a GPU that no longer exists or is already full.

Fix 3: Use Lower Precision

If the model genuinely approaches the GPU's capacity, reducing precision is the most effective lever. A 7-billion-parameter model requires roughly:

28 GiB in FP32
14 GiB in FP16 or BF16
7 GiB in INT8
~3.5 GiB in INT4

For Hugging Face Transformers, load directly in half precision:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    torch_dtype=torch.float16,
    device_map="auto",
)

For even greater savings, use 8-bit or 4-bit quantization via the bitsandbytes library:

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=bnb_config,
    device_map="auto",
)

Fix 4: Sharded and Offloaded Loading

For models that exceed a single GPU's VRAM even in half precision, use sharded loading. Hugging Face Accelerate provides device_map="auto", which splits the model across available GPUs and offloads overflow to CPU RAM or disk.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "bigscience/bloom-176b",
    device_map="auto",
    offload_folder="offload",
    offload_state_dict=True,
)

The offload_folder directory must exist and have enough disk space to hold the offloaded weights. The offload_state_dict=True option is particularly useful during loading because it streams state dict tensors to disk rather than holding them all in CPU RAM simultaneously.

You can also inspect the resulting device map to verify placement:

print(model.hf_device_map)

Fix 5: Avoid Loading Optimizer State

Training checkpoints often include optimizer state, scheduler state, and RNG state. For inference or for resuming on smaller hardware, loading all of this is wasteful and can trigger OOM. Load only what you need:

checkpoint = torch.load("checkpoint.pt", map_location="cpu", weights_only=True)

# Load only model weights, discard optimizer and scheduler
model.load_state_dict(checkpoint["model_state_dict"])

# If you must resume training, load optimizer separately
# and only after the model is already on the GPU
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

For Hugging Face models saved with save_pretrained, optimizer state is stored separately in optimizer.pt and scheduler.pt. Simply do not load those files when you only need the model.

Fix 6: Manage Fragmentation

Sometimes nvidia-smi shows enough free memory, yet PyTorch still raises OOM. This is usually fragmentation: the free memory exists, but not in a single contiguous block large enough for the requested allocation. PyTorch provides an environment variable to control the caching allocator's behavior:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

This setting allows the allocator to request expandable segments from the CUDA driver, reducing fragmentation at the cost of slightly higher overhead. It is particularly helpful when loading models in stages or when interleaving model loading with other GPU work.

You can also force a full reset of the allocator within a script:

torch.cuda.empty_cache()
# If still fragmented, the nuclear option:
torch.cuda.reset_peak_memory_stats()

Fix 7: Use Memory-Efficient Loading Libraries

For the largest models, dedicated libraries handle loading more efficiently than naive torch.load. safetensors is a popular format that supports zero-copy memory mapping, avoiding the need to load the entire checkpoint into CPU RAM before transferring to GPU:

from safetensors.torch import load_file

state_dict = load_file("model.safetensors")
model.load_state_dict(state_dict)

For extremely large models, consider ggml or llama.cpp-based loaders, which can mmap weights directly from disk and load only the layers currently needed for inference.

Best Practices

Always load to CPU first. Make GPU placement an explicit, deliberate step rather than a side effect of the loader.
Use the smallest viable precision. Default to BF16 or FP16 for loading; quantize further only when necessary.
Profile before optimizing. Use nvidia-smi and torch.cuda.memory_summary() to confirm the actual bottleneck before applying fixes.
Isolate environments. Run model loading in a fresh process or container to avoid contamination from prior allocations.
Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True in production environments prone to fragmentation.
Prefer safetensors over pickle-based checkpoints for both security and memory efficiency.
Use device_map="auto" with offloading for models that approach or exceed single-GPU capacity.
Free models explicitly. Call del model followed by gc.collect() and torch.cuda.empty_cache() before loading a replacement.
Test loading on target hardware early. Do not discover OOM at deploy time; validate the load path on the inference GPU during development.

Putting It All Together

The following snippet combines several of the techniques above into a robust loading routine suitable for large models on constrained hardware:

import gc
import os

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

# Step 1: Clear any stale state
gc.collect()
torch.cuda.empty_cache()

# Step 2: Configure quantization
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
)

# Step 3: Load with automatic device placement and offloading
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-13b-hf",
    quantization_config=quant_config,
    device_map="auto",
    offload_folder="offload",
    offload_state_dict=True,
    torch_dtype=torch.float16,
)

# Step 4: Verify placement and memory
print("Device map:", model.hf_device_map)
print("Allocated:", torch.cuda.memory_allocated() / 1e9, "GiB")
print("Reserved: ", torch.cuda.memory_reserved() / 1e9, "GiB")

Conclusion

CUDA Out of Memory during model loading is rarely a single, monolithic problem — it is usually a combination of leftover allocations, default device placement, excessive precision, and fragmentation. By following a disciplined diagnostic workflow that starts with measurement, proceeds through explicit CPU-first loading, and leverages precision reduction, quantization, sharding, and offloading as needed, you can reliably load models that initially seemed too large for your hardware. The key is to treat memory as a first-class engineering concern: measure it, control it explicitly, and never assume the framework will do the right thing by default. With these techniques in your toolkit, OOM during loading becomes a solvable puzzle rather than a project-blocking wall.