Entrenar modelos de IA a gran escala de manera eficiente en cientos de nodos es un desafío, especialmente cuando la mayoría de las cargas de trabajo están optimizadas para el hardware Nvidia y CUDA. Sin embargo, los investigadores del Laboratorio Nacional Oak Ridge (ORNL) han avanzado significativamente en el entrenamiento de modelos de billones de parámetros utilizando las GPU MI250X de AMD.
La MI250X es un potente motor de cálculo para cargas de trabajo de alto rendimiento (HPC) tradicionales, que ofrece un impresionante rendimiento de punto flotante de doble precisión. Si bien la doble precisión no es necesaria para las cargas de trabajo de IA, la MI250X ofrece un rendimiento respetable de 383 teraflops cuando se opera a precisión FP16. Con 37.888 GPU a su disposición, el equipo de ORNL se propuso optimizar la supercomputadora Frontier para el entrenamiento de grandes modelos de IA.
ORNL enfrentó múltiples desafíos en este proyecto. En primer lugar, tuvieron que portar su proyecto al tiempo de ejecución ROCm de AMD, lo que requirió la colaboración de los desarrolladores de AMD. Convertir el código de CUDA a un formato agnóstico del proveedor no es una tarea sencilla, pero se están logrando avances en esta área.
Además, los investigadores encontraron que escalar el paralelismo de tensores en varios nodos llevaba a cuellos de botella de latencia. Los mejores resultados se lograron al limitar el paralelismo de tensores a un solo nodo con ocho GPU. Además, el equipo implementó el optimizador ZeRO-1 para reducir la sobrecarga de memoria y utilizó el complemento de la biblioteca de comunicación colectiva ROCm de Amazon Web Services (RCCL) para mejorar la estabilidad de la comunicación entre nodos.
En cuanto a la eficiencia, la escala débil (aumentar el número de GPU para un tamaño de problema fijo) resultó ser 100 por ciento eficiente. Sin embargo, escalar con un tamaño de problema fijo condujo a rendimientos decrecientes debido a diversos cuellos de botella.
Si bien el éxito de ORNL en el entrenamiento de grandes modelos de IA en hardware de AMD es digno de reconocimiento, todavía queda trabajo por hacer para mejorar el rendimiento de estas cargas de trabajo. La mayoría de los marcos de entrenamiento están diseñados para el hardware de Nvidia y el soporte para la plataforma ROCm sigue siendo limitado. Sin embargo, las lecciones aprendidas de este experimento pueden servir como un modelo para otras instalaciones que operan sistemas que no son de Nvidia ni basados en CUDA, ofreciendo esperanza para una mayor adopción del hardware de AMD en el entrenamiento de IA.
The source of the article is from the blog toumai.es