Strategie voor belastingsverdeling zonder hulpondersteuning voor Mixture-of-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

Samenvatting

Voor Mixture-of-Experts (MoE)-modellen zal een ongelijke belasting van experts leiden tot routeringscollaps of een verhoogde rekenkundige overhead. Bestaande methoden gebruiken vaak een hulploss om belastingbalans te bevorderen, maar een grote hulploss introduceert niet-verwaarloosbare interferentiegradiënten in de training en schaadt daardoor de modelprestaties. Om de belastingbalans te beheersen zonder ongewenste gradiënten tijdens de training te produceren, stellen we Loss-Free Balancing voor, gekenmerkt door een strategie voor belastingbalans zonder hulploss. Concreet past Loss-Free Balancing, vóór de top-K-routeringsbeslissing, eerst een expert-specifieke bias toe op de routeringsscores van elke expert. Door de bias van elke expert dynamisch bij te werken op basis van de recente belasting, kan Loss-Free Balancing consistent een gebalanceerde verdeling van de expertbelasting handhaven. Bovendien, omdat Loss-Free Balancing geen interferentiegradiënten produceert, verhoogt het ook de bovengrens van de modelprestaties die worden behaald met MoE-training. We valideren de prestaties van Loss-Free Balancing op MoE-modellen met tot 3B parameters getraind op tot 200B tokens. Experimentele resultaten tonen aan dat Loss-Free Balancing zowel betere prestaties als een betere belastingbalans bereikt in vergelijking met traditionele strategieën voor belastingbalans die worden gecontroleerd door een hulploss.

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.