Estratégia de Balanceamento de Carga Livre de Perda de Auxiliar para Mistura de Especialistas

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

August 28, 2024

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

cs.AI

Resumo

Para modelos de Mixture-of-Experts (MoE), uma carga de especialistas desequilibrada levará a colapso de roteamento ou aumento da sobrecarga computacional. Métodos existentes comumente empregam uma perda auxiliar para incentivar o equilíbrio de carga, mas uma grande perda auxiliar introduzirá gradientes de interferência não negligenciáveis no treinamento e, portanto, prejudicará o desempenho do modelo. Para controlar o equilíbrio de carga sem produzir gradientes indesejados durante o treinamento, propomos o Balanceamento sem Perda, caracterizado por uma estratégia de balanceamento de carga livre de perda auxiliar. Especificamente, antes da decisão de roteamento top-K, o Balanceamento sem Perda aplicará primeiro um viés por especialista às pontuações de roteamento de cada especialista. Ao atualizar dinamicamente o viés de cada especialista de acordo com sua carga recente, o Balanceamento sem Perda pode manter consistentemente uma distribuição equilibrada da carga de especialistas. Além disso, como o Balanceamento sem Perda não produz quaisquer gradientes de interferência, ele também eleva o limite superior do desempenho do modelo obtido a partir do treinamento MoE. Validamos o desempenho do Balanceamento sem Perda em modelos MoE com até 3B de parâmetros treinados com até 200B de tokens. Os resultados experimentais mostram que o Balanceamento sem Perda alcança tanto melhor desempenho quanto melhor equilíbrio de carga em comparação com estratégias tradicionais de balanceamento de carga controladas por perda auxiliar.

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.