¿Te imaginas poder acelerar el entrenamiento de modelos de lenguaje de gran tamaño (LLMs) y mejorar su rendimiento? Un equipo de investigadores de Rice University ha hecho realidad esta posibilidad con un avance significativo llamado ZEN. Este nuevo sistema optimiza la comunicación entre unidades de procesamiento gráfico (GPUs), y promete transformar la forma en que entrenamos estos modelos, abordando dos grandes cuellos de botella que suelen surgir en el proceso.
¿Cuáles son los cuellos de botella en el entrenamiento de LLMs?
Los LLMs enfrentan principalmente dos problemas críticos durante su entrenamiento. El primero ocurre cuando el modelo tiene que procesar enormes volúmenes de datos, lo que puede ralentizar el sistema y consumir muchos recursos computacionales. Para solucionar esto, el equipo de investigación, liderado por Zhuang Wang y el profesor T.S. Eugene Ng, sugiere dividir los datos entre cientos o incluso miles de GPUs. De esta manera, cada unidad puede manejar múltiples muestras de datos de forma independiente y enviar los resultados al modelo.
El segundo cuello de botella se presenta cuando todas las GPUs deben sincronizarse para compartir lo que han aprendido. Esta etapa es crucial, ya que una comunicación ineficiente puede causar retrasos significativos, especialmente cuando los gradientes del modelo son muy grandes. Wang explicó que el enfoque tradicional consistía en enviar todos los datos, pero esto era ineficaz, ya que muchos de esos datos estaban llenos de valores nulos.
¿Cómo funciona el sistema ZEN?
Para enfrentar este desafío, se introduce el concepto de sparsificación, que permite eliminar esos valores nulos o casi nulos. Así, solo se sincronizan los valores relevantes durante la comunicación, conocidos como tensores dispersos. Aunque esta técnica es común en el entrenamiento de LLMs, el verdadero reto radica en manejar estos tensores en un entorno de entrenamiento distribuido.
Ng subrayó que, a pesar de que la idea de usar tensores dispersos ya estaba en el aire, no se había logrado optimizar su manejo. El equipo se enfocó en estudiar el comportamiento de estos tensores para desarrollar mejores esquemas de comunicación. A través de su investigación, descubrieron que los gradientes no nulos no se distribuyen de manera uniforme, lo que puede causar desequilibrios durante la comunicación y, por ende, ralentizar la sincronización.
Resultados y aplicaciones de ZEN
Una vez que el equipo entendió cómo diseñar ZEN, el siguiente paso fue identificar los esquemas de comunicación más efectivos. Wang y Ng analizaron diversas opciones y, al no encontrar una solución ideal, construyeron un sistema basado en sus hallazgos. ZEN se implementó en el entrenamiento de LLMs y mostró un aumento notable en la velocidad de entrenamiento.
Wang aseguró que, gracias a ZEN, el tiempo necesario para completar cada paso del entrenamiento se ha reducido considerablemente. Este descubrimiento tiene implicaciones importantes, ya que ZEN puede aplicarse a una variedad de modelos de lenguaje, incluyendo aquellos utilizados para la generación de texto e imágenes, siempre que se utilicen tensores dispersos.
No es la primera vez que Wang y Ng colaboran en un proyecto innovador. Anteriormente, habían trabajado en GEMINI, un sistema diseñado para minimizar la sobrecarga tras fallos de hardware o software en el entrenamiento de LLMs, presentado en 2023. Recientemente, Wang compartió su trabajo sobre ZEN en el 19º Simposio de USENIX sobre Diseño e Implementación de Sistemas Operativos (OSDI) celebrado en Boston.
