Stratégie d'équilibrage de charge sans perte de perte auxiliaire pour Mélange d'Experts

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

August 28, 2024

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

cs.AI

Résumé

Pour les modèles de Mixture-of-Experts (MoE), une charge d'experts déséquilibrée entraînera un effondrement du routage ou une augmentation des coûts de calcul. Les méthodes existantes utilisent couramment une perte auxiliaire pour encourager l'équilibre de charge, mais une grande perte auxiliaire introduira des gradients d'interférence non négligeables dans l'entraînement et nuira ainsi aux performances du modèle. Afin de contrôler l'équilibre de charge sans produire de gradients indésirables pendant l'entraînement, nous proposons l'Équilibrage sans Perte, caractérisé par une stratégie d'équilibrage de charge sans perte auxiliaire. Plus précisément, avant la décision de routage top-K, l'Équilibrage sans Perte appliquera d'abord un biais spécifique à chaque expert aux scores de routage de chaque expert. En mettant à jour dynamiquement le biais de chaque expert en fonction de sa charge récente, l'Équilibrage sans Perte peut maintenir de manière constante une distribution équilibrée de la charge des experts. De plus, étant donné que l'Équilibrage sans Perte ne produit aucun gradient d'interférence, il élève également la limite supérieure des performances du modèle obtenues à partir de l'entraînement MoE. Nous validons les performances de l'Équilibrage sans Perte sur des modèles MoE avec jusqu'à 3 milliards de paramètres entraînés sur jusqu'à 200 milliards de jetons. Les résultats expérimentaux montrent que l'Équilibrage sans Perte atteint à la fois de meilleures performances et un meilleur équilibre de charge par rapport aux stratégies traditionnelles d'équilibrage de charge contrôlées par perte auxiliaire.

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.