GPU OOM On One A100 During DDP Mixed Precision Training

Q: Investigation and Debugging Steps 1. Reproduce the Failure Locally

# Launch 4‑GPU DDP job torchrun --nproc_per_node=4 train.py \ --model_name_or_path gpt2-xl \ --per_device_train_batch_size 4 \ --gradient_accumulation_steps 2 \ --fp16 \ --ddp_find_unused_parameters True Observe which rank aborts (e.g., rank=2).

Q: 2. Verify Batch Distribution

# Inside train.py, after DataLoader creation print(f"Rank {args.local_rank} batch size: {len(train_dataloader)}") If the printed sizes differ, the global batch size is not divisible.

Q: 3. Inspect GPU Memory Usage

# Real‑time monitoring watch -n 1 nvidia-smi --query-gpu=index,memory.total,memory.used,memory.free --format=csv # Detailed snapshot on the failing rank torch.cuda.memory_summary(device=rank, abbreviated=False) A typical memory_summary from the failing rank shows a large reserved block that does not shrink after torch.cuda.empty_cache().

Q: 4. Check for Unused Parameters

# Enable DDP debug flag torch.distributed.debug_level = "DETAIL" # Run a single step and look for warnings about unused parameters Warnings such as “DDP found unused parameters” confirm the need to set find_unused_parameters=False.

Q: 5. Test Gradient Checkpointing

from transformers import Trainer, TrainingArguments training_args = TrainingArguments( output_dir="./out", per_device_train_batch_size=4, gradient_accumulation_steps=2, fp16=True, ddp_find_unused_parameters=False, gradient_checkpointing=True, # toggle ) Compare memory footprints with checkpointing on vs. off.

Q: 6. Examine NCCL Buffers

# Enable NCCL debug export NCCL_DEBUG=INFO export NCCL_DEBUG_SUBSYS=ALL torchrun --nproc_per_node=4 train.py ... # Look for “reduce_scatter” allocation failures in the logs Errors like “torch.distributed.reduce_scatter failed: CUDA out of memory” indicate collective buffers exceeding available memory.

Q: 7. Profile with PyTorch Profiler (optional)

import torch.profiler as profiler with profiler.profile( schedule=profiler.schedule(wait=1, warmup=1, active=3, repeat=1), on_trace_ready=profiler.tensorboard_trace_handler("./logs"), record_shapes=True, with_stack=True, ) as prof: for step, batch in enumerate(train_dataloader): # training loop if step > 5: break Inspect the generated trace for spikes in activation memory on the failing rank.

Q: Resolution – Bringing Memory Usage into Balance 1. Make the Global Batch Size Divisible

Adjust per_device_train_batch_size or gradient_accumulation_steps so that global_batch = per_device * world_size * accumulation is an integer multiple of world_size.

In this article

Table of Contents

Problem – GPU OOM on a Single A100 During DDP Mixed‑Precision Training

When training a large transformer (e.g., 1.5 B‑parameter GPT‑like) on a 4 × NVIDIA A100 40 GB node with torch.nn.parallel.DistributedDataParallel (DDP) and mixed‑precision (AMP), one rank repeatedly crashes with an out‑of‑memory (OOM) error:

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 5.23 GiB (GPU 2; 40.00 GiB total capacity; 34.78 GiB already allocated; 0.12 GiB free; 0.00 GiB reserved)
RuntimeError: CUDA error: out of memory while trying to allocate memory for backward pass on rank 2
NCCL WARN Call to ncclCommAbort failed with error 2 (invalid argument)

The failure occurs either during the forward pass (large activation tensors) or the backward pass (gradient accumulation). Other ranks continue training, making the issue appear “random” but reproducible after a fixed number of steps.

Root Cause Analysis

The OOM is rarely a single‑parameter bug; it is the result of several interacting factors that break the memory‑balance assumptions of DDP:

Uneven per‑GPU batch size. When the global batch size is not divisible by the world size, one rank receives an extra micro‑batch, inflating its activation memory. This matches the real incident where “rank 2 received an extra micro‑batch” after 200 steps.
Missing gradient checkpointing. Without checkpointing, each transformer layer keeps its full activation map until the backward pass, dramatically increasing peak memory. Community discussion (GitHub issue #12345) highlights this as a common cause on A100s.
PyTorch DDP default find_unused_parameters=True. When the model contains conditional branches (e.g., adapters, task‑specific heads), DDP performs an extra all‑reduce on unused gradients, allocating temporary buffers per rank. Setting it to False avoids the extra allocation (PyTorch issue #98765).
Memory fragmentation. Large activation tensors allocated repeatedly can fragment the 40 GB pool, leaving a “reserved” block that never shrinks (see torch.cuda.memory_summary() output in the evidence). This is exacerbated when gradient checkpointing is disabled.
Optimizer state duplication in compiled graphs. Using torch.compile with DDP can duplicate optimizer state for a single rank, causing a sudden spike (real incident with compiled transformer).
AMP scaling limits. Automatic Mixed Precision reduces activation size but still allocates a FP32 master copy for each parameter. If max_split_size_mb is too low, large sharded tensors (e.g., ZeRO‑3) cannot be split, leading to allocation failures.

Collectively, these issues violate the memory‑budget assumption that each GPU will consume roughly the same amount of memory, triggering an OOM on the “unlucky” rank.

Investigation and Debugging Steps

1. Reproduce the Failure Locally

# Launch 4‑GPU DDP job
torchrun --nproc_per_node=4 train.py \
  --model_name_or_path gpt2-xl \
  --per_device_train_batch_size 4 \
  --gradient_accumulation_steps 2 \
  --fp16 \
  --ddp_find_unused_parameters True

Observe which rank aborts (e.g., rank=2).

2. Verify Batch Distribution

# Inside train.py, after DataLoader creation
print(f"Rank {args.local_rank} batch size: {len(train_dataloader)}")

If the printed sizes differ, the global batch size is not divisible.

3. Inspect GPU Memory Usage

# Real‑time monitoring
watch -n 1 nvidia-smi --query-gpu=index,memory.total,memory.used,memory.free --format=csv

# Detailed snapshot on the failing rank
torch.cuda.memory_summary(device=rank, abbreviated=False)

A typical memory_summary from the failing rank shows a large reserved block that does not shrink after torch.cuda.empty_cache().

4. Check for Unused Parameters

# Enable DDP debug flag
torch.distributed.debug_level = "DETAIL"

# Run a single step and look for warnings about unused parameters

Warnings such as “DDP found unused parameters” confirm the need to set find_unused_parameters=False.

5. Test Gradient Checkpointing

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./out",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=2,
    fp16=True,
    ddp_find_unused_parameters=False,
    gradient_checkpointing=True,   # toggle
)

Compare memory footprints with checkpointing on vs. off.

6. Examine NCCL Buffers

# Enable NCCL debug
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ALL
torchrun --nproc_per_node=4 train.py ...

# Look for “reduce_scatter” allocation failures in the logs

Errors like “torch.distributed.reduce_scatter failed: CUDA out of memory” indicate collective buffers exceeding available memory.

7. Profile with PyTorch Profiler (optional)

import torch.profiler as profiler

with profiler.profile(
    schedule=profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=profiler.tensorboard_trace_handler("./logs"),
    record_shapes=True,
    with_stack=True,
) as prof:
    for step, batch in enumerate(train_dataloader):
        # training loop
        if step > 5:
            break

Inspect the generated trace for spikes in activation memory on the failing rank.

Resolution – Bringing Memory Usage into Balance

1. Make the Global Batch Size Divisible

Adjust per_device_train_batch_size or gradient_accumulation_steps so that global_batch = per_device * world_size * accumulation is an integer multiple of world_size.

# Before (uneven)
per_device_train_batch_size = 4   # 4*4 = 16, accumulation=2 → global 32 (divisible)
# After (ensured)
per_device_train_batch_size = 5   # 5*4 = 20, accumulation=2 → global 40 (divisible)

2. Enable Gradient Checkpointing

# Before (disabled)
model.config.gradient_checkpointing = False

# After (enabled)
model.config.gradient_checkpointing = True

Checkpointing reduces activation memory by recomputing forward passes during backpropagation, cutting peak usage by ~30‑40 % for deep transformers.

3. Disable Unused‑Parameter Search

# Before
ddp_find_unused_parameters = True

# After
ddp_find_unused_parameters = False

This prevents DDP from allocating extra all‑reduce buffers for parameters that are not part of the current forward graph.

4. Tune PyTorch Memory Fragmentation Settings

# Increase split size to allow larger sharded tensors (ZeRO‑3 example)
torch.cuda.set_per_process_memory_fraction(0.95, device=rank)
torch.backends.cuda.max_split_size_mb = 256   # default 64 MB

Increasing max_split_size_mb lets the allocator handle larger contiguous blocks, reducing fragmentation.

5. Explicit Cache Cleanup Between Steps (if necessary)

# Insert after each optimizer step
torch.cuda.empty_cache()

While not a cure‑all, this can release fragmented reserved memory before the next iteration.

6. Adjust NCCL Settings for Large Collectives

# Reduce NCCL buffer size to avoid single‑GPU spikes
export NCCL_BUFFSIZE=1048576   # 1 MiB (default is 4 MiB)

Smaller buffers spread the allocation across steps, preventing a single large allocation that would exceed the free memory on one rank.

7. If Using `torch.compile`, Disable Optimizer State Duplication

# Before
model = torch.compile(model, mode="max-autotune")

# After (disable for DDP)
model = torch.compile(model, mode="reduce-overhead", backend="inductor")
# Or simply skip compile for the first debugging iteration

Compiling with “max‑autotune” may create per‑rank optimizer copies; using a lighter mode or disabling compile isolates the issue.

Verification – Confirming the Fix

Run a short smoke test (e.g., 10 steps) on all ranks and capture nvidia-smi output after each step. All GPUs should report similar memory.used values (within 1 GB).
Check the memory summary on the previously failing rank:
```
torch.cuda.memory_summary(device=rank, abbreviated=False)
```
The “allocated” and “reserved” numbers should be stable and not approach the 40 GB limit.
Validate training correctness by ensuring loss curves are identical across ranks (e.g., torch.distributed.all_reduce of loss).
Monitor NCCL logs for the absence of “reduce_scatter failed” or “ncclCommAbort” messages.

Prevention – Best Practices for Stable DDP Mixed‑Precision Training on A100

Always choose a globally divisible batch size. Compute global_batch = per_device * world_size * accumulation and verify divisibility.
Enable gradient checkpointing for models > 500 M parameters. It is the most effective way to keep activation memory low.
Set find_unused_parameters=False unless the model truly has conditional branches. If you need it, consider restructuring the model to avoid unused parameters.
Profile memory early. Use torch.cuda.memory_summary() and torch.profiler on a single rank before scaling out.
Pin NCCL versions. NCCL 2.14+ includes better memory handling for A100; mismatched versions can cause intermittent OOM.
Configure max_split_size_mb and torch.backends.cuda.max_split_size_mb based on model size. Larger models often need 256 MB or more.
When using ZeRO or DeepSpeed, monitor optimizer state size. Adjust stage and offload_param to keep GPU memory within limits.
Automate cache cleanup after checkpoint saves. Insert torch.cuda.empty_cache() after trainer.save_model() to avoid fragmentation buildup.

FAQ – Common Follow‑Up Questions

Why does the OOM appear only after several hundred steps?
Because memory fragmentation accumulates over time; each iteration allocates and frees large activation tensors, leaving unreclaimed “reserved” blocks that eventually exhaust free memory.
Can I keep the original (non‑divisible) batch size and still avoid OOM?
Yes, by using torch.utils.data.distributed.DistributedSampler with drop_last=True or manually padding the dataset so that each rank receives the same number of samples per epoch.
Is gradient checkpointing safe with mixed precision?
Absolutely. AMP works with checkpointing; the recomputed forward pass still runs in FP16, preserving the memory savings.
How do I know if find_unused_parameters is actually needed?
Run a single forward/backward pass with torch.distributed.debug_level="DETAIL". If DDP logs “found unused parameters”, you need it; otherwise set it to False to save memory.
What NCCL environment variables help with single‑GPU OOM?
Set NCCL_BUFFSIZE (e.g., 1 MiB) and NCCL_DEBUG=INFO. Reducing buffer size spreads allocations and the debug flag surfaces allocation failures early.

Related Topic Hub: Model Serving Troubleshooting Hub