Estrategia de Equilibrio de Carga sin Pérdida de Pérdida Auxiliar para Mezcla de Expertos

Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts

August 28, 2024

Autores: Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, Damai Dai

cs.AI

Resumen

Para los modelos de Mezcla de Expertos (MoE), una carga de expertos desequilibrada conducirá a un colapso en el enrutamiento o a un aumento en la carga computacional. Los métodos existentes comúnmente emplean una pérdida auxiliar para fomentar el equilibrio de carga, pero una pérdida auxiliar grande introducirá gradientes de interferencia no despreciables en el entrenamiento y, por lo tanto, perjudicará el rendimiento del modelo. Con el fin de controlar el equilibrio de carga sin producir gradientes no deseados durante el entrenamiento, proponemos el Balance sin Pérdida, caracterizado por una estrategia de equilibrio de carga libre de pérdida auxiliar. Específicamente, antes de la decisión de enrutamiento de los mejores K, el Balance sin Pérdida aplicará primero un sesgo por experto a los puntajes de enrutamiento de cada experto. Al actualizar dinámicamente el sesgo de cada experto de acuerdo con su carga reciente, el Balance sin Pérdida puede mantener consistentemente una distribución equilibrada de la carga de expertos. Además, dado que el Balance sin Pérdida no produce gradientes de interferencia, también eleva el límite superior del rendimiento del modelo obtenido del entrenamiento de MoE. Validamos el rendimiento del Balance sin Pérdida en modelos de MoE con hasta 3 mil millones de parámetros entrenados con hasta 200 mil millones de tokens. Los resultados experimentales muestran que el Balance sin Pérdida logra tanto un mejor rendimiento como un mejor equilibrio de carga en comparación con las estrategias tradicionales de equilibrio de carga controladas por pérdida auxiliar.

English

For Mixture-of-Experts (MoE) models, an unbalanced expert load will lead to routing collapse or increased computational overhead. Existing methods commonly employ an auxiliary loss to encourage load balance, but a large auxiliary loss will introduce non-negligible interference gradients into training and thus impair the model performance. In order to control load balance while not producing undesired gradients during training, we propose Loss-Free Balancing, featured by an auxiliary-loss-free load balancing strategy. To be specific, before the top-K routing decision, Loss-Free Balancing will first apply an expert-wise bias to the routing scores of each expert. By dynamically updating the bias of each expert according to its recent load, Loss-Free Balancing can consistently maintain a balanced distribution of expert load. In addition, since Loss-Free Balancing does not produce any interference gradients, it also elevates the upper bound of model performance gained from MoE training. We validate the performance of Loss-Free Balancing on MoE models with up to 3B parameters trained on up to 200B tokens. Experimental results show that Loss-Free Balancing achieves both better performance and better load balance compared with traditional auxiliary-loss-controlled load balancing strategies.