Fix Hugging Face CUDA Out of Memory When Loading Transformer Models

Fix Hugging Face CUDA Out of Memory When Loading Transformer Models

You call from_pretrained() and the process crashes before you even run a single inference:

Traceback (most recent call last):
  File "inference.py", line 6, in <module>
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf")
  File "/usr/local/lib/python3.11/site-packages/transformers/modeling_utils.py", line 3678, in from_pretrained
    dispatch_model(model, **device_map_kwargs)
  File "/usr/local/lib/python3.11/site-packages/accelerate/big_modeling.py", line 432, in dispatch_model
    load_checkpoint_and_dispatch(...)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.86 GiB
(GPU 0; 23.69 GiB total capacity; 21.05 GiB already allocated;
1.28 GiB free; 22.14 GiB reserved in total by PyTorch)

This is distinct from the training OOM covered in the torch.cuda.OutOfMemoryError training guide. At inference time there are no gradients, no optimizer states, and no activation memory from a backward pass. The problem is simpler: the model's weights alone exceed your GPU's VRAM. Here is how to calculate exactly how much memory you need, and four concrete techniques to load the model anyway.


Why Loading a Model Runs Out of Memory

When from_pretrained() loads a model, it reads the weights from disk and places them in memory. The amount of memory consumed at load time is approximately:

VRAM (bytes) ≈ num_parameters × bytes_per_parameter × overhead_factor

Where bytes_per_parameter depends on the dtype: float32 = 4 bytes, float16/bfloat16 = 2 bytes, int8 = 1 byte, int4 = 0.5 bytes. The overhead factor (typically 1.1–1.2) accounts for buffers and temporary tensors during the load process itself.

A 13B parameter model in float32 needs approximately 52 GB of VRAM — far more than even high-end consumer GPUs. In float16 it needs 26 GB. Even 7B models at float32 require 28 GB, which exceeds most single-GPU setups.

from transformers import AutoConfig

# Estimate memory BEFORE loading the model
def estimate_model_vram(model_name_or_path, dtype_bytes=4):
    """
    Estimate VRAM needed to load a model.
    dtype_bytes: 4 for float32, 2 for float16/bfloat16, 1 for int8
    """
    config = AutoConfig.from_pretrained(model_name_or_path)

    # Most models expose num_parameters via config arithmetic
    # This is an approximation — actual count varies by architecture
    hidden = getattr(config, 'hidden_size', 768)
    layers = getattr(config, 'num_hidden_layers', 12)
    heads  = getattr(config, 'num_attention_heads', 12)
    vocab  = getattr(config, 'vocab_size', 50257)
    intermediate = getattr(config, 'intermediate_size', hidden * 4)

    # Rough parameter count for a transformer block
    attn_params = 4 * hidden * hidden          # Q, K, V, O projections
    ffn_params  = 2 * hidden * intermediate    # up + down projections
    per_layer   = attn_params + ffn_params
    total_params = layers * per_layer + vocab * hidden  # + embeddings

    vram_bytes   = total_params * dtype_bytes * 1.2    # 20% overhead
    vram_gb      = vram_bytes / (1024 ** 3)
    return total_params, vram_gb

params, vram = estimate_model_vram("bert-base-uncased", dtype_bytes=4)
print(f"Estimated parameters: {params / 1e6:.0f}M")
print(f"Estimated VRAM (float32): {vram:.1f} GB")
import torch
from transformers import AutoModelForCausalLM

# After loading, get the exact parameter count
model = AutoModelForCausalLM.from_pretrained(
    "gpt2",   # 117M params, safe to load on CPU for demonstration
    torch_dtype=torch.float32
)

total_params = sum(p.numel() for p in model.parameters())
print(f"Exact parameters: {total_params / 1e6:.1f}M")

for dtype_bytes, dtype_name in [(4, "float32"), (2, "float16"), (1, "int8")]:
    vram_gb = (total_params * dtype_bytes) / (1024 ** 3)
    print(f"  {dtype_name}: ~{vram_gb:.2f} GB")

# Exact parameters: 124.4M
#   float32: ~0.46 GB
#   float16: ~0.23 GB
#   int8:    ~0.12 GB

Fix 1: torch_dtype=torch.float16

The single most impactful change: load the model in half precision instead of float32. This halves the VRAM requirement immediately, with negligible quality loss for inference on most tasks.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Llama-2-7b-hf"

# Before — loads in float32 by default (~28 GB for 7B model)
# model = AutoModelForCausalLM.from_pretrained(model_name)

# Fix — loads in float16 (~14 GB for 7B model)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="cuda:0"     # place directly on GPU after loading
)

tokenizer = AutoTokenizer.from_pretrained(model_name)

# Verify dtype and device
print(model.dtype)          # torch.float16
print(next(model.parameters()).device)  # cuda:0

# Inference is unchanged
inputs = tokenizer("The quick brown fox", return_tensors="pt").to("cuda:0")
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

On Ampere-class GPUs (RTX 3000+, A100) use torch.bfloat16 instead — it has the same memory footprint as float16 but better numerical range, which reduces the chance of NaN outputs on very deep models:

import torch
from transformers import AutoModelForCausalLM

# bfloat16: same 2 bytes/param, wider dynamic range than float16
# Preferred on: A100, A10, RTX 3090/4090, H100
# NOT supported on: older GPUs (V100, T4, RTX 2000 series)
model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1",
    torch_dtype=torch.bfloat16,
    device_map="cuda:0"
)

Fix 2: device_map="auto" with accelerate

When the model does not fit in GPU VRAM even at float16, device_map="auto" from the accelerate library distributes the model's layers across all available hardware — GPU VRAM first, then CPU RAM, then disk — automatically. You get to run models that are larger than your GPU's VRAM at the cost of latency on the offloaded layers.

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# pip install accelerate  (required for device_map to work)

model_name = "meta-llama/Llama-2-13b-hf"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"        # accelerate decides where each layer lives
)

tokenizer = AutoTokenizer.from_pretrained(model_name)

# Inspect how layers were distributed
print(model.hf_device_map)
# {'model.embed_tokens': 0, 'model.layers.0': 0, ..., 'model.layers.18': 0,
#  'model.layers.19': 'cpu', ..., 'lm_head': 'cpu'}
# Layers that do not fit in GPU VRAM are placed on CPU RAM

To see a breakdown of what would fit before loading:

from accelerate import init_empty_weights, infer_auto_device_map
from transformers import AutoConfig, AutoModelForCausalLM
import torch

model_name = "meta-llama/Llama-2-13b-hf"
config = AutoConfig.from_pretrained(model_name)

# Create a model skeleton with no actual weights (no memory allocated)
with init_empty_weights():
    empty_model = AutoModelForCausalLM.from_config(config)

# Ask accelerate how it would distribute given available hardware
device_map = infer_auto_device_map(
    empty_model,
    max_memory={
        0: "22GiB",    # GPU VRAM limit (leave some headroom)
        "cpu": "48GiB" # CPU RAM limit
    },
    dtype=torch.float16
)
print(device_map)
# Shows exactly which layer goes where before committing to loading

Fix 3: load_in_8bit or load_in_4bit with bitsandbytes

Quantization compresses model weights to lower bit widths at the cost of a small accuracy reduction. The bitsandbytes library integrates directly with Hugging Face from_pretrained() and handles quantization transparently during inference.

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

# pip install bitsandbytes accelerate

model_name = "meta-llama/Llama-2-13b-hf"

# --- 8-bit quantization ---
# 13B model: float16 ~26 GB → int8 ~13 GB
bnb_config_8bit = BitsAndBytesConfig(load_in_8bit=True)

model_8bit = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config_8bit,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

inputs = tokenizer("Explain gradient descent in one paragraph:", return_tensors="pt").to("cuda")
with torch.no_grad():
    outputs = model_8bit.generate(**inputs, max_new_tokens=150)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

model_name = "meta-llama/Llama-2-13b-hf"

# --- 4-bit quantization (NF4) ---
# 13B model: float16 ~26 GB → int4 ~6.5 GB
# Most aggressive reduction — slight quality loss, dramatically lower VRAM
bnb_config_4bit = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",         # NormalFloat4 — best quality for LLMs
    bnb_4bit_compute_dtype=torch.float16,  # compute in fp16 for speed
    bnb_4bit_use_double_quant=True     # double quantization saves ~0.4 bits/param extra
)

model_4bit = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config_4bit,
    device_map="auto"
)

# Verify actual memory used
import torch
allocated_gb = torch.cuda.memory_allocated() / (1024 ** 3)
print(f"GPU memory allocated: {allocated_gb:.1f} GB")
# ~6-7 GB for a 13B model in 4-bit

bitsandbytes quantization is compute-on-GPU — the weights are stored in int8/int4 but are dequantized on-the-fly during matrix multiplication. This means inference latency increases slightly compared to native float16, but VRAM usage drops proportionally to the bit width.


Fix 4: Use a Smaller Model Variant

If the task allows it, the most reliable fix is using a smaller model. For most production inference tasks, a well-tuned 7B model outperforms a poorly-loaded 70B model running at extreme quantization:

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Model family size comparison — same architecture, different capacity
model_options = {
    # VRAM estimates at float16
    "facebook/opt-125m":               "~0.25 GB",
    "facebook/opt-1.3b":               "~2.6 GB",
    "facebook/opt-6.7b":               "~13 GB",
    "facebook/opt-30b":                "~60 GB",

    "mistralai/Mistral-7B-v0.1":       "~14 GB",
    "mistralai/Mixtral-8x7B-v0.1":     "~93 GB",  # MoE: activates 2 of 8 experts per token

    "meta-llama/Llama-2-7b-hf":        "~14 GB",
    "meta-llama/Llama-2-13b-hf":       "~26 GB",
    "meta-llama/Llama-2-70b-hf":       "~140 GB",
}

for name, vram in model_options.items():
    print(f"{name:<45} {vram}")

# For text classification tasks — DistilBERT instead of BERT-large
from transformers import pipeline

# BERT-large: 340M params, ~1.3 GB at float32
# DistilBERT: 66M params, ~0.26 GB at float32 — 40% faster, 97% of BERT accuracy
classifier = pipeline(
    "sentiment-analysis",
    model="distilbert-base-uncased-finetuned-sst-2-english",
    device=0
)
result = classifier("This tutorial was very helpful.")
print(result)  # [{'label': 'POSITIVE', 'score': 0.9998}]

Estimating Model Memory Before Loading

Always estimate memory requirements before loading an unfamiliar model. This avoids the OOM crash entirely:

import torch
from transformers import AutoConfig

def vram_required(model_name, dtype=torch.float16):
    """
    Estimate VRAM needed for inference (weights only, no KV cache).
    More accurate than architecture-based estimates since we read
    the weight shapes directly from the model config's param count.
    """
    dtype_bytes = {
        torch.float32: 4,
        torch.float16: 2,
        torch.bfloat16: 2,
    }.get(dtype, 4)

    # Load config only — no weights, no memory
    config = AutoConfig.from_pretrained(model_name)

    # num_parameters is available on most model configs in newer transformers
    # Fall back to a reasonable estimate based on common fields
    if hasattr(config, 'num_parameters'):
        n_params = config.num_parameters
    else:
        h = getattr(config, 'hidden_size', 768)
        layers = getattr(config, 'num_hidden_layers', 12)
        vocab = getattr(config, 'vocab_size', 50000)
        ffn = getattr(config, 'intermediate_size', h * 4)
        n_params = layers * (4 * h * h + 2 * h * ffn) + vocab * h

    weights_gb = (n_params * dtype_bytes) / (1024 ** 3)
    # KV cache for inference adds roughly 5-15% on top of weights
    total_gb = weights_gb * 1.15

    print(f"Model: {model_name}")
    print(f"Estimated parameters: {n_params / 1e9:.2f}B")
    print(f"Weights at {dtype}: {weights_gb:.1f} GB")
    print(f"Total with KV cache estimate: {total_gb:.1f} GB")

    available = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3) if torch.cuda.is_available() else 0
    print(f"Available GPU VRAM: {available:.1f} GB")
    print(f"Fits in GPU: {'YES' if total_gb < available * 0.9 else 'NO — use float16, quantization, or device_map'}")
    return total_gb

vram_required("bert-base-uncased", dtype=torch.float32)
# Model: bert-base-uncased
# Estimated parameters: 0.11B
# Weights at torch.float32: 0.4 GB
# Total with KV cache estimate: 0.5 GB
# Available GPU VRAM: 24.0 GB
# Fits in GPU: YES

Quick Reference: Model Sizes and Minimum VRAM

Model Parameters float32 float16 int8 int4
DistilBERT-base 66M 0.3 GB 0.1 GB 0.07 GB 0.03 GB
BERT-base 110M 0.4 GB 0.2 GB 0.1 GB 0.06 GB
BERT-large 340M 1.3 GB 0.7 GB 0.3 GB 0.2 GB
GPT-2 117M 0.5 GB 0.2 GB 0.1 GB 0.06 GB
GPT-2 XL 1.6B 6.1 GB 3.0 GB 1.5 GB 0.8 GB
Mistral 7B 7.3B 27 GB 14 GB 7 GB 3.5 GB
LLaMA-2 7B 7B 28 GB 14 GB 7 GB 3.5 GB
LLaMA-2 13B 13B 52 GB 26 GB 13 GB 6.5 GB
LLaMA-2 70B 70B 280 GB 140 GB 70 GB 35 GB
Mixtral 8x7B 47B 188 GB 94 GB 47 GB 24 GB

Estimates above are weights-only. Add 10–20% for KV cache during inference. GPU minimum is the amount needed with no other allocations; in practice leave 2+ GB headroom.


Summary

  • OOM at from_pretrained() is caused by model weights alone exceeding VRAM — no gradients involved.
  • Easiest fix: torch_dtype=torch.float16 — halves memory instantly, minimal quality loss.
  • Doesn't fit even at float16: add device_map="auto" with accelerate — offloads layers to CPU RAM or disk.
  • Need it all on GPU: use BitsAndBytesConfig(load_in_8bit=True) or load_in_4bit=True — 4x/8x weight compression with small accuracy cost.
  • Best long-term: choose a model that fits comfortably at float16 — a 7B model running cleanly outperforms a 70B model thrashing on disk offload.
  • Before loading: estimate VRAM with AutoConfig.from_pretrained() and parameter counting — avoid the crash before it happens.

Related articles