The CUDA out of memory error means your GPU ran out of VRAM while trying to perform a PyTorch operation. This usually happens because the model, the batch size, or the intermediate activations generated during the forward and backward passes collectively exceed the available memory on your GPU.

Here are the most common reasons and their fixes:

1. Large Batch Size:

  • Diagnosis: Monitor GPU memory usage with nvidia-smi. If memory spikes significantly when you start training and hits the limit, your batch size is likely too high.
  • Fix: Reduce your batch_size in your DataLoader. For example, if you’re using batch_size=64, try batch_size=32 or batch_size=16.
  • Why it works: A smaller batch size means fewer samples are processed simultaneously, leading to smaller tensors and thus less memory consumption for activations and gradients.

2. Large Model Parameters or Complex Architecture:

  • Diagnosis: Even with a small batch size, a very large model (e.g., many layers, wide layers, or a transformer with many attention heads) can consume too much memory. Check the total number of parameters in your model.
  • Fix:
    • Model Quantization: Use techniques like 8-bit or 4-bit quantization. For example, with bitsandbytes:
      from transformers import AutoModelForCausalLM, AutoTokenizer
      from transformers import BitsAndBytesConfig
      
      quantization_config = BitsAndBytesConfig(
          load_in_8bit=True,
          bnb_8bit_compute_dtype=torch.float16
      )
      model = AutoModelForCausalLM.from_pretrained(
          "your_model_name",
          quantization_config=quantization_config,
          device_map="auto"
      )
      
    • Gradient Checkpointing: This trades computation for memory by recomputing activations during the backward pass instead of storing them. In PyTorch:
      model.gradient_checkpointing_enable()
      
    • Model Pruning/Distillation: More advanced techniques to reduce model size.
  • Why it works: Quantization reduces the precision of model weights and activations, requiring less VRAM. Gradient checkpointing avoids storing all intermediate activations by recalculating them as needed.

3. Accumulating Gradients (Training in Chunks):

  • Diagnosis: You might be intentionally accumulating gradients over several mini-batches to simulate a larger effective batch size. If the accumulation steps are too high, the gradients from all steps can build up.
  • Fix: Reduce the number of gradient_accumulation_steps in your training loop or optimizer configuration. Ensure you call optimizer.zero_grad(set_to_none=True) at the correct frequency (usually once per effective batch, not every gradient_accumulation_steps).
    # Example training loop snippet
    optimizer.zero_grad(set_to_none=True) # Zero gradients before the first accumulation step
    for i, batch in enumerate(dataloader):
        outputs = model(batch)
        loss = outputs.loss
        loss = loss / gradient_accumulation_steps # Scale loss
        loss.backward() # Accumulates gradients
    
        if (i + 1) % gradient_accumulation_steps == 0:
            optimizer.step() # Update weights
            optimizer.zero_grad(set_to_none=True) # Zero gradients after update
    
  • Why it works: By zeroing gradients more frequently, you prevent the memory footprint of accumulated gradients from growing excessively large.

4. Memory Leaks (Unreleased Tensors):

  • Diagnosis: This is subtle. If your OOM error happens after some training steps, not immediately, it might be a leak. PyTorch’s torch.cuda.empty_cache() can temporarily free up cached memory, but it doesn’t release memory held by tensors still referenced in your Python code. Use torch.cuda.memory_allocated() and torch.cuda.max_memory_allocated() to track usage over time.
  • Fix:
    • Delete unused tensors: Explicitly del tensors and variables that are no longer needed, especially within loops.
    • Use torch.no_grad(): For inference or parts of your training loop that don’t require gradients (like evaluation), wrap them in with torch.no_grad():.
    • Avoid detaching tensors unnecessarily: Detaching a tensor can sometimes lead to it being kept in memory longer if not handled carefully.
    • Check custom __init__ or forward methods: Ensure no tensors are being stored as attributes of your model or modules when they shouldn’t be.
  • Why it works: Python’s garbage collector needs to know that a tensor is truly unreachable to free its memory. Explicitly deleting references or using no_grad ensures tensors and their associated VRAM are released when no longer required.

5. Large Input Data or Intermediate Outputs:

  • Diagnosis: Very high-resolution images, long sequences, or complex data structures can create large tensors. For example, a 4K image (3840x2160x3) can be quite large.
  • Fix:
    • Downsample/Reduce Sequence Length: Preprocess your data to reduce dimensions (e.g., resize images, truncate sequences).
    • Gradient Accumulation (as described above): This indirectly helps by allowing a smaller batch size for the same effective batch size.
    • Mixed Precision Training: Using torch.cuda.amp can reduce the memory footprint of activations.
      from torch.cuda.amp import GradScaler, autocast
      
      scaler = GradScaler()
      
      # Inside training loop:
      with autocast():
          outputs = model(inputs)
          loss = criterion(outputs, targets)
      
      scaler.scale(loss).backward()
      scaler.step(optimizer)
      scaler.update()
      
  • Why it works: Reducing the size of input tensors or using mixed precision (which uses FP16 for activations where possible) directly lowers the memory required to store them.

6. Multiple GPUs and device_map="auto":

  • Diagnosis: If you’re using device_map="auto" with Hugging Face Transformers or similar libraries on a multi-GPU system, it might be distributing your model in a way that one GPU becomes overloaded, even if others have free memory.
  • Fix: Manually define device_map to balance the load or assign specific layers to specific GPUs.
    device_map = {
        "transformer.h.0": 0, "transformer.h.1": 0, # Layers on GPU 0
        "transformer.h.2": 1, "transformer.h.3": 1, # Layers on GPU 1
        # ... and so on
    }
    model = AutoModelForCausalLM.from_pretrained("your_model_name", device_map=device_map)
    
  • Why it works: Explicitly controlling layer placement ensures that memory is distributed more evenly across available GPUs, preventing a single GPU from becoming the bottleneck.

After fixing the OOM error, the next common issue you might encounter is a RuntimeError: CUDA error: an illegal memory access was encountered if your fixes involved subtle pointer issues or if the underlying data corruption caused by the OOM has left the GPU in an unstable state.

Want structured learning?

Take the full Pytorch course →