Situation: Intermittent GPU OOMs during training can be caused by transient allocation spikes, fragmentation, non-deterministic batching, race conditions in distributed runs, or memory leaks from retained tensors. Below is an FAQ + troubleshooting guide targeted to DevOps and ML engineers.FAQ — probable root causes- Batch-size variance: last batch larger or augmentation increases memory.- Non-deterministic DataLoader yielding different shapes.- Accumulated gradients or optimizer states growing (e.g., fp16 scaler).- Memory fragmentation from many short-lived allocations.- Leaked references (keeping tensors in lists/logging).- Multi-process/device contention (shared GPU between jobs).- GPU driver/CUDA non-determinism or bug.- Distributed training race where rank occasionally holds extra tensors.Quick diagnostics (commands)- Check basic GPU usage:nvidia-smi --query-gpu=index,name,utilization.gpu,utilization.memory,memory.total,memory.used -f csv -l 2
- List processes on GPU:nvidia-smi pmon -c 1
ps -ef | grep python
- System-level metrics (sampling):nvidia-smi dmon -s u -d 2
- Container cgroup memory limits:cat /sys/fs/cgroup/memory/memory.usage_in_bytes
In-process diagnostics (PyTorch examples)- Snapshot memory stats:import torch
print(torch.cuda.memory_summary(device=0, abbreviated=False))
print("allocated:", torch.cuda.memory_allocated())
print("max_allocated:", torch.cuda.max_memory_allocated())
- Reset/track peaks:torch.cuda.reset_peak_memory_stats()
# run step
print(torch.cuda.max_memory_allocated())
- Synchronous stack trace on OOM:set environment before running:export CUDA_LAUNCH_BLOCKING=1
- Capture profiler trace:from torch.profiler import profile, ProfilerActivity
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as p:
model(input)
print(p.key_averages().table(sort_by="cuda_time_total"))
p.export_chrome_trace("trace.json")
How to capture reproducible OOM traces1. Create a minimal script that reproduces the failure deterministically: fix random seeds, use a fixed batch from disk, and run the single problematic step in a loop.2. Run with CUDA_LAUNCH_BLOCKING=1 to get accurate stack traces.3. Use torch.cuda.reset_peak_memory_stats() before the step and record torch.cuda.max_memory_allocated() after.4. Use PyTorch profiler to capture allocation shapes and kernels (export trace.json to open in Chrome tracing or Nsight).5. If intermittent only in multi-GPU, reproduce with the same number of processes and use logs per-rank.6. Collect nsys trace:nsys profile -o trace -t cuda,osrt,drm -- python train.py ...
Temporary mitigations (fast to apply)- Reduce batch size by 1–2 steps.- Enable mixed precision (AMP) to cut activation/optimizer memory:from torch.cuda.amp import GradScaler, autocast
with autocast():
out = model(x)
- Gradient accumulation: keep effective batch size without increasing memory:accumulate_steps = 4
loss = ...
loss = loss / accumulate_steps
loss.backward()
if (step+1) % accumulate_steps == 0:
optimizer.step(); optimizer.zero_grad()
- Free caches between iterations:del tensors_you_no_longer_need
torch.cuda.empty_cache()
- Disable pin_memory=True in DataLoader if CPU->GPU pinned memory causing pressure.- Reduce num_workers or prefetch_factor to avoid queued allocations.Permanent fixes / robust changes- Model architecture: reduce hidden dims, remove redundant layers, use depthwise separable convs.- Activation checkpointing (recompute activations to trade compute for memory).- Use optimizer state sharding (ZeRO from DeepSpeed/Fairscale) to reduce per-GPU memory.- Use model parallelism or pipeline parallelism when model exceeds single GPU limits.- Convert parameters/optimizer states to fp16 with careful loss-scaling or use bfloat16 if supported.- Fix code-level leaks: avoid storing tensors in global lists, ensure detach().cpu() before logging.- Upgrade drivers/CUDA/cuDNN and framework to versions with known memory fixes.- Use static shapes for inputs (pad/truncate) to avoid dynamic allocation spikes.Debugging checklist- Reproduce with minimal script and seed.- Capture GPU allocation trace (nsys/torch.profiler).- Verify process count and shared GPU usage.- Inspect logs for gradually increasing RAM/GPU memory → memory leak.- Run with reduced batch size to confirm threshold.- Try torch.backends.cudnn.benchmark=False to reduce nondeterministic workspace allocations.Examples of subtle bugs- Logging a tensor each iteration without .item() or .cpu() retains GPU memory.- Using torch.no_grad() incorrectly so gradients retained.- Mixed precision plus in-place ops causing unexpected growth.When to escalate to infra/driver- If nvidia-smi shows unexpected non-Python processes using memory.- If OOM occurs only under container orchestration — check cgroup limits and GPU sharing settings (MPS on multi-tenant).- If profiling shows CUDA runtime error or driver crashes — collect driver logs (dmesg, /var/log/nvidia-installer.log) and open an issue with framework and driver versions and captured traces.Example minimal reproducible pattern (PyTorch)import torch, random
torch.manual_seed(0)
x = torch.randn(8,3,224,224, device='cuda')
model = MyModel().cuda()
torch.cuda.reset_peak_memory_stats()
with torch.no_grad():
for i in range(100):
out = model(x) # log shapes, probe memory
print(torch.cuda.max_memory_allocated())
Final notes- Start with simplest checks (batch size, other processes). Use CUDA_LAUNCH_BLOCKING and profilers to obtain deterministic traces. Apply temporary mitigations to continue training while implementing longer-term fixes like checkpointing, ZeRO, or architecture changes. Collect and attach profiler traces and nsys outputs when filing issues.