Gpu Containers

GPU Containers

A comprehensive guide to running GPU-accelerated workloads in containers for research computing. This skill covers NVIDIA CUDA and AMD ROCm container setup, driver compatibility, base image selection, framework-specific images for PyTorch, TensorFlow, and JAX, multi-GPU configuration, and Docker Compose integration. GPU containers are essential for machine learning training, scientific simulation, and any compute-intensive workload that benefits from hardware acceleration.

Resources in This Skill

This skill includes supporting materials for GPU container tasks:

References (detailed guides -- consult the table of contents in each file and read specific sections as needed):

• references/nvidia-cuda-setup.md - NVIDIA CUDA setup: nvidia-container-toolkit installation, CUDA version matrix, base image variants (base/runtime/devel), cuDNN and NCCL integration, multi-GPU with CUDA_VISIBLE_DEVICES, profiling in containers, and driver compatibility
• references/amd-rocm-setup.md - AMD ROCm setup: ROCm installation, amdgpu driver, ROCm Docker images, HIP compatibility layer, device mapping, and ROCm version matrix

Assets (ready-to-use Dockerfile templates):

• assets/cuda-dockerfile.dockerfile - Multi-stage CUDA Dockerfile with PyTorch, optimized layer ordering, and commented CUDA/cuDNN choices
• assets/rocm-dockerfile.dockerfile - ROCm Dockerfile with PyTorch, device mapping instructions, and HIP compatibility notes

Quick Reference Card

CUDA Version Compatibility

Quick GPU Check

# Verify GPU is visible from host
nvidia-smi

# Run a container with GPU access
docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi

# Run with specific GPUs
docker run --rm --gpus '"device=0,1"' nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi

# Run with AMD GPU
docker run --rm --device=/dev/kfd --device=/dev/dri \
  rocm/pytorch:latest rocm-smi

NVIDIA Container Toolkit Install (Quick)

# Ubuntu/Debian
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
  sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.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-get update && sudo apt-get install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

When to Use

Use this skill when you need to:

• Run NVIDIA GPU workloads (CUDA) inside Docker containers
• Run AMD GPU workloads (ROCm) inside Docker containers
• Choose the right CUDA base image variant (base, runtime, or devel)
• Build Dockerfiles for PyTorch, TensorFlow, or JAX with GPU support
• Configure multi-GPU access for distributed training
• Set up Docker Compose services with GPU resources
• Debug GPU access issues inside containers (driver mismatch, device not found)
• Profile GPU workloads inside containers
• Run GPU containers on HPC clusters (cross-reference the singularity-apptainer skill)

NVIDIA CUDA Setup

nvidia-container-toolkit

The NVIDIA Container Toolkit is the bridge between Docker and your NVIDIA GPUs. It must be installed on the host (not in the container).

What it provides:

• --gpus flag for docker run
• Automatic mounting of GPU drivers into containers
• CUDA library injection so containers do not need to bundle host drivers

Verify installation:

# Check the toolkit is configured
nvidia-ctk --version

# Test GPU access in a container
docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi

If nvidia-smi inside the container shows your GPUs, the toolkit is working.

See references/nvidia-cuda-setup.md for full installation instructions and troubleshooting.

Driver Compatibility

The host NVIDIA driver must be compatible with the CUDA version used in the container. The NVIDIA Container Toolkit handles forward compatibility (newer driver with older CUDA), but the driver must meet the minimum version for the CUDA toolkit in the container.

Rule of thumb: Keep host drivers up to date. A driver from the 550+ series supports all CUDA 12.x containers.

# Check host driver version
nvidia-smi --query-gpu=driver_version --format=csv,noheader

AMD ROCm Setup

amdgpu Driver and ROCm Runtime

AMD GPU containers require the amdgpu kernel driver on the host and pass devices via --device flags (not --gpus).

# Verify AMD GPU on host
rocm-smi

# Run a container with AMD GPU
docker run --rm \
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  rocm/pytorch:latest \
  python -c "import torch; print(torch.cuda.is_available())"

Required device mappings:

• /dev/kfd -- Kernel Fusion Driver (compute)
• /dev/dri -- Direct Rendering Infrastructure (GPU access)

See references/amd-rocm-setup.md for installation and version compatibility.

Base Image Selection

NVIDIA CUDA Image Variants

NVIDIA provides three variants of each CUDA version:

# For inference (no compilation needed)
FROM nvidia/cuda:12.4.1-runtime-ubuntu22.04

# For training with cuDNN (most ML workloads)
FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04

# For compiling custom CUDA extensions
FROM nvidia/cuda:12.4.1-devel-ubuntu22.04

Multi-stage pattern: Use devel to compile, then copy artifacts to runtime:

FROM nvidia/cuda:12.4.1-devel-ubuntu22.04 AS builder
# ... compile CUDA code ...

FROM nvidia/cuda:12.4.1-runtime-ubuntu22.04
COPY --from=builder /app/build/myapp /usr/local/bin/myapp

AMD ROCm Images

# ROCm base
FROM rocm/dev-ubuntu-22.04:6.1

# ROCm with PyTorch pre-installed
FROM rocm/pytorch:latest

# ROCm with TensorFlow pre-installed
FROM rocm/tensorflow:latest

Framework Images

PyTorch

NVIDIA (NGC container):

# Pre-built PyTorch with CUDA, cuDNN, NCCL
FROM nvcr.io/nvidia/pytorch:24.04-py3

Custom build with CUDA:

FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

AMD ROCm:

FROM rocm/pytorch:latest
# Or install manually
FROM rocm/dev-ubuntu-22.04:6.1
RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.1

See assets/cuda-dockerfile.dockerfile and assets/rocm-dockerfile.dockerfile for complete templates.

TensorFlow

NVIDIA:

# Official TensorFlow GPU image
FROM tensorflow/tensorflow:2.16.1-gpu

# Or NGC container with optimizations
FROM nvcr.io/nvidia/tensorflow:24.04-tf2-py3

AMD ROCm:

FROM rocm/tensorflow:latest

JAX

NVIDIA:

FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
RUN pip install --upgrade "jax[cuda12]"

AMD ROCm:

FROM rocm/dev-ubuntu-22.04:6.1
RUN pip install --upgrade "jax[rocm]" -f https://storage.googleapis.com/jax-releases/rocm/jaxlib-latest.html

Multi-GPU Configuration

Selecting GPUs

Control which GPUs a container can access:

# All GPUs
docker run --gpus all myapp

# Specific GPUs by index
docker run --gpus '"device=0,2"' myapp

# Specific GPUs by UUID
docker run --gpus '"device=GPU-abc123,GPU-def456"' myapp

CUDA_VISIBLE_DEVICES

Inside the container, use CUDA_VISIBLE_DEVICES to further restrict GPU visibility for multi-process workloads:

# In Docker run
docker run --gpus all -e CUDA_VISIBLE_DEVICES=0,1 myapp

# In Dockerfile
ENV CUDA_VISIBLE_DEVICES=0,1

Distributed Training

For multi-GPU training with PyTorch DDP or Horovod, all GPUs must be visible to the container:

docker run --gpus all \
  --ipc=host \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  myapp torchrun --nproc_per_node=4 train.py

Key flags for distributed training:

• --ipc=host -- shared memory for inter-process communication (required for PyTorch DataLoader with num_workers > 0)
• --ulimit memlock=-1 -- allow unlimited locked memory (required for NCCL)
• --ulimit stack=67108864 -- increase stack size for deep models
• --shm-size=8g -- alternative to --ipc=host for controlling shared memory size

Docker Compose with GPU

NVIDIA GPU in Compose

services:
  training:
    build: .
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all        # or a specific number: 1, 2
              capabilities: [gpu]
    volumes:
      - ./data:/app/data
    ipc: host
    ulimits:
      memlock: -1
      stack: 67108864

Specific GPUs in Compose

services:
  model-a:
    build: .
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ['0']
              capabilities: [gpu]

  model-b:
    build: .
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ['1']
              capabilities: [gpu]

AMD GPU in Compose

services:
  training:
    build: .
    devices:
      - /dev/kfd:/dev/kfd
      - /dev/dri:/dev/dri
    group_add:
      - video

Testing GPU Access

Verification Commands

# NVIDIA: check GPU visibility
docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi

# NVIDIA: check CUDA version
docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 \
  nvcc --version 2>/dev/null || echo "nvcc not in base image, use devel variant"

# PyTorch GPU check
docker run --rm --gpus all my-pytorch-image \
  python -c "import torch; print(f'CUDA available: {torch.cuda.is_available()}, Devices: {torch.cuda.device_count()}')"

# TensorFlow GPU check
docker run --rm --gpus all my-tf-image \
  python -c "import tensorflow as tf; print(f'GPUs: {tf.config.list_physical_devices(\"GPU\")}')"

# AMD: check ROCm
docker run --rm --device=/dev/kfd --device=/dev/dri rocm/pytorch:latest rocm-smi

Common Failures and Fixes

HPC GPU Containers

For running GPU containers on HPC clusters with Singularity/Apptainer, GPU passthrough works differently than Docker. Singularity automatically binds NVIDIA drivers from the host.

# Convert Docker image to Singularity
singularity pull docker://nvcr.io/nvidia/pytorch:24.04-py3

# Run with GPU
singularity run --nv pytorch_24.04-py3.sif python train.py

# Apptainer (Singularity successor)
apptainer run --nv pytorch_24.04-py3.sif python train.py

See the singularity-apptainer skill for complete HPC container guidance.

Common Mistakes

1. Installing NVIDIA drivers inside the container -- The nvidia-container-toolkit injects host drivers at runtime. Installing drivers in the container causes conflicts.

2. Using the wrong CUDA base image variant -- base does not include cuDNN. Use runtime for inference or devel for compiling CUDA extensions.

3. CUDA version mismatch with host driver -- The container's CUDA version must be compatible with the host's NVIDIA driver. Check the compatibility matrix.

4. Forgetting --ipc=host for training -- PyTorch DataLoader with multiple workers requires shared memory. Without --ipc=host or --shm-size, training crashes with bus errors.

5. Not pinning CUDA versions -- Using nvidia/cuda:latest can break builds when CUDA updates. Pin to specific versions like 12.4.1.

6. Using devel images in production -- The devel variant includes compilers and headers (~3.5 GB). Use multi-stage builds to compile in devel and run in runtime.

7. Ignoring GPU memory limits -- Unlike CPU containers, GPU memory is not limited by Docker. A single container can consume all GPU memory, starving other containers.

8. AMD: forgetting device flags -- AMD GPUs require --device=/dev/kfd --device=/dev/dri and --group-add video. There is no --gpus equivalent for ROCm.

Best Practices

• Install nvidia-container-toolkit on the host, never install drivers in containers
• Pin CUDA and cuDNN versions in your Dockerfile for reproducibility
• Use multi-stage builds: compile in devel, run in runtime
• Verify GPU access with nvidia-smi and framework checks before running workloads
• Use --ipc=host or --shm-size=8g for PyTorch training workloads
• Set CUDA_VISIBLE_DEVICES to control GPU allocation in multi-GPU systems
• Use NGC containers (nvcr.io/nvidia/pytorch) for optimized framework builds
• Test on the same CUDA version locally that your production/HPC cluster runs
• Document minimum driver version requirements in your README
• For AMD GPUs, always map both /dev/kfd and /dev/dri devices
• Use Docker Compose deploy.resources.reservations.devices for declarative GPU access
• Monitor GPU memory usage with nvidia-smi or nvitop during development

Resources

Official Documentation

• NVIDIA Container Toolkit: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/
• NVIDIA CUDA Docker Hub: https://hub.docker.com/r/nvidia/cuda
• NVIDIA NGC Catalog: https://catalog.ngc.nvidia.com/
• AMD ROCm Docker: https://rocm.docs.amd.com/projects/install-on-linux/en/latest/how-to/docker.html

Framework Images

• PyTorch NGC: https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch
• TensorFlow Docker: https://www.tensorflow.org/install/docker
• JAX Installation: https://jax.readthedocs.io/en/latest/installation.html

Compatibility

• CUDA Compatibility Matrix: https://docs.nvidia.com/deploy/cuda-compatibility/
• cuDNN Support Matrix: https://docs.nvidia.com/deeplearning/cudnn/support-matrix/
• ROCm Compatibility: https://rocm.docs.amd.com/en/latest/compatibility/compatibility-matrix.html

Tools

• nvitop (GPU monitoring): https://github.com/XuehaiPan/nvitop
• gpustat: https://github.com/wookayin/gpustat