Дьявол кроется в деталях: о реализации потерь балансировки нагрузки для обучения специализированных моделей смеси экспертов.

Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models

January 21, 2025

Авторы: Zihan Qiu, Zeyu Huang, Bo Zheng, Kaiyue Wen, Zekun Wang, Rui Men, Ivan Titov, Dayiheng Liu, Jingren Zhou, Junyang Lin

cs.AI

Аннотация

Эта статья пересматривает реализацию балансировки нагрузки (Load-balancing Loss, LBL) при обучении моделей смеси экспертов (Mixture-of-Experts, MoEs). Конкретно, LBL для MoEs определяется как сумма N_E sum_{i=1}^{N_E} f_i p_i, где N_E - общее количество экспертов, f_i представляет собой частоту выбора эксперта i, а p_i обозначает средний балл шлюза эксперта i. Существующие рамки обучения MoE обычно используют стратегию параллельного обучения, чтобы f_i и LBL были рассчитаны в пределах микро-пакета, а затем усреднены по параллельным группам. По сути, микро-пакет для обучения миллиардных LLM обычно содержит очень немного последовательностей. Поэтому микро-пакет LBL практически на уровне последовательности, и маршрутизатор вынужден равномерно распределять токены в каждой последовательности. При таком строгом ограничении даже токены из последовательности, специфичной для определенной области (например, код), равномерно направляются ко всем экспертам, тем самым препятствуя специализации экспертов. В данной работе мы предлагаем рассчитывать LBL с использованием глобального пакета, чтобы снять это ограничение. Поскольку глобальный пакет содержит намного более разнообразные последовательности, чем микро-пакет, это будет способствовать балансировке нагрузки на уровне корпуса. Конкретно, мы вводим дополнительный шаг коммуникации для синхронизации f_i между микро-пакетами, а затем используем его для расчета LBL. Через эксперименты по обучению LLM на основе MoEs (до 42,8 млрд общего количества параметров и 400 млрд токенов) мы удивительно обнаруживаем, что стратегия глобального пакета LBL приводит к отличному улучшению производительности как в предварительной перплексии, так и в задачах на следующем этапе. Наш анализ показывает, что глобальный пакет LBL также значительно улучшает специализацию области экспертов MoE.

English

This paper revisits the implementation of Load-balancing Loss (LBL) when training Mixture-of-Experts (MoEs) models. Specifically, LBL for MoEs is defined as N_E sum_{i=1}^{N_E} f_i p_i, where N_E is the total number of experts, f_i represents the frequency of expert i being selected, and p_i denotes the average gating score of the expert i. Existing MoE training frameworks usually employ the parallel training strategy so that f_i and the LBL are calculated within a micro-batch and then averaged across parallel groups. In essence, a micro-batch for training billion-scale LLMs normally contains very few sequences. So, the micro-batch LBL is almost at the sequence level, and the router is pushed to distribute the token evenly within each sequence. Under this strict constraint, even tokens from a domain-specific sequence (e.g., code) are uniformly routed to all experts, thereby inhibiting expert specialization. In this work, we propose calculating LBL using a global-batch to loose this constraint. Because a global-batch contains much more diverse sequences than a micro-batch, which will encourage load balance at the corpus level. Specifically, we introduce an extra communication step to synchronize f_i across micro-batches and then use it to calculate the LBL. Through experiments on training MoEs-based LLMs (up to 42.8B total parameters and 400B tokens), we surprisingly find that the global-batch LBL strategy yields excellent performance gains in both pre-training perplexity and downstream tasks. Our analysis reveals that the global-batch LBL also greatly improves the domain specialization of MoE experts.