Demons in the Detail: On Implementing Load Balancing Loss for Training
  Specialized Mixture-of-Expert Models
Authors: name: Zihan Qiu; name: Zeyu Huang; name: Bo Zheng; name: Kaiyue Wen; name: Zekun Wang; name: Rui Men; name: Ivan Titov; name: Dayiheng Liu; name: Jingren Zhou; name: Junyang Lin
Published: 2025-01-21 04:04:39

Abstract: 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.