Modelos de entrenamiento de IA de 100 billones de parámetros

Los sistemas de IA de recomendación son un componente importante de los servicios de Internet en la actualidad: los negocios de ingresos de miles de millones de dólares como Amazon y Netflix están directamente impulsados ​​por los servicios de recomendación.

Los recomendadores de IA mejoran a medida que crecen. Recientemente se han lanzado varios modelos con miles de millones de parámetros hasta incluso billones. Cada salto en la capacidad del modelo ha traído una mejora significativa en la calidad. La era de los 100 billones de parámetros está a la vuelta de la esquina.

La red neuronal de descanso complicada y densa requiere cada vez más computación con más de 100 TFLOP en cada iteración de entrenamiento. Por lo tanto, es importante contar con algún mecanismo sofisticado para administrar un clúster con recursos heterogéneos para tales tareas de capacitación.

Recientemente, Kwai Seattle AI Lab y DS3 Lab de ETH Zurich han colaborado para proponer un sistema novedoso llamado "Persia" para abordar este problema mediante un diseño conjunto cuidadoso tanto del algoritmo de entrenamiento como del sistema de entrenamiento. A nivel de algoritmo, Persia adopta un algoritmo de entrenamiento híbrido para manejar la capa de incrustación y los módulos de red neuronal densa de manera diferente. La capa de incrustación se entrena de forma asíncrona para mejorar el rendimiento de las muestras de entrenamiento, mientras que el resto de la red neuronal se entrena de forma síncrona para preservar la eficiencia estadística. A nivel del sistema, se ha implementado una amplia gama de optimizaciones del sistema para la gestión de la memoria y la reducción de la comunicación para liberar todo el potencial del algoritmo híbrido.

Recursos en la nube para 100 billones de modelos de parámetros de IA

La carga de trabajo de IA de 100 billones de parámetros de Persia se ejecuta en los siguientes recursos heterogéneos:

3,000 núcleos de máquinas virtuales con uso intensivo de cómputo
8 máquinas virtuales A2 que suman un total de 64 GPU A100 Nvidia
30 máquinas virtuales de alta memoria, cada una con 12 TB de RAM, con un total de 360 ​​TB
Orquestación con Kubernetes
Todos los recursos debían lanzarse simultáneamente en la misma zona para minimizar la latencia de la red. Google Cloud pudo proporcionar la capacidad requerida con muy poco aviso.

AI Training necesita recursos en ráfagas.

Se utilizó Google Kubernetes Engine (GKE) para orquestar la implementación de 138 máquinas virtuales y contenedores de software. Tener la carga de trabajo en contenedores también permite la portabilidad y la repetibilidad de la capacitación.

Resultados y conclusiones
Con el apoyo de la infraestructura de Google Cloud, el equipo demostró la escalabilidad de Persia hasta 100 billones de parámetros. El algoritmo de entrenamiento distribuido híbrido introdujo relajaciones elaboradas del sistema para la utilización eficiente de clústeres heterogéneos, al mismo tiempo que convergía tan rápido como SGD estándar. Google Cloud fue esencial para superar las limitaciones del hardware en las instalaciones y demostró ser un entorno informático óptimo para la capacitación en aprendizaje automático distribuido a gran escala.

Persia se lanzó como un proyecto de código abierto en github con instrucciones de configuración para Google Cloud: a todos, tanto del mundo académico como de la industria, les resultaría fácil entrenar modelos recomendadores de aprendizaje profundo a una escala de 100 billones de parámetros.