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

TL;DR

This work identifies a fundamental constraint in training mixture-of-experts models: micro-batch load-balancing losses force uniform router behavior within small parallel groups, hindering domain specialization. It introduces a global-batch LBL via cross-group synchronization and a buffering mechanism to approximate global statistics when full global batches are unavailable, and demonstrates substantial gains in both language modeling perplexity and downstream tasks across large MoE-LMs. The results show improved model performance and clearer domain-specific expert specialization, with scalable benefits that grow with Balance BSZ and only modest overhead. The proposed approach offers a practical, integrable way to realize more diverse and interpretable expert usage in MoE-based systems across domains.

Abstract

This paper revisits the implementation of $\textbf{L}$oad-$\textbf{b}$alancing $\textbf{L}$oss (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 $\textbf{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 ($\textit{e.g.}$, code) are uniformly routed to all experts, thereby inhibiting expert specialization. In this work, we propose calculating LBL using a $\textbf{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 $\textbf{42.8B}$ total parameters and $\textbf{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.

Figures (4)

• Figure 1: The impact of the Balance BSZ on (a) model performance and (b) expert specialization. (a) When employing micro-batch level load balance, methods based on LBL and based on auxiliary-loss-free wang2024auxiliary approaches perform worse than employing the global-batch balance. (b) When employing micro-batch balance, there is no significant difference in the selection frequency of different domain-specific data, and the selection frequency of different experts within the same domain is approximately the same. With global-batch balance, there is a noticeable difference in the selection frequency of experts on different domain data, and within the same domain, there are experts with high selection frequency (marked in blue).
• Figure 2: The performance of MoE-3.4A0.6B trained on 400B tokens with different Balance BSZ.
• Figure 3: The LBL (Top) and language modeling loss (Bottom) curve for MoE-3.4A0.6B trained on 400B tokens under different Balance BSZ, with a zoom-in of the last 15k steps shown below. In addition to the settings with Balance BSZ of 512 and 2, there are settings where the Balance BSZ changes from 2 to 512 at 10k, 30k, and 50k steps and a setting where it changes from 512 to 2 at the 50k step.
• Figure 4: The topK score sums across layers (a), and the distribution of high-frequency experts on different domains for models using global-batch balance (b). The topK sum of global-batch balance is higher than other methods and shows a similar distribution of high-frequency experts on closer domains.