Cómo hacer que los modelos de IA funcionen hasta cuatro veces más rápido fundiendo sus capas
Hugging Face muestra que combinar varias operaciones matemáticas en una sola puede multiplicar por cuatro la velocidad de inferencia en PyTorch, sin tocar ni una línea del código del modelo.
Los modelos de inteligencia artificial generativa procesan información a través de capas sucesivas de cálculos matemáticos. Cada una de esas capas —que en el código de PyTorch (la biblioteca de programación más usada para entrenar y ejecutar modelos de IA) se representa con bloques llamados nn.Linear o nn.MLP— realiza operaciones por separado: primero multiplica números, luego suma, después aplica una función de activación (una fórmula que decide qué información pasa a la siguiente capa). El problema es que cada operación por separado obliga al ordenador a mover datos de la memoria a la GPU (el chip especializado que hace los cálculos) y viceversa, lo que ralentiza todo el proceso.
Ingenieros de Hugging Face (una plataforma donde desarrolladores comparten y descargan modelos de IA ya entrenados) han publicado un tutorial técnico en el que demuestran cómo fusionar esas operaciones —es decir, combinarlas en una sola para que la GPU las ejecute de un tirón— puede acelerar la velocidad de inferencia (el tiempo que tarda un modelo en generar una respuesta) hasta cuatro veces. La técnica, llamada kernel fusion, no es nueva, pero el artículo muestra paso a paso cómo aplicarla en PyTorch usando la herramienta torch.compile, que desde 2023 viene incluida de serie en la biblioteca.
El tutorial compara el rendimiento de un MLP (multilayer perceptron, un bloque básico de redes neuronales que apila varias capas lineales) en tres versiones: sin optimización, con compilación automática y con fusión manual escrita en Triton (un lenguaje de programación que permite escribir operaciones personalizadas para GPUs). Los resultados muestran que la versión fundida procesa lotes de datos hasta 3,8 veces más rápido que la versión estándar, sobre todo cuando el tamaño del modelo y los datos de entrada son grandes.
La parte más relevante para quien desarrolla modelos es que no hace falta reescribir el código del modelo. Basta con pasarlo por torch.compile antes de usarlo, y el compilador se encarga de detectar operaciones que se pueden fusionar. En muchos casos eso ya da una mejora notable de velocidad. Para quien quiera exprimir aún más el rendimiento, el artículo explica cómo escribir kernels fusionados a mano con Triton, algo que requiere más conocimiento técnico pero que puede ser necesario en aplicaciones donde cada milisegundo cuenta, como sistemas de asistentes de voz o moderación de contenido en tiempo real.
El tutorial forma parte de una serie sobre perfilado de rendimiento en PyTorch (técnicas para medir dónde un modelo pierde tiempo y optimizarlo). Hugging Face lo publica en abierto como parte de su estrategia de compartir conocimiento técnico con la comunidad de desarrolladores, algo habitual en la empresa, que basa gran parte de su negocio en facilitar el acceso a modelos y herramientas de IA de código abierto.
