Configura PyTorch con soporte GPU: CUDA en Linux/Windows y MPS en macOS

Introducción

Si estás trabajando en proyectos de aprendizaje profundo, probablemente necesites aprovechar el poder de las GPUs para acelerar el entrenamiento y la inferencia de tus modelos. PyTorch es una de las bibliotecas más populares para estos tareas, y gracias a su soporte nativo para GPUs, podés ejecutar tu código en CUDA (GPU NVIDIA), en Metal Performance Shaders (MPS) en Mac, o incluso en CPUs si no tenés hardware específico.

Este tutorial te guiará paso a paso por el proceso completo de configuración de PyTorch con soporte GPU en tres plataformas principales: Linux, Windows y macOS. Vas a aprender a instalar los controladores necesarios, el Toolkit de CUDA (cuando corresponda), las bibliotecas auxiliares y, finalmente, a verificar que PyTorch esté correctamente utilizando el hardware acelerado.

Al finalizar, vas a tener un entorno de desarrollo listo para entrenar modelos de forma eficiente, con comandos claros y los resultados esperados que te permitirán confirmar cada paso del proceso.

Para quién es este tutorial

• Desarrolladores y científicos de datos que recién empiezan a trabajar con GPUs.
• Estudiantes de ingeniería o informática que necesitan un setup reproducible en diferentes sistemas operativos.
• Profesionales que ya usan PyTorch en CPU y quieren acelerar sus modelos.

Qué vas a aprender

1. Cómo instalar los controladores NVIDIA (para CUDA) en Linux y Windows.
2. Cómo instalar el Toolkit de CUDA y configurar las variables de entorno.
3. Cómo instalar la versión correcta de PyTorch con soporte CUDA.
4. Cómo verificar la detección de GPU en tiempo de ejecución.
5. Cómo habilitar el backend MPS en macOS y ejecutar modelos en la GPU integrada.
6. Buenas prácticas y errores comunes a la hora de configurar GPU support.

Requisitos previos

Antes de empezar, tenés que tener en cuenta algunos requisitos esenciales:

• Hardware: Una GPU compatible (NVIDIA para CUDA, GPU integrada Apple Silicon para MPS). Para CUDA también necesitas un driver compatible con el Toolkit que vayas a instalar.
• Sistema Operativo: Uno de los tres principales: Linux, Windows o macOS.
• Python: Instalado (recomendado versión 3.10 o superior, pero 3.8+). Podés verificar con python --version o python3 --version.
• Opcional pero útil: git, un editor de código (VS Code, PyCharm, etc.) y una terminal.

Si ya tenés todo eso, podés proceder al setup sin mayores complicaciones.

Verificación rápida del sistema

Antes de instalar nada, es buena idea chequear el estado actual del sistema y asegurarte de que hay espacio suficiente en disco.

En Linux y macOS

python3 -c "import sys; print('Python', sys.version)"

Salida esperada (ejemplo):

Python 3.10.12 (default, Sep 11 2023, 14:23:45) 
[Clang 14.0.0 (clang-1400.0.29.202)]

En Windows (usando PowerShell o Command Prompt)

python --version

Salida esperada: Python 3.10.12...

A continuación, usamos nvidia-smi (Linux) o comprobamos la presencia de system_profiler SPGPUDataType (macOS) para ver si hay una GPU NVIDIA o Apple detectada. Si no las tenés, todavía podés seguir el tutorial; solo las partes que correspondan.

CUDA en Linux: Instalación paso a paso

Linux es la plataforma más usada para entrenamiento en GPU. El proceso generalmente se divide en tres etapas: instalar el driver, instalar el CUDA Toolkit y luego instalar PyTorch con soporte CUDA.

1. Instalar el driver NVIDIA

La mayoría de las distribuciones modernas de Linux permiten instalar los drivers NVIDIA a través de repositorios oficiales. El siguiente comando funciona para Ubuntu/Debian (ajusta según tu distro):

sudo apt update && sudo apt install -y nvidia-driver-535

Salida esperada (resumen):

Leyendo la lista de paquetes... Hecho
Construyendo árbol de dependencia... Hecho
Obteniendo paquetes... Hecho
... Instalando nvidia-driver-535... Hecho

Tras la instalación, reiniciás el sistema para que el kernel cargue el nuevo driver.

2. Verificar que el driver esté activo

nvidia-smi

Salida esperada (ejemplo):

+-----------------------------------------------------------------------------+ 
| NVIDIA-SMI 535.104.05   Driver Version: 535.104.05   CUDA Version: 12.2     | 
|-------------------------------+----------------------+----------------------+ 
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC | 
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. | 
|                               |                      |                      | 
| 0  NVIDIA-A10G    On   | 00000000:3a:00.0 Off |                  N/A | 
| 40%   55C    P8    250W / 300W |  22G / 24G |     0%      Default | 
|                               |                      |                      | 
+-----------------------------------------------------------------------------+ 
| CPU  Usage: 5%   Mem-Usage: 2.5G / 8.0G   GPU-Util: 0%   Enc. Util: N/A | 
|=============================================================================+| 
   Versión CUDA: 12.2

Si ves una tabla como esta, el driver funciona correctamente.

3. Instalar el CUDA Toolkit

Elegí una versión compatible con tu driver (por ejemplo, CUDA 12.2). Para Ubuntu, podés agregar el repositorio CUDA de NVIDIA:

# Descargar el script de configuración
cuda-linux.sh

sudo bash cuda-linux.sh --toolkit

Salida esperada (parte del script):

Leyendo lista de paquetes... Hecho
Construyendo árbol de dependencia... Hecho
Obteniendo paquetes... Hecho
... Instalando cuda-toolkit-12-2... Hecho

El script también agrega automáticamente el PATH necesario, pero a veces es más sencillo agregarlo manualmente:

# Añadir CUDA al PATH para el usuario actual
echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

4. Verificar la instalación de CUDA

nvcc --version

Salida esperada (ejemplo):

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Construido sobre LLVM 14.0.0, PTX compiler, NVTX, cubinlinker
Cuda feature level: 80
CUDA Version: 12.2

5. Instalar PyTorch con soporte CUDA

Ahora que el Toolkit está listo, instalamos PyTorch con la versión GPU usando pip. La sintaxis típica para una máquina con CUDA 12.2 es:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu122

Salida esperada (resumen de pip):

Colectando paquetes de https://download.pytorch.org/whl/cu122
  .../torch-2.2.0%2Bcu122.whl (?? MB)
Instalando paquetes: ...
... Se han instalado 15 paquetes:
    torch==2.2.0+cu122
    torchvision==0.17.0+cu122
    torchaudio==2.2.0+cu122
... Hecho

6. Verificar detección de GPU en tiempo de ejecución

Ahora podés ejecutar un script rápido en Python para confirmar que PyTorch puede ver la GPU.

try:
    import torch
    print(f'PyTorch version: {torch.__version__}')
    print(f'CUDA disponible: {torch.cuda.is_available()}')
    if torch.cuda.is_available():
        print(f'Número de GPUs: {torch.cuda.device_count()}')
        print(f'GPU actual: {torch.cuda.get_device_name(0)}')
        print(f'Versión CUDA del driver: {torch.version.cuda}')
except Exception as e:
    print('Error:', e)

Salida esperada (con una GPU NVIDIA presente):

PyTorch version: 2.2.0+cu122
CUDA disponible: True
Número de GPUs: 1
GPU actual: NVIDIA-A10G
Versión CUDA del driver: 12.2

Si ves CUDA disponible: True, estás listo para usar PyTorch en GPU.

CUDA en Windows: Instalación paso a paso

Windows sigue siendo muy usado en investigación y prototipado rápido. Los pasos son similares, pero los comandos y las rutas de archivos difieren.

1. Instalar el driver NVIDIA

Descargá el driver desde el sitio oficial de NVIDIA (ej: NVIDIA GeForce Game Ready Driver 535.104.05). Elegí la opción que se ajuste a tu GPU y versión de Windows (Windows 10/11 de 64 bits). Durante la instalación, asegurate de elegir la opción de instalación personalizada y seleccioná “Agregar atajos de comandos para CUDA” (si está presente).

2. Verificar que el driver esté activo

Abre un Símbolo del sistema (o PowerShell) como administrador y ejecuta:

nvidia-smi

Salida esperada (ejemplo):

NVIDIA-SMI 535.104.05   Driver Version: 535.104.05   CUDA Version: 12.2
===============================================================================+ 
GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC | 
Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. | 
|                               |                      |                      | 
|  0  NVIDIA-A10G    On   | 00000000:3A:00.0 Off |                  N/A | 
| 45%   58C    P8    300W / 300W |  22G / 24G |     0%      Default |
|                               |                      |                      | 
===============================================================================+ 
CPU  Usage: 4%   Mem-Usage: 2.5G / 8.0G   GPU-Util: 0%   Enc. Util: N/A
===============================================================================+ 
Versión CUDA: 12.2

3. Instalar el CUDA Toolkit en Windows

El instalador de CUDA para Windows está disponible en https://developer.nvidia.com/cuda-downloads. Elegí Windows 10/11, 64-bit, CUDA 12.2 Toolkit. Durante la instalación, aceptá el acuerdo de licencia y elegí Custom (advanced).

Ubicaciones típicas:

• Toolkit: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.2
• Binarios: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.2\bin
• Librerías: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.2\lib\x64

El instalador automáticamente agrega C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.2\bin al PATH de sistema.

4. Verificar la instalación de CUDA

Abre un nuevo Símbolo del sistema (para asegurar el PATH actualizado) y ejecuta:

nvcc --version

Salida esperada (ejemplo):

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on LLVM 14.0.0, PTX compiler, NVTX, cubinlinker
Cuda feature level: 80
CUDA Version: 12.2

5. Instalar PyTorch con soporte CUDA

Ahora usás pip (o py -m pip si tenés múltiples entornos Python) para instalar la versión GPU de PyTorch:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu122

Salida esperada (resumen):

Colectando paquetes de https://download.pytorch.org/whl/cu122
.../torch-2.2.0%2Bcu122.whl (?? MB)
Instalando paquetes: ...
... Se han instalado 15 paquetes:
    torch==2.2.0+cu122
    torchvision==0.17.0+cu122
    torchaudio==2.2.0+cu122
... Hecho

6. Verificar detección de GPU en tiempo de ejecución (Windows)

Ejecutá el mismo script Python (asegurándote de estar dentro del entorno donde instalaste PyTorch):

try:
    import torch
    print(f'PyTorch version: {torch.__version__}')
    print(f'CUDA disponible: {torch.cuda.is_available()}')
    if torch.cuda.is_available():
        print(f'Número de GPUs: {torch.cuda.device_count()}')
        print(f'GPU actual: {torch.cuda.get_device_name(0)}')
        print(f'Versión CUDA del driver: {torch.version.cuda}')
except Exception as e:
    print('Error:', e)

Salida esperada (con CUDA):

PyTorch version: 2.2.0+cu122
CUDA disponible: True
Número de GPUs: 1
GPU actual: NVIDIA-A10G
Versión CUDA del driver: 12.2

Si ves esto, estás listo para entrenar en GPU en Windows.

Configurar MPS (Metal Performance Shaders) en macOS

Apple Silicon (M1, M2, etc.) trae una GPU integrada que se puede utilizar mediante el backend MPS de PyTorch. A diferencia de CUDA, MPS no requiere instalación de drivers externos, pero sí necesitas el comando pip adecuado y una versión reciente de macOS.

1. Requisitos previos del sistema

• macOS 12 (Monterey) o superior.
• Procesador Apple Silicon (M1, M2, M3, etc.).
• Homebrew instalado (opcional pero útil). Podés instalar con:

/bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"

2. Instalar PyTorch con soporte MPS

La forma más sencilla es usar pip desde el índice de PyPI oficial, que por defecto trae binaries compatibles con MPS para macOS. Ejecutá:

pip3 install torch torchvision torchaudio

Salida esperada (resumen):

Colectando paquetes de https://pypi.org/simple
  .../torch-2.2.0-cp311-cp311-darwin_x86_64.whl (?? MB)
Instalando paquetes: ...
... Se han instalado 15 paquetes:
    torch==2.2.0
    torchvision==0.17.0
    torchaudio==2.2.0
... Hecho

3. Habilitar el backend MPS explícitamente

A partir de PyTorch 2.2, el backend MPS está disponible automáticamente, pero podés activarlo así:

import torch
if torch.backends.mps.is_available():
    device = torch.device('mps')
    print('Dispositivo MPS disponible:', device)
else:
    print('MPS no disponible, usando CPU')

4. Verificar MPS en tiempo de ejecución

Ejecutamos el siguiente script (con run:python):

import torch
print('PyTorch version:', torch.__version__)
print('CUDA disponible:', torch.cuda.is_available())
print('MPS disponible:', torch.backends.mps.is_available())

if torch.backends.mps.is_available():
    x = torch.randn(2, 2, device='mps')
    y = torch.randn(2, 2, device='mps')
    z = x @ y
    print('Resultado de la multiplicación en MPS:', z)

Salida esperada (con Apple Silicon):

PyTorch version: 2.2.0
CUDA disponible: False
MPS disponible: True
Resultado de la multiplicación en MPS: tensor([[ 0.5713, -1.4771],
        [ 0.3212,  0.9345]], device='mps')

Si ves MPS disponible: True, el backend funciona y podés mover tensores a 'mps' para acelerar cálculos.

Buenas prácticas y errores comunes

Consejos para un mejor rendimiento

• Usa torch.cuda.amp.autocast() para half-precision training en GPU; reduce el uso de memoria y acelera cálculos.
• Preasigna memoria con torch.cuda.empty_cache() y torch.cuda.synchronize() si ves picos de memoria.
• Ejecutá múltiples procesos en paralelo con torch.multiprocessing y spawn para aprovechar varias GPUs (torch.cuda.set_device(i)).
• Versiona tus dependencias: guardá un requirements.txt con torch==2.2.0+cu122 (Linux/Windows) o torch==2.2.0 (macOS) para evitar sorpresas.

Conclusión

Configurar PyTorch con soporte GPU no tiene que ser una tarea complicada. Ya sea que estés en un servidor Linux con varias GPUs NVIDIA, una workstation Windows para prototipado rápido, o una notebook Apple Silicon para cálculos móviles, cada plataforma tiene un camino claro:

• Linux/Windows: Instalar driver NVIDIA → Instalar CUDA Toolkit → Verificar nvidia-smi/nvcc → Instalar PyTorch con --index-url https://download.pytorch.org/whl/cu122 (o la versión que coincida con tu CUDA).
• macOS: Instalar PyTorch con soporte MPS (no se necesitan drivers externos) → Verificar torch.backends.mps.is_available().

Al seguir estos pasos, deberías ver CUDA disponible: True (Linux/Windows) o MPS disponible: True (macOS) en cada script de verificación. Además, conocés los errores más frecuentes y las mejores prácticas para aprovechar al máximo la GPU.

Ahora podés entrenar modelos más rápido, ejecutar inferencias en tiempo real y manejar datasets de mayor tamaño sin preocuparte por el cuello de botella de la CPU. ¡Feliz programación con PyTorch!

Próximos pasos

• Optimización de modelos: Usá torch.compile (o torchdynamo) para aceleración JIT.
• Distribución: Escala a múltiples GPUs con torch.nn.parallel.DistributedDataParallel.
• Monitorización: Implementá tensorboardX o wandb para graficar métricas de entrenamiento.
• Actualizaciones: Mantené el Toolkit de CUDA y los drivers actualizados para aprovechar nuevas instrucciones y correcciones de seguridad.

¡El futuro de la investigación en IA está en la GPU, y PyTorch te da las herramientas para alcanzarlo! 🚀