Hilfsverlustfreie Lastenausgleichsstrategie für das Mixture-of-Experts-Modell

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

August 28, 2024

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

cs.AI

Zusammenfassung

Für Mixture-of-Experts (MoE) Modelle führt eine unausgeglichene Expertenlast zu einem Routing-Zusammenbruch oder erhöhtem Rechenaufwand. Gängige Methoden verwenden eine Hilfsverlustfunktion, um eine ausgewogene Last zu fördern, jedoch führt ein großer Hilfsverlust zu nicht zu vernachlässigenden Interferenzgradienten während des Trainings und beeinträchtigt somit die Modellleistung. Um die Lastenausgeglichenheit zu kontrollieren, ohne unerwünschte Gradienten während des Trainings zu erzeugen, schlagen wir Loss-Free Balancing vor, das durch eine Hilfsverlustfreie Lastenausgleichsstrategie gekennzeichnet ist. Konkret wendet Loss-Free Balancing vor der Top-K-Routing-Entscheidung zunächst eine expertenweises Bias auf die Routing-Scores jedes Experten an. Durch die dynamische Aktualisierung des Bias jedes Experten gemäß seiner aktuellen Last kann Loss-Free Balancing kontinuierlich eine ausgewogene Verteilung der Expertenlast aufrechterhalten. Darüber hinaus hebt Loss-Free Balancing aufgrund der fehlenden Erzeugung von Interferenzgradienten auch die Obergrenze der Modellleistung, die aus dem MoE-Training erzielt wird, an. Wir validieren die Leistung von Loss-Free Balancing an MoE-Modellen mit bis zu 3 Milliarden Parametern, die auf bis zu 200 Milliarden Tokens trainiert wurden. Experimentelle Ergebnisse zeigen, dass Loss-Free Balancing im Vergleich zu traditionellen Strategien zur lastenausgeglichenen Steuerung mit Hilfsverlust sowohl eine bessere Leistung als auch eine bessere Lastenausgeglichenheit erzielt.

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.