A batch job uses too much GPU memory - rFronteddu/general_wiki GitHub Wiki

1. Symptoms
2. Common causes
3. Fixes
4. Metrics

Symptoms

A script runs a model over many inputs and either crashes with CUDA OOM or uses more memory than expected.

• Add watch:

watch -n 1 nvidia-smi

• Measure peak GPU memory

import torch

torch.cuda.reset_peak_memory_stats()

# run workload

print("allocated GB:", torch.cuda.memory_allocated() / 1e9)
print("reserved GB:", torch.cuda.memory_reserved() / 1e9)
print("peak allocated GB:", torch.cuda.max_memory_allocated() / 1e9)

• Print summary

print(torch.cuda.memory_summary())

Common causes

• Batch size is too large.
• Code tracks gradients during inference.
• The model is in training mode.
• Outputs are stored without detaching from computation graph.
• A list keeps accumulating GPU tensors.
• Data is copied unnecessarily.
• Multiple model copies exist on GPU.
• Intermediate tensors are not released.
• The job uses float32 when float16/bfloat16 is enough.
• First safe fixes

Fixes

• For inference

model.eval()

with torch.inference_mode():
    output = model(input)

• If they are using:

with torch.no_grad():

That is also okay, but inference_mode() is usually stronger for inference.

• Reduce batch size

batch_size = batch_size // 2

• Use mixed precision if appropriate

with torch.inference_mode():
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        output = model(input)

• Or on newer NVIDIA GPUs, sometimes:

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    output = model(input)

• Do not store GPU tensors forever Bad:

outputs = []

for batch in loader:
    output = model(batch)
    outputs.append(output)

Better:

outputs = []

for batch in loader:
    with torch.inference_mode():
        output = model(batch)

    outputs.append(output.detach().cpu())

• If you only need final values, move them to CPU or write them to disk.
• do not treat torch.cuda.empty_cache() as a real fix. It can help release cached memory back to CUDA, but it usually does not solve the real memory bug.

Metrics

Baseline peak GPU memory: X GB
OOM occurs at batch size: N

Findings
- gradients were being tracked during inference
- outputs were stored on GPU
- batch size was too large
- memory grows each iteration, suggesting accumulation

Change
- added inference_mode
- moved outputs to CPU
- reduced batch size
- tested autocast

Result
- peak memory reduced from X GB to Y GB
- job completed successfully