Aceleración CUDA en AMDGPU… o de como logre ejecutar Stable Diffusion en una vieja RADEON RX 5500XT con 4GB de VRAM, y no morir en el intento.

Tiempos difíciles ¿post?-pandemia: recesión, sequías por un lado, inundaciones por otro, guerra, muerte y destrucción; y lo peor de todo: escases de produccion chips mas un montón de cryptobros acaparándose de las todas las GPU decentes, bueno, al menos esto ultimo ya está mejorando gracias a que su burbuja comenzó a reventar.

Pero hace un año, debido a eso, me toco cambiarme hacia el lado rojo de AMD, ya que los precio de las GPU nVidia estaban por los cielos, si es que lograba encontrar alguna; lo cual implicaba ya no poder juguetear decentemente en mi PC con cuestiones que implicaran aceleración gráfica por hardware, entiéndase principalmente juegos, pero también experimentos locales con TensorFlow/keras o PyTorch, por lo que toco temporalmente subir a la nube voladora de Colab y Kaggle gratuitamente mientras tanto.

Mis intereses se renovaron hace unas semanas con salida de Dall-E, proyecto de «Open»AI que, al igual que GPT-3, son intersantes para juguetear un poco con ese montón de numeritos decimales ordenados de cierta manera que generan cosas interesantes, hasta donde lo «Open» de OpenAI se termina al quererte imponer restricciones moralmente estúpidas, y encima de eso cobrar su uso y abrirlo unicamente a las mega corporaciones como Microsoft entre otras. Afortunadamente hace una semana unos tipos alemanes agrupados bajo el nombre de StabilityAI sacaron al mundo un opción realmente competente realmente abierta: Stable Diffusion, algo que tenia que probar.

El problema era que en mi actual configuración libre de nVidia y 100% AMD, el API CUDA de AMD ROCm parecía apestar, comparadolo con API de nVidia que esta ampliamente disponible para casi todas sus GPU recientes y cuenta con soporte completo por la mayoría de frameworks de ML actuales; ROCm por otro lado, «oficialmente» soporta solo unas pocas GPUs. Para iniciar los experimentos me tocó algunas opciones CPU-only de Stable Diffusion, sin embargo no era muy cómodo esperar mas de 2 minutos usando un CPU AMD Zen2 para generar una sola imagen de mala calidad.

Gracias a los amigos de reddit, 4chan y hackernews, noté que varios usuarios no sólo ya habían logrado ejecutar Stable Diffusion en algunos AMD GPU recientes como las RX 6600, sino ademas versiones optimizadas que levantan con al menos 4GB de VRAM. por lo que me puse a echar a andar el experimento en mi vieja (pero actualizada) y fiel instalación de Arch Linux (est. 2009) de mi PC AMD actual.

Siguiendo mas o menos una serie de pasos hace un par de días, me percate que mi entorno de python estaba de cabeza, por lo que además tuve que arreglarlo.

Luego de eso, habría que instalar una serie de paquetes desde los repositorios de usuarios de Arch Linux (AUR) para soportar ROCm, en ese entonces aun no estaba seguro si funcionaria, ya que aunque la documentación «oficial» solo listaba unas pocas GPU soportadas, ya otros habian logrado ejecutarlas en otras GPU compatibles con las arquitecturas RDNA.

El primer consejo era instalar rocminfo y validar el id de mi GPU (gfx1012), para colocarla como variable de entorno antes de compilar ROCm, para que de esa manera, LLVM compilara hacia esa GPU… supongo.

Luego de varias horas de compilación, en donde me sentí como hace 18 años cuando compile el kernel de Liinux, gnome y mplayer para poder soporta ntfs y mp3 en fedora core 3 sobre mi vieja pentium4. Y luego de una lluvia de caracteres del stdout de CMAKE, GCC y CLANG, falla en un paquete especifico: rocalution.

El problema era que había una condición del pre-procesador en hip_utils.hpp en el cual solo tomaba en cuenta las gpu gfx1030 para compilar cierta lógica de multiplicación de matrices (o algo así) a CUDA. Despues de reportar el incidente especificamente para mi GPU en pull request del open issue del repo oficial, manualmente me dirigí al directorio de cache de paru, agregue mi gpu al #ifndef en dicha linea de código del hpp en vez del gfx1030, y luego de un makepkg -e logre instalarlo sin problema alguno, abriendo así la posibilidad de, al menos, procesar pytorch por CUDA en mi AMD, por fin.

Todo bien hasta el momento pero el otro problema que tenia era la limitada VRAM de mi GPU, lo recomendado eran mas de 8 GB oficialmente para Stable Diffusion, afortunadamente en estos últimos dias, gracias al poder del open source, mucha gente se dio a la tarea, no solo de optimizarlo para GPUs mas discretas, sino ademas agregar otros cosillas como levantar un frontend gradio.

Luego de seguir todos los pasos recomendados en cada uno de los proyectos «optimizados», como instalar la versión correcta de PyTorch compilada para ROCm, verificar que realmente se use cuda, bajar e instalar los modelos, y de varios intentos fallidos usando otras versiones, al final estaba listo para comenzar a generar imágenes de gatitos aterrizando en marte a alrededor 2 iteraciones por segundo, una gran mejora comparado al rendimiento de usar unicamente el CPU,

Al final en esta pseudo-cronica, mi interés no esta en generar gatitos voladores o chicas con tres tetas, sino simbolizar mi protesta a herramientas privativas e ironicamente cerradas como las de OpenAI, que solo buscan beneficiar a las «megacorps» (…aunque la licencia sobre la que se publica esta herramienta no es la mas abierta por sus posturas morales y politicas propias del nuevo orden mundial de inclusion), y al mismo tiempo experimentar como seria correr un DNN en mi actual configuración, que ademas de alimentarme de ricos neurotransmisores que se generan con el «Sí se pudo», abrirá las puertas a poder correr otro sin fin de proyectos en mi actual GPU, que últimamente estuvo relegada sólo a correr juegos de Steam sobre Proton durante los últimos años.


Deja un comentario