Стратегия балансировки нагрузки без вспомогательных потерь для смеси экспертов

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

August 28, 2024

Авторы: Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, Damai Dai

cs.AI

Аннотация

Для моделей Mixture-of-Experts (MoE) неравномерная нагрузка экспертов приведет к сбою маршрутизации или увеличению вычислительной нагрузки. Существующие методы обычно используют вспомогательные потери для поощрения баланса нагрузки, однако большие вспомогательные потери приводят к значительному вмешательству градиентов в процессе обучения и, следовательно, ухудшают производительность модели. Для контроля баланса нагрузки без появления нежелательных градиентов во время обучения мы предлагаем Loss-Free Balancing, основанный на стратегии балансировки нагрузки без вспомогательных потерь. Конкретно, перед принятием решения о маршрутизации для топ-K, Loss-Free Balancing сначала применяет смещение для оценок маршрутизации каждого эксперта. Динамически обновляя смещение каждого эксперта в соответствии с его недавней нагрузкой, Loss-Free Balancing может постоянно поддерживать сбалансированное распределение нагрузки экспертов. Кроме того, поскольку Loss-Free Balancing не порождает никаких вмешательств градиентов, он также повышает верхнюю границу производительности модели, полученной в результате обучения MoE. Мы проверяем производительность Loss-Free Balancing на моделях MoE с до 3 миллиардов параметров, обученных на до 200 миллиардов токенов. Экспериментальные результаты показывают, что Loss-Free Balancing достигает как лучшей производительности, так и лучшего баланса нагрузки по сравнению с традиционными стратегиями балансировки нагрузки с контролем вспомогательных потерь.

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.