How to Handle GPU Memory Limits and OOM in GitHub Actions

Cap batch size, call torch.cuda.empty_cache between tests, and assert peak memory so a job stays within the runner GPU and OOM is caught as a test failure.

Reset the peak-memory counter, run the step, then assert max_memory_allocated stays under a budget. Clear the cache between tests so one heavy test does not OOM the next.

Steps

• Reset peak memory with torch.cuda.reset_peak_memory_stats().
• Run the step, then assert max_memory_allocated is under budget.
• Call torch.cuda.empty_cache() between tests.

Workflow

steps:
  - name: Peak memory gate
    run: |
      python - <<'PY'
      import torch
      from model import build, batch
      torch.cuda.reset_peak_memory_stats()
      m = build().cuda(); x, _ = batch()
      m(x.cuda()).mean().backward()
      peak = torch.cuda.max_memory_allocated() / 1e9
      print(f"peak {peak:.2f} GB")
      assert peak < 14.0, "exceeds GPU memory budget"
      torch.cuda.empty_cache()
      PY

Gotchas

• Fragmentation can OOM below the nominal limit; set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to reduce it.
• On a shared self-hosted GPU, another job can consume memory; serialize GPU jobs with a concurrency group.

Related guides