How to fix OOM crashes when running large open-source LLMs locally

Engineering Stable Local LLM Inference: Mastering VRAM Budgeting Beyond Weight Sizing

Current Situation Analysis

The most pervasive failure mode in local large language model (LLM) deployment is the "Phantom Memory" crash. Developers routinely encounter a scenario where a model's documented weight size appears to fit comfortably within available GPU VRAM, yet the inference process terminates abruptly with an Out of Memory (OOM) error the moment a prompt is processed. This is not a bug in the model or the hardware; it is a fundamental misunderstanding of VRAM composition during inference.

The industry pain point stems from marketing and documentation that focus exclusively on weight storage. A 13-billion parameter model in FP16 requires approximately 26 GB for weights alone. On a 24 GB GPU, this looks impossible, but on a 48 GB GPU, it looks safe. However, weight storage is only one component of the VRAM equation. The actual memory footprint at runtime is the sum of three distinct allocations:

1. Model Weights: Static parameters loaded into VRAM.
2. KV Cache: Dynamic memory storing Key and Value tensors for every token in the context window. This scales linearly with sequence length and batch size.
3. Activation Memory: Intermediate tensors generated during the forward pass, required for gradient computation (if fine-tuning) or intermediate layer outputs.

The KV cache is the primary culprit for "fits but crashes" failures. For a transformer architecture, the KV cache memory requirement is calculated as:

KV_Cache_Bytes = 2 × Layers × Batch_Size × Seq_Length × Heads × Head_Dim × Bytes_Per_Element

Consider a 13B model with 40 layers, 40 attention heads, and a head dimension of 128. In FP16 (2 bytes per element), a single sequence at 4,096 context length consumes roughly 2.5 GB for the KV cache. Extending this to 32,768 tokens balloons the cache to approximately 20 GB. When combined with 26 GB of weights, the total requirement hits 46 GB, causing an immediate OOM on hardware with less capacity. This dynamic scaling is often overlooked during initial capacity planning, leading to production instability when context lengths increase.

WOW Moment: Key Findings

The critical insight is that weight quantization alone often fails to resolve OOM issues at long context lengths because the KV cache remains uncompressed. The most effective strategy involves a dual approach: quantizing weights to reduce the static footprint and quantizing the KV cache to control dynamic scaling.

The table below illustrates the VRAM impact of different optimization strategies for a 13B parameter model at 32K context length on a single GPU.

Optimization Strategy	Estimated VRAM Usage	Perplexity Impact	Inference Latency
FP16 Baseline	~46 GB	Baseline	Fast
NF4 Weights Only	~27 GB	Near Lossless	Fast
NF4 Weights + FP8 KV	~12 GB	Near Lossless	Fast
NF4 Weights + Q4 KV	~8 GB	Degraded Recall	Fast
NF4 Weights + CPU Offload	~7 GB (GPU)	Lossless	5-10x Slower

Why this matters: The data reveals that FP8 KV cache quantization is the "sweet spot" for most production workloads. It reduces the dominant dynamic memory component by 2-4x with negligible quality loss, allowing models to run at long context lengths on consumer hardware. Conversely, dropping KV cache to 4-bit (Q4) saves memory but introduces measurable degradation in tasks requiring precise token recall, making it unsuitable for code generation or complex reasoning without careful evaluation.

Core Solution

Resolving VRAM constraints requires a systematic audit and optimization pipeline. The following steps provide a reproducible methodology for stabilizing local LLM inference.

Step 1: Audit VRAM Allocation with Memory Snapshots

Before applying optimizations, you must identify which component is consuming VRAM. PyTorch provides a built-in memory profiler that captures allocation events. This is superior to monitoring tools like nvidia-smi because it attributes memory to specific tensors and operations.

Implement a snapshot capture routine to dump memory state after a representative forward pass.

import torch
import pickle
from transformers import AutoTokenizer, AutoModelForCausalLM

def capture_vram_snapshot(model_id: str, prompt: str, output_path: str):
    """
    Captures a detailed VRAM allocation snapshot for debugging OOM issues.
    """
    # Enable memory history recording
    torch.cuda.memory._record_memory_history(max_entries=100000)
    
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForCausalLM.from_pretrained(
            model_id, 
            torch_dtype=torch.float16, 
            device_map="auto"
        )
        
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        # Execute a generation pass to populate KV cache
        with torch.no_grad():
            model.generate(
                **inputs, 
                max_new_tokens=128, 
                do_sample=False
            )
            
        # Dump snapshot for analysis
        torch.cuda.memory._dump_snapshot(output_path)
        print(f"Snapshot saved to {output_path}")
        
    finally:
        # Always disable recording to avoid performance overhead
        torch.cuda.memory._record_memory_history(enabled=None)

# Usage
# capture_vram_snapshot("meta-llama/Llama-2-13b-hf", "Explain quantum computing.", "vram_debug.pickle")

Analyze the resulting pickle file using the PyTorch Memory Visualizer. Look for large allocations attributed to self_attn or k_cache/v_cache tensors. If these dominate the profile, weight quantization will not solve the problem; KV cache optimization is required.

Step 2: Implement Targeted Quantization

Quantization must be applied based on the bottleneck identified in Step 1.

Weight Quantization: For weight-dominated scenarios, use 4-bit NormalFloat (NF4) quantization via bitsandbytes. NF4 is statistically optimized for normally distributed weights and outperforms standard FP4 quantization in perplexity benchmarks. Enable double quantization to compress the quantization constants themselves, saving additional memory.

from transformers import BitsAndBytesConfig, AutoModelForCausalLM
import torch

def load_quantized_weights(model_id: str):
    """
    Loads a model with NF4 weight quantization and double quantization.
    """
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True
    )
    
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        quantization_config=quantization_config,
        device_map="auto"
    )
    return model

KV Cache Quantization: If the KV cache is the bottleneck, quantize the cache rather than the weights. This is typically handled by the inference engine rather than the model loader.

• vLLM: Use the --kv-cache-dtype fp8 flag. This maintains near-lossless quality while reducing cache memory by 2x.
• llama.cpp: Use --cache-type-k q4_0 --cache-type-v q4_0 for aggressive compression, or --cache-type-k q8_0 --cache-type-v q8_0 for a balance of compression and quality.

Step 3: Resolve Memory Fragmentation

PyTorch's default CUDA allocator can suffer from fragmentation after repeated allocation and deallocation cycles, particularly during long generation sequences. This results in OOM errors even when nvidia-smi reports free memory, because the free memory is split into small, non-contiguous blocks.

Enable the expandable segments allocator via an environment variable. This allows the allocator to grow memory blocks dynamically, significantly reducing fragmentation.

# Set in shell profile or startup script
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

This setting is safe for production and has negligible overhead. It resolves a class of OOM errors that are otherwise difficult to diagnose.

Step 4: Strategic Offloading

When VRAM is insufficient even after quantization, offload layers to CPU memory using the accelerate library. This allows models to run on GPU-constrained hardware at the cost of latency due to PCIe bandwidth limitations.

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

def setup_offloaded_model(model_id: str, gpu_limit_gb: int, cpu_limit_gb: int):
    """
    Configures a device map with explicit GPU headroom and CPU offloading.
    """
    config = AutoConfig.from_pretrained(model_id)
    
    with init_empty_weights():
        model = AutoModelForCausalLM.from_config(config)
        
    # Reserve GPU memory for KV cache and activations
    # Example: Leave 6GB free on GPU 0 for runtime overhead
    device_map = infer_auto_device_map(
        model,
        max_memory={
            0: f"{gpu_limit_gb}GiB",
            "cpu": f"{cpu_limit_gb}GiB"
        },
        no_split_module_classes=["LlamaDecoderLayer"]
    )
    
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map=device_map,
        torch_dtype=torch.float16
    )
    return model

Rationale: Always reserve headroom on the GPU. If you allocate all GPU memory to weights, there is no space for the KV cache, and the model will crash. The infer_auto_device_map function distributes layers intelligently, but manual limits ensure the runtime has sufficient buffer.

Pitfall Guide

1. The Smoke Test Trap:

   • Explanation: Testing with a single short prompt passes, but the model crashes under production load. Smoke tests often fail to populate the KV cache sufficiently to trigger OOM.
   • Fix: Stress test with maximum expected context length and batch size. Use tools like vLLM's benchmarking suite to simulate concurrent requests.

2. KV Cache Blindness:

   • Explanation: Assuming VRAM usage is constant regardless of sequence length.
   • Fix: Calculate KV cache size using the formula provided in the Current Situation Analysis. Plan for the maximum context length, not the average.

3. Over-Quantizing the KV Cache:

   • Explanation: Using 4-bit KV cache quantization to save memory, resulting in degraded performance on tasks requiring precise recall, such as code completion or multi-step reasoning.
   • Fix: Use FP8 KV cache for production. Reserve Q4 KV cache only for experimental runs where memory is critically constrained and quality degradation is acceptable.

4. Fragmentation Neglect:

   • Explanation: Ignoring memory fragmentation leads to intermittent OOMs that are hard to reproduce.
   • Fix: Always set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. This is a best practice for any PyTorch inference workload.

5. CPU Offload Latency Shock:

   • Explanation: Expecting GPU-like speeds when layers are offloaded to CPU. PCIe transfer bottlenecks can cause 5-10x slowdowns.
   • Fix: Use offloading only when necessary. If latency is critical, consider model parallelism across multiple GPUs or upgrading hardware rather than offloading.

6. Batch Size Oversight:

   • Explanation: Optimizing for single-sequence inference but deploying with a batch size > 1. KV cache scales linearly with batch size.
   • Fix: Include batch size in all VRAM calculations. Test with the actual batch size used in production.

7. Double Quantization Omission:

   • Explanation: Enabling 4-bit quantization but forgetting double quantization, wasting memory on quantization constants.
   • Fix: Always enable bnb_4bit_use_double_quant=True when using bitsandbytes 4-bit quantization.

Production Bundle

Action Checklist

• Calculate KV Budget: Compute KV cache memory for max context length and batch size before deployment.
• Enable Expandable Segments: Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True in the environment.
• Profile Baseline: Run a memory snapshot on the unoptimized model to identify bottlenecks.
• Apply Weight Quantization: Use NF4 with double quantization if weights dominate VRAM usage.
• Apply KV Quantization: Use FP8 KV cache if context length dominates VRAM usage.
• Reserve GPU Headroom: Ensure device maps leave sufficient memory for runtime allocations.
• Stress Test: Validate stability with max context and batch size using monitoring tools.
• Select Serving Engine: Use vLLM or llama.cpp for production to leverage PagedAttention and optimized memory management.

Decision Matrix

Scenario	Recommended Approach	Why	Cost Impact
Short Context, High Quality	FP16/BF16 Weights, No KV Quant	Maximum quality; VRAM sufficient for weights and short cache.	Low (Standard GPU)
Long Context, Limited VRAM	NF4 Weights + FP8 KV Cache	Balances memory reduction with near-lossless quality.	Medium (Optimized GPU usage)
Max Context, Budget Hardware	NF4 Weights + Q4 KV Cache	Aggressive memory reduction; acceptable quality trade-off for extreme lengths.	Low (Consumer GPU viable)
Model Too Large for GPU	NF4 Weights + CPU Offload	Enables running models exceeding VRAM capacity; latency penalty accepted.	Low (Hardware cost) / High (Latency)
High Throughput Serving	vLLM with PagedAttention	Efficient KV cache management; automatic batching; production-ready.	Medium (Engineering effort)

Configuration Template

Use this template to configure a stable inference environment with quantization and fragmentation fixes.

# inference_config.py
import os
import torch
from transformers import BitsAndBytesConfig, AutoModelForCausalLM, AutoTokenizer

# Environment setup for fragmentation prevention
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

class LLMInferenceConfig:
    def __init__(self, model_id: str, quantize_weights: bool = True, kv_cache_dtype: str = "fp16"):
        self.model_id = model_id
        self.quantize_weights = quantize_weights
        self.kv_cache_dtype = kv_cache_dtype
        
    def get_quantization_config(self) -> BitsAndBytesConfig:
        if not self.quantize_weights:
            return None
            
        return BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True
        )
        
    def load_model(self):
        config = self.get_quantization_config()
        
        model = AutoModelForCausalLM.from_pretrained(
            self.model_id,
            quantization_config=config,
            torch_dtype=torch.bfloat16 if config is None else None,
            device_map="auto"
        )
        
        tokenizer = AutoTokenizer.from_pretrained(self.model_id)
        return model, tokenizer

# Example usage
# config = LLMInferenceConfig("meta-llama/Llama-2-13b-hf", quantize_weights=True)
# model, tokenizer = config.load_model()

Quick Start Guide

1. Install Dependencies: Ensure torch, transformers, bitsandbytes, and accelerate are installed.

   pip install torch transformers bitsandbytes accelerate
   

2. Set Environment Variable: Add export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to your shell profile.
3. Run Profiler: Execute the snapshot script from Step 1 to analyze current memory usage.
4. Apply Quantization: Update your model loading code to use NF4 weight quantization and FP8 KV cache if supported by your inference engine.
5. Verify Stability: Run a stress test with maximum context length and monitor VRAM usage. Confirm no OOM errors occur.