Les démons dans les détails : Sur la mise en œuvre de la perte d'équilibrage de charge pour l'entraînement de modèles spécialisés de mélange d'experts.

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

January 21, 2025

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

cs.AI

Résumé

Ce document revisite la mise en œuvre de la perte d'équilibrage de charge (LBL) lors de l'entraînement des modèles de Mélange d'Experts (MoEs). Plus précisément, la LBL pour les MoEs est définie comme la somme de N_E pour i allant de 1 à N_E de f_i p_i, où N_E est le nombre total d'experts, f_i représente la fréquence de sélection de l'expert i, et p_i désigne le score de sélection moyen de l'expert i. Les cadres d'entraînement MoE existants utilisent généralement la stratégie d'entraînement parallèle de sorte que f_i et la LBL sont calculés au sein d'une micro-batch puis moyennés à travers des groupes parallèles. En essence, une micro-batch pour l'entraînement des LLM à l'échelle du milliard contient normalement très peu de séquences. Ainsi, la LBL de la micro-batch est presque au niveau de la séquence, et le routeur est poussé à distribuer le jeton de manière égale au sein de chaque séquence. Sous cette contrainte stricte, même les jetons d'une séquence spécifique au domaine (par exemple, du code) sont acheminés uniformément vers tous les experts, inhibant ainsi la spécialisation des experts. Dans ce travail, nous proposons de calculer la LBL en utilisant une global-batch pour lever cette contrainte. Étant donné qu'une global-batch contient beaucoup plus de séquences diverses qu'une micro-batch, cela encouragera l'équilibrage de charge au niveau du corpus. Plus précisément, nous introduisons une étape de communication supplémentaire pour synchroniser f_i à travers les micro-batches, puis l'utilisons pour calculer la LBL. À travers des expériences sur l'entraînement de LLM basés sur MoEs (jusqu'à 42,8 milliards de paramètres au total et 400 milliards de jetons), nous constatons de manière surprenante que la stratégie de LBL de la global-batch entraîne d'excellents gains de performance à la fois en termes de perplexité de pré-entraînement et de tâches en aval. Notre analyse révèle que la LBL de la global-batch améliore également grandement la spécialisation de domaine des experts 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.