Strategia di Bilanciamento del Carico Senza Perdita di Loss Ausiliaria per la Miscela di Esperti

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

August 28, 2024

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

cs.AI

Abstract

Per i modelli Mixture-of-Experts (MoE), un carico di esperti sbilanciato porterà al collasso del routing o a un aumento del carico computazionale. I metodi esistenti comunemente utilizzano una perdita ausiliaria per incoraggiare il bilanciamento del carico, ma una grande perdita ausiliaria introdurrà gradienti di interferenza non trascurabili durante l'addestramento e comprometterà quindi le prestazioni del modello. Al fine di controllare il bilanciamento del carico senza produrre gradienti indesiderati durante l'addestramento, proponiamo il Bilanciamento Senza Perdita, caratterizzato da una strategia di bilanciamento del carico senza perdita ausiliaria. Per essere specifici, prima della decisione di routing top-K, il Bilanciamento Senza Perdita applicherà innanzitutto un bias specifico per esperto ai punteggi di routing di ciascun esperto. Aggiornando dinamicamente il bias di ciascun esperto in base al suo carico recente, il Bilanciamento Senza Perdita può mantenere costantemente una distribuzione bilanciata del carico degli esperti. Inoltre, poiché il Bilanciamento Senza Perdita non produce alcun gradiente di interferenza, aumenta anche il limite superiore delle prestazioni del modello ottenute dall'addestramento MoE. Confermiamo le prestazioni del Bilanciamento Senza Perdita su modelli MoE con fino a 3 miliardi di parametri addestrati su fino a 200 miliardi di token. I risultati sperimentali mostrano che il Bilanciamento Senza Perdita ottiene sia migliori prestazioni sia un miglior bilanciamento del carico rispetto alle tradizionali strategie di bilanciamento del carico controllate da perdite ausiliarie.

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.