a golden statue of a man and a woman on a boat

GPU OOM errors during model training after Prometheus alert in sandbox

by

Problem

During iterative development in the shared sandbox, data‑science engineers observed frequent GPU out‑of‑memory (OOM) failures when launching PyTorch training jobs. The failures coincided with a Prometheus alert that fired on the gpu_memory_usage_bytes metric exceeding a static threshold. The alert itself did not trigger any remediation, and the training jobs aborted with the following error pattern:


Traceback (most recent call last):
  File "train.py", line 112, in <module>
    trainer.fit(model, dataloader)
  File "/opt/conda/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1245, in fit
    self._run_stage()
  File "/opt/conda/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1328, in _run_stage
    self._run_train()
  File "/opt/conda/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1653, in _run_train
    self._run_training_epoch()
  File "/opt/conda/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1739, in _run_training_epoch
    self._run_training_batch()
  File "/opt/conda/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1765, in _run_training_batch
    self._run_backward()
  File "/opt/conda/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1792, in _run_backward
    self.backward(loss, optimizer, optimizer_idx)
  File "/opt/conda/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1820, in backward
    self._run_backward_pass(loss, optimizer, optimizer_idx)
  File "/opt/conda/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1845, in _run_backward_pass
    loss.backward()
RuntimeError: CUDA out of memory. Tried to allocate 2.00 GiB. 
GPU 0:  12.34 GiB total,  11.12 GiB allocated,  0.00 GiB free,  0.00 GiB cached

The sandbox hosts multiple concurrent training pods, each requesting a single GPU via the NVIDIA device plugin. The alert threshold is set to 90 % of total GPU memory, which was intended as an early‑warning signal but appears to be a catalyst for the OOM condition.

Root Cause

The OOM events stem from a combination of:

  • Prometheus‑driven GPU memory reclamation script: A sidecar container watches gpu_memory_usage_bytes and, when the threshold is crossed, issues nvidia-smi --gpu-reset on the affected device to “free” memory. The reset command forces the driver to drop all active contexts, including those belonging to running training jobs.
  • Improper isolation of the reset operation: The reset script runs with --all-gpus flag due to a copy‑paste error, resetting every GPU on the node regardless of which one triggered the alert.
  • Training jobs do not checkpoint before OOM: The training loop lacks a torch.cuda.empty_cache() guard or periodic checkpoint, so when the driver resets, the process receives a CUDA driver error: unknown error that propagates as an OOM exception.

Thus, the alert is not merely a symptom; the remediation logic actively destroys the GPU context, turning a high‑memory warning into a fatal OOM.

Debug

Step‑by‑step investigation that reproduced the failure:

  1. Confirm alert firing

$ curl -s http://prometheus:9090/api/v1/query?query=gpu_memory_usage_bytes
{
  "status":"success",
  "data":{
    "resultType":"vector",
    "result":[
      {
        "metric":{"instance":"node-01:9100","gpu":"0"},
        "value":[1697445600,"12456789012"]
      }
    ]
  }
}

The value translates to ~11.6 GiB used on a 12 GiB card (≈96 %).

  1. Inspect the reset script logs

2024-10-15T14:23:07Z INFO  gpu‑reset‑watcher: GPU 0 usage 96% > threshold 90%
2024-10-15T14:23:07Z WARN  gpu‑reset‑watcher: Executing nvidia-smi --gpu-reset --all-gpus
2024-10-15T14:23:08Z INFO  gpu‑reset‑watcher: Reset completed, returning 0

The --all-gpus flag is evident.

  1. Correlate pod logs with reset timestamps

2024-10-15T14:23:06Z INFO  trainer: Starting epoch 3
2024-10-15T14:23:07Z INFO  trainer: Forward pass completed, loss=0.842
2024-10-15T14:23:08Z ERROR trainer: CUDA driver error: unknown error
2024-10-15T14:23:08Z ERROR trainer: RuntimeError: CUDA out of memory.

The error appears immediately after the reset command.

  1. Validate driver state before and after reset

# Before reset
$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap| Memory-Usage | GPU-Util  Compute M. |
|                               |                  |               |
|  0  Tesla T4            Off  | 00000000:00:1E.0 Off | 0                  |
| N/A   45C    P0    30W /  70W | 11258MiB / 15109MiB |   98%      Default |
+-------------------------------+----------------------+----------------------+

# After reset (triggered by script)
$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap| Memory-Usage | GPU-Util  Compute M. |
|                               |                  |               |
|  0  Tesla T4            Off  | 00000000:00:1E.0 Off | 0                  |
| N/A   40C    P8    10W /  70W | 0MiB / 15109MiB |    0%      Default |
+-------------------------------+----------------------+----------------------+

The reset clears all memory allocations, confirming the driver‑level impact.

Solution

Two complementary changes eliminate the OOM cascade:

1. Fix the reset script

Restrict the reset to the offending GPU only and add a safety guard that aborts if any container has an active CUDA context.

# Before (buggy)
#!/usr/bin/env bash
THRESHOLD=0.90
while true; do
  USAGE=$(curl -s http://prometheus:9090/api/v1/query?query=gpu_memory_usage_bytes | jq -r '.data.result[0].value[1]')
  TOTAL=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits)
  PCT=$(awk "BEGIN {print $USAGE/$TOTAL}")
  if (( $(echo "$PCT > $THRESHOLD" | bc -l) )); then
    logger "GPU usage $PCT > $THRESHOLD, resetting all GPUs"
    nvidia-smi --gpu-reset --all-gpus
  fi
  sleep 30
done

# After (corrected)
#!/usr/bin/env bash
THRESHOLD=0.90
while true; do
  # Query per‑GPU usage
  RESPONSE=$(curl -s http://prometheus:9090/api/v1/query?query=gpu_memory_usage_bytes)
  for ROW in $(echo "$RESPONSE" | jq -c '.data.result[]'); do
    GPU=$(echo "$ROW" | jq -r '.metric.gpu')
    USAGE=$(echo "$ROW" | jq -r '.value[1]')
    TOTAL=$(nvidia-smi --id=$GPU --query-gpu=memory.total --format=csv,noheader,nounits)
    PCT=$(awk "BEGIN {print $USAGE/$TOTAL}")
    if (( $(echo "$PCT > $THRESHOLD" | bc -l) )); then
      # Verify no active CUDA processes on this GPU
      PROCS=$(nvidia-smi --id=$GPU --query-compute-apps=pid --format=csv,noheader)
      if [ -z "$PROCS" ]; then
        logger "GPU $GPU usage $PCT > $THRESHOLD, resetting GPU $GPU"
        nvidia-smi --id=$GPU --gpu-reset
      else
        logger "GPU $GPU over threshold but active processes detected (pids: $PROCS); skipping reset"
      fi
    fi
  done
  sleep 30
done

Key changes:

  • Iterate over each GPU metric instead of assuming a single device.
  • Use --id=$GPU to target the specific card.
  • Check for active compute processes before resetting; if any are present, log and skip.

2. Make training jobs resilient to transient OOM

Introduce checkpointing and explicit cache clearing after each batch. The following snippet demonstrates a robust training loop:

# Before (no checkpoint, no cache clear)
for epoch in range(num_epochs):
    for batch in dataloader:
        optimizer.zero_grad()
        loss = model(batch)
        loss.backward()
        optimizer.step()

# After (adds checkpoint & cache clear)
import torch
import os

CHECKPOINT_DIR = "/mnt/checkpoints"
os.makedirs(CHECKPOINT_DIR, exist_ok=True)

def save_checkpoint(epoch, model, optimizer):
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
    }, os.path.join(CHECKPOINT_DIR, f"ckpt_epoch_{epoch}.pt"))

for epoch in range(num_epochs):
    for batch_idx, batch in enumerate(dataloader):
        optimizer.zero_grad()
        loss = model(batch)
        loss.backward()
        optimizer.step()

        # Periodic cache clear to avoid fragmentation
        torch.cuda.empty_cache()

        # Save checkpoint every 500 batches
        if batch_idx % 500 == 0:
            save_checkpoint(epoch, model, optimizer)

These changes ensure that if a reset does occur, the job can resume from the latest checkpoint without losing all progress.

Verify

After deploying the corrected reset script and updating the training code, perform the following validation steps:

  1. Simulate high memory usage by running a synthetic workload that allocates 95 % of GPU memory.
  2. Observe Prometheus metrics to confirm the alert fires.
  3. Check the reset script log – it should now report “active processes detected; skipping reset”.
  4. Confirm that the synthetic workload continues running without termination.
  5. Run a real training job to completion and verify that checkpoint files appear in /mnt/checkpoints and that the final loss converges.

Sample verification output:


2024-10-15T15:02:10Z INFO  gpu‑reset‑watcher: GPU 0 usage 0.96 > 0.90, active processes detected (pids: 31245 31246); skipping reset
2024-10-15T15:02:12Z INFO  trainer: Epoch 5 completed, checkpoint saved to /mnt/checkpoints/ckpt_epoch_5.pt
2024-10-15T15:02:12Z INFO  trainer: GPU memory after torch.cuda.empty_cache(): 2.1 GiB free

Prevent

Implementing guardrails reduces the chance of recurrence:

  • Metric‑driven alerts only: Configure Prometheus alerts to fire without automatic remediation. Use Alertmanager to route to a human on‑call.
  • Separate reset responsibilities: Run the reset script on a dedicated management node with read‑only access to GPU metrics; never grant it --all-gpus privileges.
  • GPU memory headroom: Adjust the alert threshold to 80 % and add a hysteresis period (e.g., 5 minutes) before any action is considered.
  • Monitoring of reset actions: Export a custom Prometheus metric gpu_reset_attempts_total and gpu_reset_success_total to track how often resets are attempted and whether they succeed.
  • Container resource limits: Enforce limits.nvidia.com/gpu: 1 and set requests.nvidia.com/memory via the NVIDIA device plugin (if supported) to prevent a single pod from monopolizing the GPU.

Related Topic Hub: Observability Troubleshooting Hub

FAQ

  1. Why does the OOM only appear after the Prometheus alert? The alert triggers a reset script that forcibly clears GPU contexts. If the script is mis‑configured (e.g., uses --all-gpus), it resets the GPU while a training job is still active, causing the driver to report OOM.
  2. Can I keep the reset script but make it safe? Yes. Restrict resets to GPUs with no active compute processes, add a grace period, and avoid the --all-gpus flag. Consider using nvidia-smi --gpu-reset --id=<GPU> only after confirming nvidia-smi --query-compute-apps returns empty.
  3. Is torch.cuda.empty_cache() enough to prevent OOM? It helps mitigate fragmentation but does not replace proper memory budgeting. Use it in conjunction with checkpointing and realistic batch sizes.
  4. How should I set the Prometheus alert threshold? A conservative threshold (≈80 %) with a 5‑minute sustained period reduces false positives while still providing early warning. Pair the alert with a runbook that suggests manual intervention rather than automatic resets.
  5. What if I need to reset a GPU that is truly stuck? Perform a manual reset from the node console after confirming no critical workloads are running: nvidia-smi --gpu-reset --id=<GPU>. Document the event and investigate why the GPU entered a stuck state (e.g., driver bug, hardware fault).