CUDA out-of-memory during training/inference

CUDA out of memory is the most frustrating way to lose a multi-hour run: everything was going fine, and on a long batch or a long sequence the allocator hit the VRAM ceiling and the process died. If it happened at 3 a.m., you find out at 9 a.m. — half a day gone. A push catches the moment of the crash and immediately suggests what to shrink: batch size, seq length, or enable gradient checkpointing.

There are three places to catch it: in the Python code itself, from the process exit logs, and from the card’s free memory in advance.

Option 1: catch OOM in the training/inference code

PyTorch raises torch.cuda.OutOfMemoryError (a subclass of RuntimeError). We wrap the step and pull out the context — the sizes of the tensor that did not fit:

import os, torch, requests

NOTIFLY_URL   = os.environ["NOTIFLY_URL"]
NOTIFLY_TOKEN = os.environ["NOTIFLY_TOKEN"]

def notify(title, msg, prio):
    requests.post(f"{NOTIFLY_URL}/message",
                  params={"token": NOTIFLY_TOKEN},
                  json={"title": title, "message": msg, "priority": prio},
                  timeout=5)

def run_step(model, batch, run_name):
    try:
        return model(**batch)
    except torch.cuda.OutOfMemoryError as e:
        # what it was at the moment of the crash
        bs      = batch["input_ids"].shape[0]
        seq_len = batch["input_ids"].shape[1]
        free, total = torch.cuda.mem_get_info()
        notify("💥 CUDA OOM",
               f"{run_name}\nbatch_size={bs}, seq_len={seq_len}\n"
               f"VRAM: свободно {free/1e9:.1f}/{total/1e9:.1f} ГБ\n"
               f"Уменьшить batch/seq или включить gradient checkpointing.",
               prio=9)
        torch.cuda.empty_cache()
        raise

priority=9 — the crash costs you GPU-hours and possibly a ruined run, so it is loud. The batch_size/seq_len are the most valuable part: usually it is enough to halve one of them.

Option 2: a guard on exit code and logs

Sometimes an OOM kills the whole process (the kernel OOM-killer, CUDA_ERROR_OUT_OF_MEMORY from the C++ layer) — and the Python handler never runs. Then we guard the process from the outside: launch training through a wrapper that catches the death and greps the log:

#!/usr/bin/env bash
# train_guard.sh — runs training and alerts on OOM
LOG=/var/log/train.log
python train.py "$@" > "$LOG" 2>&1
code=$?

if [[ $code -ne 0 ]]; then
  # pull the last OOM line and its context
  oom=$(grep -iE "out of memory|CUDA_ERROR_OUT_OF_MEMORY|Killed process" "$LOG" | tail -3)
  title="💥 Тренировка упала (exit $code)"
  [[ -n "$oom" ]] && title="💥 CUDA OOM (exit $code)"
  msg=$(printf 'Процесс: train.py\nExit: %s\n\n%s' "$code" "${oom:-см. лог}")
  curl -fsS "$NOTIFLY_URL/message?token=$NOTIFLY_TOKEN" \
    -H "Content-Type: application/json" \
    -d "$(jq -n --arg t "$title" --arg m "$msg" '{title:$t, message:$m, priority:9}')"
fi
exit $code

dmesg/journalctl will show Out of memory: Killed process ... for the kernel OOM-killer — the same grep catches it too.

Option 3: a preventive watchdog on free VRAM

It is better not to crash at all. A separate timer script watches free memory and warns before the OOM, while the run can still be saved:

import subprocess
def free_vram_mib():
    out = subprocess.check_output(
        ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"],
        text=True)
    return min(int(x) for x in out.split())

if free_vram_mib() < 500:        # less than 500 MiB left — an OOM is imminent
    notify("🟠 VRAM почти кончилась",
           f"Свободно {free_vram_mib()} МиБ. Скоро OOM — снизьте нагрузку.",
           prio=7)

Combine with a heartbeat: if after an OOM the process dies silently and stops sending pings, Notifly sends a second, independent alert on timeout.

What to put in the alert text

• the process/run name and the exit code;
• the batch_size and seq_len at the moment of the crash — the key to a fast fix;
• the free/total VRAM on the card;
• a hint: shrink the batch/seq, gradient checkpointing, offload, a smaller dtype.