Docker GPU containers with the NVIDIA Container Toolkit

Running Docker GPU workloads with the NVIDIA Container Toolkit takes one flag, but real setups need a little more: choosing GPUs per container, reserving them in Docker Compose, picking a base image that matches the driver and giving PyTorch enough shared memory. Docker and the toolkit are preinstalled on our Ubuntu images. AMD servers use device flags instead, covered near the end.

Check Docker and the toolkit

docker --version
nvidia-ctk --version
docker info | grep -i runtimes
docker run --rm --gpus all ubuntu nvidia-smi

The last command should print the same GPU table as on the host. The toolkit mounts the host's driver into the container, so images never include a driver; they only need the CUDA libraries your application uses. As a non-root user, prefix docker with sudo, or join the docker group, which is equivalent to root access.

Give a container all GPUs or only some

# all GPUs
docker run --rm --gpus all ubuntu nvidia-smi -L
# GPUs 0 and 1 (the inner quotes are required because of the comma)
docker run --rm --gpus '"device=0,1"' ubuntu nvidia-smi -L
# one GPU by UUID, as printed by nvidia-smi -L on the host
docker run --rm --gpus device=GPU-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx ubuntu nvidia-smi -L

Inside the container, the selected GPUs are numbered from 0, so host GPUs 2 and 3 appear as 0 and 1. UUIDs are the safest way to pin a service to a specific GPU. To run eight independent jobs on an 8-GPU server, start one container per GPU with --gpus device=N.

The toolkit also reads the NVIDIA_VISIBLE_DEVICES variable when a container runs with the NVIDIA runtime:

docker run --rm --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=0,1 ubuntu nvidia-smi -L

If nvidia is the default runtime in /etc/docker/daemon.json, images that set this variable, as NVIDIA's CUDA images do, get every GPU even without --gpus, so set it explicitly when you want a limit. CUDA_VISIBLE_DEVICES is different: it filters, inside the container, among the GPUs already mounted.

Reserve GPUs in Docker Compose

services:
  trainer:
    image: ubuntu:24.04
    command: nvidia-smi -L
    ipc: host
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0", "1"]
              capabilities: [gpu]

Run it with docker compose up; the output should list two GPUs. To take GPUs by number rather than by ID, replace device_ids with count: all or count: 2; the two keys cannot be combined, and capabilities: [gpu] is required. If docker compose version fails, install the Compose plugin: docker-compose-plugin when Docker comes from Docker's repository, docker-compose-v2 when it comes from Ubuntu's.

Pick a CUDA or NGC base image

• nvidia/cuda images on Docker Hub come in three flavors: base (minimal CUDA runtime), runtime (adds the CUDA libraries) and devel (adds the compiler and headers), some with cuDNN. Tags follow VERSION-FLAVOR-OS, for example 12.8.1-runtime-ubuntu24.04.
• NGC images such as nvcr.io/nvidia/pytorch bundle a tuned framework stack, with monthly tags in the form YY.MM-py3. Their release notes list the driver each tag needs.

Pick a tag whose CUDA version is at or below the CUDA Version shown by nvidia-smi on the host, and CUDA 12.8 or later for Blackwell GPUs. If an image needs a newer driver than the host has, the container refuses to start with an unsatisfied condition: cuda>= error. Driver and CUDA versions explains the rule.

Give PyTorch enough shared memory

Docker limits /dev/shm to 64 MB by default. PyTorch DataLoader workers and NCCL rely on shared memory, so training fails with a "bus error" or an "insufficient shared memory (shm)" message. Share the host's IPC namespace or raise the limit (replace YY.MM with a current tag):

# share the host IPC namespace, as NVIDIA recommends for its PyTorch images
docker run --rm -it --gpus all --ipc=host \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /data:/data nvcr.io/nvidia/pytorch:YY.MM-py3

# or keep IPC isolated with a larger /dev/shm
docker run --rm -it --gpus all --shm-size=16g \
  -v /data:/data nvcr.io/nvidia/pytorch:YY.MM-py3

Inside the container, df -h /dev/shm shows the space available. In Compose, use ipc: host or shm_size: "16gb". The vLLM image needs the same treatment, as shown in serving an LLM with vLLM.

AMD GPUs: the ROCm equivalent

docker run --rm --device=/dev/kfd --device=/dev/dri --group-add video --ipc=host \
  rocm/pytorch:latest python3 -c "import torch; print(torch.cuda.device_count())"

On an MI355X server this prints 8. The standard way to run ROCm containers passes device files rather than --gpus: /dev/kfd is the compute interface and /dev/dri holds the GPU device nodes. To limit a container to some GPUs, add -e HIP_VISIBLE_DEVICES=0,1. See ROCm setup for the rest.

Troubleshooting

• could not select device driver "" with capabilities: [[gpu]], or unknown or invalid runtime name: nvidia. Docker is not set up for the NVIDIA runtime. Run sudo nvidia-ctk runtime configure --runtime=docker, then sudo systemctl restart docker.
• unsatisfied condition: cuda>=X.Y. The image needs a newer driver. Use an older image tag, or ask us about the driver.
• permission denied while trying to connect to the Docker daemon socket. Use sudo, or add your user to the docker group and log in again.