Skip to content

GPU

Last reviewed: 2026-06-16

Purpose: Reference notes on GPU computing stacks for AI/ML workloads — drivers, CUDA, container toolkits, and monitoring.

Contents

Overview

The GPU compute stack is the foundational software layer required to run AI/ML workloads on NVIDIA hardware. The canonical stack (bottom to top) is:

  1. NVIDIA GPU Driver — Kernel-mode driver that allows the OS to communicate with the GPU.
  2. CUDA Toolkit — Compiler, runtime libraries, and development tools for GPU-accelerated computing.
  3. cuDNN — Deep neural network primitives library (convolution, pooling, normalisation, etc.).
  4. nvidia-container-toolkit — Allows Docker/Podman containers to access host GPUs.
  5. ML Framework — PyTorch, TensorFlow, JAX, etc., each shipped with pre-built CUDA wheels.

Compatibility is critical. Each layer must be built against a compatible CUDA version. PyTorch wheels are pinned to specific CUDA minor versions (e.g., CUDA 12.4, 12.1). Always check the framework's official compatibility matrix before installing.

NVIDIA Driver Installation (Ubuntu)

1. Identify your GPU

lspci | grep -i nvidia

If the output is empty, you may need to install pciutils: 

sudo apt update && sudo apt install -y pciutils

2. Remove existing or conflicting drivers

sudo apt purge --auto-remove nvidia-* libnvidia-*
sudo apt autoremove

Reboot after removal.

# Detect Ubuntu version
ubuntu_version=$(lsb_release -cs)

# Add NVIDIA driver repository
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${ubuntu_version}/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt update

Then install the metapackage that pulls in the latest stable driver and CUDA together (see CUDA section below), or install the driver alone:

# Install driver only (latest branch)
sudo apt install -y nvidia-driver-550

# Or list available versions
apt search nvidia-driver

Branch notes (as of mid-2026): - nvidia-driver-550 — Latest Production Branch (recommended for most users) - nvidia-driver-545 — Previous Production Branch - nvidia-driver-535 — Long-lived Production Branch

For data-centre GPUs (A100/H100/B200), use the datacentre branch: 

sudo apt install -y nvidia-driver-550-server

4. Reboot and verify

sudo reboot
# After reboot:
nvidia-smi

Expected output shows GPU model, driver version, CUDA version, and process list:

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 550.xxx   Driver Version: 550.xxx   CUDA Version: 12.4          |
|-------------------------------+----------------------+----------------------+

Alternative: Ubuntu's ubuntu-drivers tool

sudo ubuntu-drivers autoinstall

This picks the driver version recommended by Ubuntu. Usually adequate but may trail the latest NVIDIA release.

CUDA Toolkit Setup

Toolkit vs. Driver CUDA version

The driver reports a "CUDA Version" — this is the maximum CUDA toolkit version the driver supports, not the installed toolkit version. You can install any CUDA toolkit ≤ the driver's reported version.

# After adding cuda-keyring (see Driver section), install the full toolkit:
sudo apt install -y cuda-toolkit-12-4

This installs nvcc, cuda-runtime, cudart, cublas, curand, cufft, etc. under /usr/local/cuda-12.4/.

Set environment variables

Add to ~/.bashrc (or ~/.zshrc):

export PATH=/usr/local/cuda-12.4/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

Apply:

source ~/.bashrc

Verify installation

nvcc --version

Should print the CUDA release version and build info.

Using the runfile installer (isolated, no root)

Download the CUDA runfile from https://developer.nvidia.com/cuda-downloads and run:

wget https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run
sudo sh cuda_12.4.0_550.54.14_linux.run --toolkit --silent

Note: The runfile also bundles a driver. Pass --toolkit to install only the toolkit.

cuDNN Installation

cuDNN (CUDA Deep Neural Network library) provides highly tuned implementations of standard routines — forward/backward convolution, pooling, normalisation, activation layers.

Prerequisites

Local .deb install (Ubuntu)

# 1. Download from NVIDIA Developer portal (example filename)
#    For CUDA 12.x, cuDNN 9.x:
#    cudnn-local-repo-ubuntu2204-9.3.0_1.0-1_amd64.deb

sudo dpkg -i cudnn-local-repo-ubuntu2204-9.3.0_1.0-1_amd64.deb
sudo cp /var/cudnn-local-repo-ubuntu2204/cudnn-*-keyring.gpg /usr/share/keyrings/
sudo apt update

# 2. Install the runtime and dev libraries
sudo apt install -y libcudnn9-cuda-12 libcudnn9-dev-cuda-12

Verify cuDNN

# Check installed files
dpkg -l | grep cudnn

# Or test via a small C program:
cat << 'EOF' > /tmp/cudnn_test.c
#include <cudnn.h>
#include <stdio.h>

int main() {
    cudnnHandle_t handle;
    cudnnCreate(&handle);
    printf("cuDNN version: %d\n", cudnnGetVersion());
    cudnnDestroy(handle);
    return 0;
}
EOF

nvcc -o /tmp/cudnn_test /tmp/cudnn_test.c -lcudnn
/tmp/cudnn_test

Expected output: cuDNN version: 9300 (for cuDNN 9.3.0).

nvidia-container-toolkit (Docker GPU Passthrough)

This tool configures Docker (or containerd, Podman) so containers can see and use host GPUs.

1. Install the NVIDIA Container Toolkit

# Add the repository
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
  sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
  sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

sudo apt update
sudo apt install -y nvidia-container-toolkit

2. Configure the container runtime

sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

3. Test GPU passthrough

docker run --rm --gpus all nvidia/cuda:12.4.0-base-ubuntu22.04 nvidia-smi

If successful, this prints the GPU info from inside the container.

For Podman (instead of Docker):

sudo nvidia-ctk runtime configure --runtime=podman
sudo systemctl restart podman

Docker Compose example

services:
  train:
    image: pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    volumes:
      - .:/workspace

Monitoring with nvidia-smi

nvidia-smi is the primary CLI tool for GPU monitoring.

Common invocations

# Live dashboard (refreshes every 1 second)
nvidia-smi --query-gpu=index,name,temperature.gpu,utilization.gpu,memory.used,memory.total --format=csv -l 1

# One-shot with key stats
nvidia-smi

# GPU process list
nvidia-smi pmon -s p -d 2

# JSON output for scripting
nvidia-smi --query-gpu=index,name,driver_version,temperature.gpu,utilization.gpu,memory.used --format=csv,noheader,nounits

Persistence Mode

By default, NVIDIA drivers unload from GPU memory when no process is actively using the card. This adds latency on first access. Persistence mode keeps the driver loaded.

# Enable persistence mode
sudo nvidia-smi -pm 1

# Verify
nvidia-smi -q -d PERSISTENCE_MODE | grep "Persistence Mode"

🚨 Important: Persistence mode is reset on reboot. Add it to a startup script or set via systemd:

sudo tee /etc/systemd/system/nvidia-persistence.service > /dev/null << 'EOF'
[Unit]
Description=Enable NVIDIA Persistence Mode
Before=nvidia-fabricmanager.service

[Service]
Type=oneshot
ExecStart=nvidia-smi -pm 1
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
EOF

sudo systemctl enable --now nvidia-persistence.service

DCGM (NVIDIA Data Center GPU Manager)

For datacentre GPUs (A100/H100/B200), DCGM offers richer telemetry:

sudo apt install -y datacenter-gpu-manager
sudo systemctl enable --now nvidia-dcgm
dcgmi discovery -l
dcgmi stats -v

PyTorch / TensorFlow GPU Verification

PyTorch

python3 -c "
import torch
print('PyTorch version:', torch.__version__)
print('CUDA available:', torch.cuda.is_available())
if torch.cuda.is_available():
    print('GPU count:', torch.cuda.device_count())
    for i in range(torch.cuda.device_count()):
        print(f'  [{i}] {torch.cuda.get_device_name(i)}')
    # Quick tensor test
    x = torch.randn(3, 3).cuda()
    print('Tensor on GPU:', x.device)
"

Expected outcome: CUDA available: True and a tensor on cuda:0.

TensorFlow

python3 -c "
import tensorflow as tf
print('TF version:', tf.__version__)
print('GPU devices:', tf.config.list_physical_devices('GPU'))
if tf.config.list_physical_devices('GPU'):
    with tf.device('/GPU:0'):
        a = tf.constant([[1.0, 2.0], [3.0, 4.0]])
        b = tf.constant([[1.0, 0.0], [0.0, 1.0]])
        c = tf.matmul(a, b)
        print('Matrix multiply on GPU:\n', c.numpy())
"

Expected outcome: GPU devices: [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] and a successful matmul.

JAX

python3 -c "
import jax
print('JAX devices:', jax.devices())
print('Default device:', jax.default_device())
"

Troubleshooting

"nvidia-smi: command not found" or "Failed to initialize NVML"

Cause: Driver not installed, kernel module not loaded, or NVML library missing.

# Check kernel module
lsmod | grep nvidia

# Load manually
sudo modprobe nvidia

# Rebuild initramfs (if installed but not loading on boot)
sudo update-initramfs -u
sudo reboot

If still failing, reinstall the driver.

Driver / CUDA version mismatch

Symptom: Framework says CUDA error: no kernel image is available for execution on the device or silently falls back to CPU.

Root cause: The PyTorch/TF wheel was compiled for a different CUDA minor version than what your driver supports.

Fix: Install a driver that supports the required CUDA version, or use the appropriate framework wheel.

# Check max CUDA version your driver supports
nvidia-smi | grep "CUDA Version"

# Match framework wheel to that version
# PyTorch wheels: pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
# (where cu124 = CUDA 12.4)

"libcuda.so.1: cannot open shared object file"

Cause: The CUDA runtime or driver library path is not in LD_LIBRARY_PATH.

Fix:

# Find the library
find /usr/local/cuda -name "libcuda*" 2>/dev/null

# Or use ldconfig
sudo ldconfig -p | grep cuda

# Add to LD_LIBRARY_PATH
export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64:$LD_LIBRARY_PATH

"nvidia-container-cli: initialization error: driver error"

Cause: nvidia-container-toolkit installed but the Docker runtime hook is misconfigured, or the host driver version doesn't match the container's expected CUDA version.

Fix:

# Check current runtime
docker info | grep -i runtime

# Re-run the toolkit configuration
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

# Test with the exact CUDA version matching your driver
docker run --rm --gpus all nvidia/cuda:12.4.0-runtime-ubuntu22.04 nvidia-smi

GPU not detected inside container (but nvidia-smi works on host)

# 1. Ensure nvidia-container-toolkit is installed
dpkg -l | grep nvidia-container-toolkit

# 2. Check the runtime config is present in /etc/docker/daemon.json
cat /etc/docker/daemon.json
# Should contain: "runtimes": {"nvidia": ...}

# 3. Ensure you're passing --gpus all or using deploy.reservations in compose

Thermal throttling / power capping

# Check current power draw and limit
nvidia-smi --query-gpu=power.draw,power.limit --format=csv

# Check thermal throttle reasons
nvidia-smi -q -d TEMPERATURE | grep -i throttle

# For compute workloads, set performance mode to maximum
sudo nvidia-smi -pm 1
sudo nvidia-smi -ac 5001,1590   # (example clock values for A100)

"CUDA out of memory"

Not a driver issue, but common:

# Clear GPU memory (kill zombie processes)
nvidia-smi | grep python | awk '{print $5}' | xargs -I{} kill -9 {}

# Or more aggressively:
sudo fuser -v /dev/nvidia*

To avoid OOM in PyTorch:

torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()

Quick Reference: Version Compatibility Matrix

Driver Branch Max CUDA Toolkit Recommended cuDNN PyTorch Wheel
550.x 12.4 / 12.5 9.3.x cu124
545.x 12.3 9.1.x cu121
535.x 12.2 8.9.x cu118 / cu121
525.x 12.0 8.8.x cu118

⚠️ Always verify framework compatibility at the official sources: - https://pytorch.org/get-started/locally/ - https://www.tensorflow.org/install/source#gpu - https://jax.readthedocs.io/en/latest/installation.html