GPU memory is the most constrained resource in deep learning. When you scale from a single GPU to distributed training using DistributedDataParallel (DDP) or Fully Sharded Data Parallel (FSDP), memory overhead increases significantly. You often encounter the dreaded RuntimeError: CUDA out of memory even when your hardware seems sufficient. This guide provides actionable fixes for PyTorch 2.0+ to stabilize your training loops and maximize hardware utilization.
TL;DR — To stop CUDA OOM, reduce per_device_batch_size and use Gradient Accumulation. Enable torch.amp for mixed precision and checkpoint for trade-off memory for compute. Set max_split_size_mb via environment variables to prevent fragmentation.
Table of Contents
Common Symptoms of Distributed OOM
💡 Analogy: Think of your GPU memory as a kitchen counter. Training a model is like cooking a complex meal. If the counter is cluttered with dirty dishes (old gradients) and too many ingredients (large batches), you run out of space to actually chop the vegetables (compute activations). Distributed training adds more chefs (GPUs), each needing their own space plus a shared table for communication.
A typical OOM error in a distributed environment looks like this: RuntimeError: CUDA out of memory. Tried to allocate 512.00 MiB (GPU 0; 23.65 GiB total capacity; 21.50 GiB already allocated; 128.50 MiB free; 21.80 GiB reserved in total by PyTorch). In multi-GPU setups, you might notice that only Rank 0 crashes, or that memory usage slowly creeps up over several epochs before a sudden failure.
Pay close attention to the difference between Allocated and Reserved memory. If reserved memory is much higher than allocated, your GPU is suffering from fragmentation. If they are nearly equal, you are simply pushing too much data through the pipeline. In distributed training, the NCCL backend also requires a buffer for gradient synchronization, which consumes a few hundred MBs of "invisible" overhead on each device.
Root Causes of GPU Memory Spikes
Improper Batch Distribution
In DistributedDataParallel, the batch_size you define in your DataLoader is usually the per-GPU batch size, not the global batch size. If you set a batch size of 32 and scale to 8 GPUs, your effective batch size is 256. Beginners often forget that each GPU needs to fit those 32 samples plus their corresponding activations and gradients. If Rank 0 is also handling logging or validation, it will hit the OOM threshold faster than other ranks.
Activation Storage during Forward Pass
The forward pass stores "activations" for every layer to calculate gradients during the backward pass. For deep models like Transformers, these activations often take up 3x to 5x more space than the model weights themselves. In distributed settings, if you are using long sequence lengths (e.g., 2048 tokens), the activation tensors grow quadratically with the sequence length, leading to immediate crashes during the first iteration.
Memory Fragmentation
PyTorch uses a caching allocator to speed up memory management. However, if your model frequently allocates and deallocates tensors of varying sizes, the "holes" between allocated blocks become too small to fit a new large tensor. This results in an OOM error even if the total "free" memory appears sufficient in nvidia-smi.
Effective Fixes for PyTorch OOM
1. Employ Automatic Mixed Precision (AMP)
Using 16-bit precision (FP16 or BF16) instead of 32-bit (FP32) can reduce memory consumption by nearly 50%. BF16 is preferred on NVIDIA Ampere (A100/A30) and newer GPUs because it has the same dynamic range as FP32, eliminating the need for loss scaling.
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for inputs, targets in dataloader:
optimizer.zero_grad()
# Casts operations to mixed precision
with autocast(device_type='cuda', dtype=torch.bfloat16):
outputs = model(inputs)
loss = criterion(outputs, targets)
# Scales loss and calls backward
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
2. Apply Gradient Checkpointing
If your model is too deep, use torch.utils.checkpoint. This technique drops activations during the forward pass and re-computes them during the backward pass. It trades a ~20% increase in computation time for a massive reduction in memory usage.
from torch.utils.checkpoint import checkpoint
def forward(self, x):
# Instead of: x = self.block(x)
x = checkpoint(self.block, x, use_reentrant=False)
return x
3. Use Gradient Accumulation
Instead of a per-GPU batch size of 32, use a batch size of 8 and accumulate gradients over 4 steps. This mimics the mathematical behavior of a larger batch without the memory pressure.
accumulation_steps = 4
for i, (inputs, targets) in enumerate(dataloader):
outputs = model(inputs)
loss = criterion(outputs, targets) / accumulation_steps
loss.backward()
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
⚠️ Common Mistake: Do not use model.zero_grad() inside the accumulation loop. Only clear gradients after the optimizer.step() call, or you will lose the accumulated progress from previous steps.
Verification and Monitoring
To verify if your fixes are working, you must look beyond nvidia-smi. Use PyTorch's internal memory summary to identify exactly where the "peak" allocation occurs. Run this command after the first 10 iterations of your training loop:
print(torch.cuda.memory_summary(device=None, abbreviated=False))
Check the "Max Reserved Memory". If this value is close to your GPU's physical limit, you have no buffer for unexpected spikes during validation or logging. You should also verify that the NCCL backend is not timing out. In distributed training, if one GPU OOMs, the others will hang indefinitely. Set export NCCL_DEBUG=INFO in your shell to see detailed logs of the communication state during a crash.
Prevention Strategies
Preventing OOM errors requires proactive environment configuration. Set the PYTORCH_CUDA_ALLOC_CONF environment variable to manage how PyTorch interacts with the CUDA driver. This is the most effective way to combat fragmentation without changing your code.
# Set this in your .bashrc or at the top of your script
export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128"
📌 Key Takeaways
- Mixed Precision: Always use BF16 or FP16 for modern LLM/Vision training.
- FSDP over DDP: For models > 2B parameters, use
FullyShardedDataParallelto shard weights across GPUs. - Clear Cache: Use
torch.cuda.empty_cache()between training and validation loops to reset the allocator. - Profile Early: Use the PyTorch Profiler to find memory-heavy layers before scaling to a cluster.
Frequently Asked Questions
Q. Why does OOM only happen on Rank 0 in PyTorch DDP?
A. Rank 0 often handles additional tasks like calculating metrics, logging to Weights & Biases, or saving checkpoints. These tensors stay in memory on GPU 0. To fix this, ensure you move logging tensors to the CPU using .detach().cpu() before processing.
Q. Does torch.cuda.empty_cache() actually solve OOM errors?
A. No, it does not increase the total memory available. It simply releases "reserved" memory back to the OS. It is helpful for fragmentation issues but will not help if your model and batch size genuinely exceed your GPU's capacity.
Q. How much memory does gradient checkpointing actually save?
A. It can reduce activation memory from O(N) to O(sqrt(N)), where N is the number of layers. For a typical Transformer, this can allow you to double or triple your sequence length on the same hardware.
Post a Comment