Compatibilidad entre versiones de PyTorch, CUDA y xFormers

Aprende cómo solucionar conflictos de dependencias en software de Inteligencia Artificial.

Introducción

Para ejecutar modelos generativos en nuestra máquina local recurrimos a software que nos facilita el proceso de inferencia.

Para los modelos generativos de texto como Falcon, Llama 2 o modelos de la familia Llama (Vicuna o Alpaca), solemos utilizar la interfaz Oobabooga text-generation-webui o FastChat.

En el caso de los modelos generativos de imagen, o más bien, en el caso de Stable Diffusion, las interfaces más populares son Stable Diffusion web UI de Automatic1111, Invoke AI o ComfyUI. Para el entrenamiento o fine-tuning de estos modelos se hace uso de aplicaciones/extensiones como Dreambooth o Kohya SS.

Todas estas aplicaciones tienen una pieza en común. Utilizan PyTorch, una librería de machine learning escrita en Python. A su vez, PyTorch es dependencia para otras librerías utilizadas en algunas de estas aplicaciones (como xFormers). Y por si no fuera suficiente, para sacar el máximo rendimiento utilizando la GPU en vez de la CPU, necesitaremos librerías como CUDA en el caso de tarjetas gráficas NVIDIA.

Todas estas librerías necesitan entenderse entre sí y no causar problemas entre versiones no compatibles, algo difícil hasta que le pillamos el truco. En este artículo veremos como instalar las dependencias adecuadas para todas las librerías que necesitemos y solucionar de una vez por todas los conflictos.

Cuando ejecutes comandos de instalación mediante pip, recuerda estar en un virtual environment (venv) para que las librerías no se instalen de manera global. Por supuesto, también necesitas tener instalado Python 3.8 o superior y python-pip.

Elegir versión de PyTorch

Lo principal es seleccionar la versión de PyTorch que necesitamos ya que esta elección condicionará a todas las demás librerías. La versión depende de la aplicación que utilicemos. Por ejemplo, en el caso de Stable Diffusion web UI de Automatic1111, la última versión utiliza PyTorch 2.0.1. Por otro lado, la última versión de InvokeAI requiere la versión 2.1.0 de PyTorch. Podemos encontrar esta información en el fichero requirements.txt del proyecto o en su documentación.

De momento no instalaremos nada, solo anotamos la versión de PyTorch que utilizaremos. Para este artículo utilizaremos la versión 2.0.1.

CUDA

Si vamos a utilizar nuestra tarjeta gráfica lo primero es instalar CUDA toolkit, una plataforma creada por NVIDIA para el procesamiento gráfico en paralelo. Si no tenemos una gráfica NVIDIA o queremos utilizar la CPU (aunque es mucho más lento), podemos saltarnos este paso. Además, es necesario tener instalados los drivers de NVIDIA.

Para saber qué versión de CUDA es compatible con una versión específica de PyTorch, acudimos a la página web de PyTorch y encontraremos una tabla. Si la versión que necesitamos es la versión estable actual, la seleccionamos y nos fijamos más abajo en la línea Compute Platform. Aquí aparecen las versiones de CUDA compatibles con esta versión.

Tabla mostrando las versiones de CUDA compatibles con PyTorch 2.1.1: CUDA 11.8 y CUDA 12.1
PyTorch 2.1.1 es compatible con CUDA 11.8 y CUDA 12.1

Ahora mismo la versión estable oficial es 2.1.1 por lo que esta no es la versión que queremos. Para encontrar otras versiones accedemos a https://pytorch.org/get-started/previous-versions/.

En nuestro caso buscamos la sección llamada v2.0.1, después la subsección Wheel y dentro de Linux and Windows encontraremos esto:

# ROCM 5.4.2 (Linux only)
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/rocm5.4.2
# CUDA 11.7
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2
# CUDA 11.8
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118
# CPU only
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cpu

Aquí podemos ver como esta versión es compatible con CUDA 11.7 y 11.8. Vamos a elegir CUDA 11.8. Reservamos el comando de instalación de PyTorch para más tarde, primero vamos a instalar CUDA.

Si utilizamos Windows o distribuciones Linux basadas en Debian (Ubuntu) o Red Hat (Fedora o CentOS), podemos acudir a la página de descargas oficial, seleccionar la versión de CUDA deseada y seguir las instrucciones de instalación que nos aparecerán.

En el caso de Arch Linux, descargaremos el fichero .tar.* de la versión que necesitemos, a través del archivo de paquetes de archlinux: https://archive.archlinux.org/packages/c/cuda/.

Para instalar una versión especifica hay que ejecutar el comando sudo pacman -U PACKAGE.tar.xz, ya que pacman -S cuda siempre instalará la versión más moderna disponible.

Además, es imprescindible añadir cuda a la variable IgnorePkg dentro del fichero /etc/pacman.conf (IgnorePkg = cuda), para así así evitar que este paquete se actualice cuando actualicemos el sistema.

Por lo tanto, como la última versión de CUDA 11.8 disponible en Arch Linux es cuda-11.8.0-1-x86_64.pkg.tar.zst, descargamos este fichero y ejecutamos sudo pacman -U cuda-11.8.0-1-x86_64.pkg.tar.zst.

Comprobación

En sistemas Linux, para comprobar que CUDA está instalado correctamente, mucha gente cree que se utiliza el comando nvidia-smi. Esto no es correcto ya que este comando muestra la versión más moderna de CUDA que es compatible con los drivers instalados. Se mostrará una versión incluso si no tenemos instalado CUDA toolkit.

nvidia-smi
Mon Nov 20 21:21:13 2023
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 545.29.02    Driver Version: 545.29.02    CUDA Version: 12.3     |
|-------------------------------+----------------------+----------------------+
...

El método correcto para comprobar la versión de CUDA instalada es utilizar el comando nvcc --version. Obtendremos una salida como ésta:

nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Sep_21_10:33:58_PDT_2022
Cuda compilation tools, release 11.8, V11.8.89
Build cuda_11.8.r11.8/compiler.31833905_0

En algunas distribuciones Linux el comando nvcc --version detecta la versión instalada de CUDA pero, al arrancar una aplicación que utiliza CUDA, veremos errores del tipo:

CUDA_SETUP: WARNING! libcudart.so not found in any environmental path. Searching /usr/local/cuda/lib64...

Esto sucede porque la variable LD_LIBRARY_PATH no contiene la ruta necesaria. Busca el fichero libcudart.so en tu sistema y añade esa ruta a la variable. En el caso de Arch Linux yo tuve que exportar esta ruta: LD_LIBRARY_PATH=/opt/cuda/targets/x86_64-linux/lib:$LD_LIBRARY_PATH.

cuDNN

cuDNN es opcional, pero se trata de una librería creada por NVIDIA que ofrece mejoras cuando trabajamos con redes neuronales. Es decir, aumenta la eficiencia de ciertos algoritmos y obtendremos un extra de velocidad.

Para acceder a las descargas necesitaremos una cuenta gratuita en el portal de desarrolladores de NVIDIA. También puedes probar a encontrar la versión que necesitas sin necesidad de registro en la siguiente página: https://developer.download.nvidia.com/compute/redist/cudnn/.

Para instalar cuDNN en Windows o distribuciones Linux basadas en Debian/Red Hat, accedemos al archivo del portal de desarrolladores de NVIDIA (https://developer.nvidia.com/rdp/cudnn-download) y seleccionamos la última versión de cuDNN compatible con la versión de CUDA que hayamos instalado. A día de hoy, cuDNN v8.9.6 es la última versión disponible para la rama de CUDA 11.x.

Download cuDNN v8.9.6 (November 1st, 2023), for CUDA 11.x

Otras versiones más antiguas de cuDNN se pueden encontrar aquí: https://developer.nvidia.com/rdp/cudnn-archive.

En Arch Linux no necesitamos descargar nada del portal de desarrolladores, ni tampoco necesitamos una cuenta (aunque viene bien para consultar qué versión de cuDNN que necesitamos). Dercargaremos el paquete del archivo de archlinux como hicimos con CUDA. En este caso accedemos a https://archive.archlinux.org/packages/c/cudnn/ y buscamos la versión que necesitamos. En nuestro caso descargamos cudnn-8.9.6.50-1-x86_64.pkg.tar.zst y lo instalamos con sudo pacman -U cudnn-8.9.6.50-1-x86_64.pkg.tar.zst.

Si tienes problemas para encontrar las versiones compatibles puedes consultar la página de documentación cuDNN Support Matrix, donde encontrarás tablas de compatibilidad entre distintas combinaciones de sistemas operativos, drivers y versiones de CUDA/cuDNN.

Comprobación

cuDNN se limita a copiar unos ficheros en el sistema, así que si buscamos en el sistema el fichero cudnn.h y lo encontramos, se habrá instalado correctamente. En Linux podemos buscarlo mediante find / -name cudnn.h.

PyTorch

Ahora sí, vamos a instalar la versión de PyTorch elegida. Accedemos a la página web de PyTorch y encontraremos una tabla cuya primera opción es la actual versión estable de PyTorch. Si ésta es la versión que necesitamos, basta con seguir eligiendo opciones hasta recibir el comando de instalación que ejecutaremos. En la fila Your OS elegimos nuestro sistema operativo, en la fila Package elegimos Pip, en la fila Language elegimos Python y en la fila Compute Platform elegimos la versión de CUDA que instalamos previamente (o ROCm/CPU si procede).

Actualmente la versión estable es la 2.1.1, por lo que esa no es la versión que necesitamos para este artículo. Para encontrar otras versiones accedemos a la página https://pytorch.org/get-started/previous-versions/. En nuestro caso buscamos la sección llamada v2.0.1, después la subsección Wheel y dentro de Linux and Windows (OSX si usamos Apple) encontraremos el comando de instalación según nuestra plataforma de computación:

# ROCM 5.4.2 (Linux only)
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/rocm5.4.2
# CUDA 11.7
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2
# CUDA 11.8
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118
# CPU only
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cpu

En este artículo optamos por utilizar CUDA 11.8 así que ejecutamos dicho comando.

Si estás utilizando Stable Diffusion web UI de Automatic1111, cada vez que se ejecuta la aplicación se actualizarán las librerías. Puedes evitar la actualización de torch mediante la variable TORCH_COMMAND, utilizando como valor el mismo comando de instalación.

Ejemplo: TORCH_COMMAND="pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118"

Comprobación

Para comprobar que PyTorch se ha instalado correctamente, ejecutamos el siguiente comando:

python -c "import torch; print(torch.__version__)"
2.0.1+cu118

Y si queremos comprobar que PyTorch tiene acceso a la GPU mediante CUDA:

python -c "import torch; print(torch.cuda.is_available())"
True

xFormers

xFormers es una librería desarrollada por Facebook que brinda una serie de mejoras y optimizaciones cuando trabajamos con Transformers. En el caso de modelos generativos como Stable Diffusion, utilizar xFormers puede reducir considerablemente los requerimientos de memoria. La mayoría de aplicaciones citadas en la introducción son capaces de utilizar xFormers. Consulta la documentación del proyecto para descubrirlo.

La lista de versiones de xformers se puede encontrar en https://pypi.org/project/xformers/#history.

Lamentablemente no he encontrado ninguna tabla de compatibilidad entre versiones de PyTorch y versiones de xFormers. Si hacemos pip install xformers, se instalará la última versión de xFormers y además se actualizará PyTorch a la última versión, algo que no queremos.

Aún así es bastante sencillo de solucionar, únicamente hay que añadir xformers al comando con el que instalamos PyTorch:

pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 xformers --index-url https://download.pytorch.org/whl/cu118

De esta manera forzamos la instalación de una versión de xformers que sea compatible con torch==2.0.1.

Si estás utilizando Stable Diffusion web UI de Automatic1111, cada vez que se ejecuta la aplicación se actualizarán las librerías. Puedes evitar la actualización de xformers mediante la variable XFORMERS_PACKAGE.

Ejemplo: XFORMERS_PACKAGE="xformers==0.0.22.post7+cu118".

Puedes ver la versión que se ha instalado mediate pip list | grep xformers.

Comprobación

Para comprobar que xFormers está instalado correctamente, ejecutamos el siguiente comando:

python -m xformers.info
xFormers 0.0.22
memory_efficient_attention.cutlassF:               available
memory_efficient_attention.cutlassB:               available
memory_efficient_attention.decoderF:               available
memory_efficient_attention.flshattF@v2.3.0:        available
memory_efficient_attention.flshattB@v2.3.0:        available
...

Veremos más información de la aquí mostrada, pero lo importante es ver available en las líneas de memory efficient attention.

Si vemos el texto WARNING[XFORMERS]: xFormers can't load C++/CUDA extensions., significa que no se ha instalado correctamente.

Conclusión

En este artículo hemos visto como resolver cualquier problema de versiones en torno a estas dependencias. Resumiendo:

  1. Elegir una versión de PyTorch según las necesidades de la aplicación que vamos a utilizar.
  2. Instalar CUDA si queremos aprovechar el rendimiento que nos ofrece una GPU NVIDIA. Instalar cuDNN para acelerar más aún el software.
  3. Instalar PyTorch con el comando de instalación que nos brinda su sitio web, eligiendo la plataforma de computación adecuada.
  4. Instalar xFormers para reducir los requerimientos de memoria si la aplicación lo soporta.

Puedes apoyarme para que pueda dedicar aún más tiempo a escribir artículos y tener recursos para crear nuevos proyectos. ¡Gracias!