ES
Un sistema de entrenamiento adapta la memoria a los picos de carga de los expertos
BaGuaLu-2: Taming Dynamic Memory Fluctuation for Efficient MoE Model Training over 10,000 GPUs · arXiv:2610.23064
En palabras sencillas
Algunos modelos de lenguaje grandes envían cada token a redes expertas seleccionadas. La cantidad de tokens que recibe cada experto puede volverse desigual de repente y aumentar la memoria temporal necesaria durante el entrenamiento.
Una estrategia fija de memoria puede desperdiciar recursos en los grupos habituales o arriesgarse a sufrir fallos durante un pico. BaGuaLu-2 es un sistema de entrenamiento que cambia las medidas de ahorro de memoria para cada grupo pequeño y solapa parte de ese trabajo con otras tareas. En una ejecución con 512 GPU, la medida de uso del cómputo que informó se mantuvo estable durante picos que redujeron la medida en dos sistemas de comparación.
La investigación
El entrenamiento de modelos de mezcla de expertos a gran escala distribuye los tokens entre muchas redes expertas. Los picos infrecuentes de enrutamiento pueden aumentar la demanda de memoria de activaciones, mientras que una estrategia fija de memoria debe ralentizar el trabajo habitual o dejar la ejecución expuesta a fallos por falta de memoria.
¿Cómo puede responder el entrenamiento a esos cambios breves sin asumir el mismo costo en cada grupo?
BaGuaLu-2 elige medidas de ahorro de memoria para cada grupo pequeño de entrenamiento y solapa ese trabajo entre grupos cuando cambia la presión de memoria.
BaGuaLu-2 estima las necesidades de memoria de activaciones y, para cada microgrupo, elige entre volver a calcular, trasladar los valores fuera de la GPU, cuantizarlos o no aplicar medidas. Solapa ese trabajo con operaciones de grupos cercanos para reducir el tiempo dedicado a la gestión de memoria.
En las configuraciones evaluadas de modelos y clústeres, de 128 a 12.288 GPU, los autores informan un rendimiento promedio de principio a fin 1,14 veces mayor que el de Megatron-Kuai y 1,13 veces mayor que el de AdaPipe.
El rendimiento aquí promedia la utilización de FLOPS del modelo (MFU), una medida de qué parte de la capacidad de cálculo disponible de las GPU se utiliza, en cada iteración de toda la ejecución, incluido el calentamiento de las máquinas después de los fallos. Los resultados respaldan un mayor rendimiento general en estas pruebas.
La comparación promedio abarca las configuraciones evaluadas, que comenzaron con la misma configuración de paralelismo seleccionada fuera de línea; el artículo no especifica si compartían los mismos puntos de control iniciales.
El cálculo de los reinicios incluye el calentamiento de las máquinas, pero omite el diagnóstico y la intervención de las personas, por lo que subestima esa parte del costo de reiniciar.
Qué muestran las pruebas
- Picos de enrutamiento infrecuentes
Los modelos de mezcla de expertos a gran escala distribuyen los tokens entre redes expertas. La carga de esos expertos puede variar de una iteración de entrenamiento a otra y cambiar la cantidad de memoria temporal de activaciones necesaria.
Los autores midieron un desequilibrio superior a 1,2 veces la media en el 2,9 % de las iteraciones, con una tasa máxima de 2,6 veces. La pregunta central es cómo puede el entrenamiento gestionar picos ocasionales de memoria sin recurrir todo el tiempo a una estrategia fija y costosa.
- Ajuste de memoria para cada grupo
BaGuaLu-2 es un sistema de entrenamiento que estima la demanda de memoria de cada microgrupo y luego elige qué operaciones volver a calcular, trasladar fuera de la GPU, reducir en precisión o dejar sin cambios. Su presupuesto de memoria durante la ejecución se adapta a las condiciones actuales de la cadena de procesamiento.
También intercala el trabajo de microgrupos consecutivos para ocultar parte del tiempo que añade la gestión de memoria.
- Comparación de ejecuciones completas
Los autores compararon BaGuaLu-2 con Megatron-Kuai y AdaPipe en configuraciones de 128 a 12.288 GPU. Los sistemas comenzaron con la misma configuración de paralelismo seleccionada fuera de línea, mientras que los de comparación usaron estrategias estáticas de memoria y podían reiniciarse manualmente después de los fallos.
El rendimiento total promedió la utilización de FLOPS del modelo (MFU), una medida de qué parte de la capacidad de cálculo disponible de las GPU se utiliza, en cada iteración de toda la ejecución, incluido el calentamiento de las máquinas tras los reinicios. El promedio de BaGuaLu-2 fue 1,14 veces el de Megatron-Kuai y 1,13 veces el de AdaPipe.
- Evidencia sobre escala y componentes
La ejecución evaluada a mayor escala entrenó un modelo de 1.040 mil millones de parámetros en 12.288 GPU e informó un rendimiento total de entrenamiento de 2,012 exaFLOPS. Una comparación progresiva de componentes también asoció la incorporación sucesiva de medidas estáticas más precisas, medidas dinámicas y planificación en tiempo de ejecución que solapa tareas con aumentos del rendimiento respecto de la configuración anterior.
Estos resultados describen la ejecución a gran escala y el análisis de ablación informados; no aíslan el efecto de cada componente en todas las cargas de trabajo.
- Rendimiento durante los picos
En una ejecución representativa de un modelo de 423 mil millones de parámetros en 512 GPU, la medida de utilización de FLOPS del modelo de los sistemas de comparación disminuyó durante los picos de desequilibrio, mientras que la de BaGuaLu-2 se mantuvo estable.
Durante esa fase, BaGuaLu-2 promedió un rendimiento 1,28 veces mayor que el de los sistemas de comparación. El resultado se refiere a la condición concreta para la que se diseñó el sistema, pero procede de una sola configuración de modelo y clúster.
Detalles
- Autores del paper original
- Mingshu Zhai, Minkun Guo, Hongrui Lian, Shuangyu Li, Qi Zhang
- Enviado a arXiv
- Licencia
- CC BY 4.0