Detalles en la Práctica: Sobre la Implementación de la Pérdida de Equilibrio de Carga para Entrenar Modelos Especializados de Mezcla de Expertos

Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models

January 21, 2025

Autores: Zihan Qiu, Zeyu Huang, Bo Zheng, Kaiyue Wen, Zekun Wang, Rui Men, Ivan Titov, Dayiheng Liu, Jingren Zhou, Junyang Lin

cs.AI

Resumen

Este documento revisita la implementación del Balance de Carga de Pérdida (LBL, por sus siglas en inglés) al entrenar modelos de Mezcla de Expertos (MoEs). Específicamente, LBL para MoEs se define como la suma de N_E sum_{i=1}^{N_E} f_i p_i, donde N_E es el número total de expertos, f_i representa la frecuencia de selección del experto i, y p_i denota la puntuación de enrutamiento promedio del experto i. Los marcos de entrenamiento de MoE existentes suelen emplear la estrategia de entrenamiento paralelo para que f_i y el LBL se calculen dentro de un micro lote y luego se promedien en grupos paralelos. En esencia, un micro lote para entrenar LLMs a escala de miles de millones normalmente contiene muy pocas secuencias. Por lo tanto, el LBL del micro lote está casi a nivel de secuencia, y el enrutador se ve obligado a distribuir el token de manera uniforme dentro de cada secuencia. Bajo esta estricta restricción, incluso los tokens de una secuencia específica de dominio (por ejemplo, código) se dirigen uniformemente a todos los expertos, lo que inhibe la especialización de los expertos. En este trabajo, proponemos calcular LBL usando un lote global para liberar esta restricción. Dado que un lote global contiene secuencias mucho más diversas que un micro lote, lo que fomentará el equilibrio de carga a nivel de corpus. Específicamente, introducimos un paso de comunicación adicional para sincronizar f_i entre micro lotes y luego usarlo para calcular el LBL. A través de experimentos en el entrenamiento de LLMs basados en MoEs (hasta 42.8B parámetros totales y 400B tokens), encontramos sorprendentemente que la estrategia de LBL de lote global produce excelentes ganancias de rendimiento tanto en la perplejidad de pre-entrenamiento como en las tareas posteriores. Nuestro análisis revela que el LBL de lote global también mejora en gran medida la especialización de dominio de los expertos de MoE.

English

This paper revisits the implementation of Load-balancing Loss (LBL) when training Mixture-of-Experts (MoEs) models. Specifically, LBL for MoEs is defined as N_E sum_{i=1}^{N_E} f_i p_i, where N_E is the total number of experts, f_i represents the frequency of expert i being selected, and p_i denotes the average gating score of the expert i. Existing MoE training frameworks usually employ the parallel training strategy so that f_i and the LBL are calculated within a micro-batch and then averaged across parallel groups. In essence, a micro-batch for training billion-scale LLMs normally contains very few sequences. So, the micro-batch LBL is almost at the sequence level, and the router is pushed to distribute the token evenly within each sequence. Under this strict constraint, even tokens from a domain-specific sequence (e.g., code) are uniformly routed to all experts, thereby inhibiting expert specialization. In this work, we propose calculating LBL using a global-batch to loose this constraint. Because a global-batch contains much more diverse sequences than a micro-batch, which will encourage load balance at the corpus level. Specifically, we introduce an extra communication step to synchronize f_i across micro-batches and then use it to calculate the LBL. Through experiments on training MoEs-based LLMs (up to 42.8B total parameters and 400B tokens), we surprisingly find that the global-batch LBL strategy yields excellent performance gains in both pre-training perplexity and downstream tasks. Our analysis reveals that the global-batch LBL also greatly improves the domain specialization of MoE experts.