Posted on
• Artificial Intelligence

Optimising GPU Memory Usage

Author
  • User
    linuxbash
    Posts by this author
    Posts by this author

Optimising GPU Memory Usage on Linux (A Bash-First Guide)

Ever lost hours to a mysterious “CUDA out of memory” error even though your model ran fine yesterday? On shared servers and workstations, GPU VRAM can vanish quickly to zombie processes, greedy frameworks that pre-allocate everything, and fragmentation. The good news: with a few Bash-friendly habits and environment tweaks, you can reclaim stability and squeeze more work out of the same GPU.

This post explains why GPU memory pressure happens, how to see it clearly, and 3–5 practical steps you can apply today. All examples are Linux-friendly and simple to automate in your shell workflow.

Why this matters

  • VRAM is scarce and expensive: A 24 GB GPU can still OOM if your framework pre-allocates memory or fragments small allocations.

  • Shared systems: Other users’ processes (or your old ones) can retain VRAM long after you think they’ve finished.

  • Framework defaults: TensorFlow and JAX often pre-allocate aggressively; PyTorch can fragment memory under heavy allocation churn.

  • Predictability: Consistent, scripted memory hygiene reduces training failures, job preemptions, and wasted queue time.

1) Measure what matters: Live monitoring and quick audits

First, get clear, low-latency visibility into who is holding your VRAM.

  • NVIDIA: nvidia-smi for snapshots; nvtop for a top-like live view.

  • AMD: radeontop for usage and memory.

  • Intel: intel_gpu_top (from intel-gpu-tools) for activity and memory stats on compatible chips.

Example one-liners:

# NVIDIA: one-line summary every 2s
watch -n 2 'nvidia-smi --query-gpu=index,name,memory.used,memory.total --format=csv,noheader'

# NVIDIA: log usage to a file (great for long runs / CI)
while true; do
  date
  nvidia-smi --query-gpu=timestamp,index,memory.used,memory.free --format=csv
  sleep 5
done | tee -a gpu-mem.log

# AMD: live overview
radeontop

# Intel: live overview (needs sudo on some systems)
sudo intel_gpu_top

Quickly list which PIDs hold the most NVIDIA VRAM:

nvidia-smi --query-compute-apps=gpu_uuid,pid,process_name,used_memory --format=csv,noheader | sort -t, -k4 -nr

2) Hunt and evict “ghost” processes

Long-running sessions, crashed notebooks, or detached screens can leave VRAM allocated.

  • Find processes using NVIDIA devices:
nvidia-smi pmon -c 1
nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv,noheader
  • Find any process keeping /dev/nvidia* open:
fuser -v /dev/nvidia* 2>/dev/null
  • AMD/Intel (processes using the render node):
lsof /dev/dri/renderD128 2>/dev/null | awk 'NR==1 || /python|java|ray|tensor|torch|jax|cuda/'
  • Kill gracefully, then force if needed:
kill -15 <PID>    # SIGTERM first
sleep 2
kill -9 <PID>     # SIGKILL if it refuses to exit (last resort)
  • Reset a wedged NVIDIA GPU (danger: requires root, exclusive mode, and no display clients):
sudo nvidia-smi -pm 1        # enable persistence mode
sudo nvidia-smi --gpu-reset  # reset (only safe if no active contexts)

3) Automatically pick the right GPU before you launch

Don’t guess which GPU is “free”—choose by VRAM headroom. Drop this helper into your shell:

choose_free_gpu() {
  # Requires nvidia-smi
  local idx
  idx=$(nvidia-smi --query-gpu=index,memory.free --format=csv,noheader | sort -t, -k2 -nr | head -n1 | cut -d',' -f1 | tr -d ' ')
  if [ -z "$idx" ]; then
    echo "No NVIDIA GPU detected or nvidia-smi not available." >&2
    return 1
  fi
  export CUDA_VISIBLE_DEVICES="$idx"
  # For ROCm users, also:
  export HIP_VISIBLE_DEVICES="$idx"
  echo "Using GPU $idx (CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES)"
}

Usage:

choose_free_gpu && python train.py

This simple step prevents accidental collisions and spreads jobs across devices.

4) Stop frameworks from grabbing all the VRAM

Many frameworks default to “take it all, just in case.” You can tame that with environment variables:

  • TensorFlow:
# Enable incremental growth instead of large upfront allocation:
export TF_FORCE_GPU_ALLOW_GROWTH=true

# Use the async CUDA allocator (often reduces fragmentation on recent TF/CUDA):
export TF_GPU_ALLOCATOR=cuda_malloc_async
  • JAX/XLA:
# Avoid preallocating nearly all memory:
export XLA_PYTHON_CLIENT_PREALLOCATE=false

# Or cap the fraction (e.g., 80% of total VRAM):
export XLA_PYTHON_CLIENT_MEM_FRACTION=0.8
  • PyTorch (CUDA memory allocator tuning):
# Reduce fragmentation by adjusting split size; optionally enable GC threshold:
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128,garbage_collection_threshold:0.7,expandable_segments:True

Tip: Put these in your project’s env.sh and source them per-project, not globally:

set -a
# ...env vars...
set +a

5) Reduce fragmentation and keep sessions stable

  • Reuse process lifetimes: Frequent short runs can fragment memory—prefer longer-lived processes or workers when iterating fast.

  • Free caches between experiments (in code):

    • PyTorch: torch.cuda.empty_cache() after large deallocations.
    • TensorFlow 2: wrap subgraphs in tf.function and delete large tensors explicitly.

  • Enable NVIDIA Persistence Mode so the driver stays initialized (reduces context churn):

sudo nvidia-smi -pm 1
  • On A100/H100 systems with multiple tenants, consider MIG to hard-slice memory/SMs. Example to enable MIG on all GPUs (requires root, supported SKUs only):
sudo nvidia-smi -mig 1

Then create instances with nvidia-smi mig -cgi/-cci as appropriate for your SKU. This enforces isolation and eliminates “neighbor” OOMs at the cost of flexibility.

Real-world mini-scenarios

  • Shared lab server: Two users fight for a 24 GB GPU. Adding choose_free_gpu and setting XLA_PYTHON_CLIENT_MEM_FRACTION=0.8 stopped surprise OOMs and spread jobs more fairly.

  • Notebook churn: PyTorch users restarting cells caused fragmentation and OOM mid-epoch. Setting PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 plus occasional torch.cuda.empty_cache() restored stability.

  • CI pipeline: Sporadic failures vanished after adding a preflight cleanup that kills stale PIDs holding /dev/nvidia* and running with TF_FORCE_GPU_ALLOW_GROWTH=true.

Installation quick-reference (apt, dnf, zypper)

The tools below are widely available. Commands may need sudo.

  • nvtop (live GPU monitor for NVIDIA/AMD/Intel):
# apt (Debian/Ubuntu)
apt update && apt install -y nvtop

# dnf (Fedora/RHEL-compatible)
dnf install -y nvtop

# zypper (openSUSE)
zypper refresh && zypper install -y nvtop
  • radeontop (AMD live monitor):
# apt
apt update && apt install -y radeontop

# dnf
dnf install -y radeontop

# zypper
zypper refresh && zypper install -y radeontop
  • intel-gpu-tools (includes intel_gpu_top):
# apt
apt update && apt install -y intel-gpu-tools

# dnf
dnf install -y intel-gpu-tools

# zypper
zypper refresh && zypper install -y intel-gpu-tools

  • NVIDIA driver (provides nvidia-smi). Package names and repos vary; here are common paths:

    • apt (Ubuntu/Debian):
    # On Ubuntu, the simplest route:
ubuntu-drivers devices
sudo ubuntu-drivers autoinstall
# Or install a specific driver series, e.g.:
sudo apt update && sudo apt install -y nvidia-driver-535
    • dnf (Fedora; enable RPM Fusion first):
    sudo dnf install -y \
https://download1.rpmfusion.org/free/fedora/rpmfusion-free-release-$(rpm -E %fedora).noarch.rpm \
https://download1.rpmfusion.org/nonfree/fedora/rpmfusion-nonfree-release-$(rpm -E %fedora).noarch.rpm

sudo dnf install -y akmod-nvidia xorg-x11-drv-nvidia-cuda
# Reboot after install
    • zypper (openSUSE; add NVIDIA repo matching your distro):
    # Example for Tumbleweed (replace URL for Leap if needed):
sudo zypper addrepo --refresh https://download.nvidia.com/opensuse/tumbleweed NVIDIA
sudo zypper refresh
sudo zypper install -y nvidia-computeG06 nvidia-glG06
# Reboot after install

Note: After installing drivers, nvidia-smi should be available:

nvidia-smi

A simple “preflight” script you can reuse

Save this as gpu_preflight.sh and run it before launching jobs:

#!/usr/bin/env bash
set -euo pipefail

echo "[1/4] Checking NVIDIA GPUs..."
if command -v nvidia-smi >/dev/null 2>&1; then
  nvidia-smi --query-gpu=index,name,memory.used,memory.total --format=csv
  echo "[2/4] Stale processes holding NVIDIA devices:"
  fuser -v /dev/nvidia* 2>/dev/null || true
else
  echo "nvidia-smi not found (skipping NVIDIA checks)."
fi

echo "[3/4] AMD/Intel processes on render node:"
lsof /dev/dri/renderD128 2>/dev/null | awk 'NR==1 || NR<=10' || true

echo "[4/4] Choosing GPU with most free memory (if NVIDIA present)..."
if command -v nvidia-smi >/dev/null 2>&1; then
  idx=$(nvidia-smi --query-gpu=index,memory.free --format=csv,noheader | sort -t, -k2 -nr | head -n1 | cut -d',' -f1 | tr -d ' ')
  if [ -n "${idx}" ]; then
    export CUDA_VISIBLE_DEVICES="$idx"
    export HIP_VISIBLE_DEVICES="$idx"
    echo "Selected GPU $idx (CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES)"
  fi
fi

echo "Tip: export TF_FORCE_GPU_ALLOW_GROWTH=true or set PYTORCH_CUDA_ALLOC_CONF as needed."

Conclusion and next steps

Optimising GPU memory use is equal parts visibility, hygiene, and sane defaults. Start by monitoring in real time, clean up stale contexts, pick the right GPU automatically, and configure your framework to avoid pre-allocation and fragmentation.

Your next steps:

  • Install nvtop and your vendor’s tooling, then add the preflight script to your workflow.

  • Add the environment variables that match your stack (TF/JAX/PyTorch) to your project’s env.sh.

  • Share the choose_free_gpu function with your team to stop accidental collisions.

Have your own battle-tested tips? Send a snippet or gist—we’ll feature the best in a follow-up post.

Related posts