Cómo activar GPU para poder acceder desde un contenedor Docker por ejemplo para un modelo LLM

Tiempo de lectura: 2 minutos

Activar el acceso a la GPU es imprescindible si necesitamos arrancar un modelo LLM y utilizar GPU con VRAM.

Requisitos previos

Asegúrate de tener una GPU NVIDIA en tu máquina con los drivers instalados y actualizados.
Verifica con:

nvidia-smi

Deberías ver tu GPU y versión de driver.
También necesitas Docker instalado.

Instalar NVIDIA Container Toolkit

Esto permite que Docker detecte y use la GPU.

En Ubuntu o Debian:

distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

sudo apt-get update
sudo apt-get install -y nvidia-docker2

sudo systemctl restart docker

Verifica que funcione:

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

Deberías ver tu GPU listada dentro del contenedor.

Ejecutar contenedor con GPU

Usando Docker CLI:

docker run --gpus all -it --rm \
  vllm/vllm-openai:latest \
  --model meta-llama/Meta-Llama-3-8B-Instruct

--gpus all da acceso a todas las GPUs disponibles.
vLLM detectará la GPU automáticamente y la usará.

Usando Docker Compose

Ejemplo mínimo docker-compose.yml:

version: "3.9"
services:
  vllm:
    image: vllm/vllm-openai:latest
    container_name: vllm_gpu
    ports:
      - "8000:8000"
    volumes:
      - ./models:/model
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    command: >
      --model /model
      --host 0.0.0.0
      --port 8000

Ejecuta:

docker compose up -d

Verificar GPU dentro del contenedor

docker exec -it vllm_gpu bash
python -c "import torch; print(torch.cuda.is_available())"

Si devuelve True, la GPU está activa.
También puedes usar:

nvidia-smi

Consejos para producción

Para usar varias GPUs:

--gpus '"device=0,1"'

Ajusta la memoria GPU en vLLM:

--gpu-memory-utilization 0.9

Montar modelos locales en /model evita descargar cada vez:

volumes:
  - ./models:/model