Source-linked AI summary
Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models
Zihan Qiu, Zeyu Huang, Bo Zheng, Kaiyue Wen, Zekun Wang, Rui Men, Ivan Titov, Dayiheng Liu, Jingren Zhou, Junyang Lin
TL;DR
The paper addresses the overly strict sequence-level routing induced by micro-batch LBL, which can hinder MoE expert specialization. It replaces this with global-batch LBL by synchronizing expert-selection frequencies, and reports improved performance together with clearer domain specialization. The approach adds limited latency but is evaluated mainly for language-model pre-training.
Problem
Micro-batch LBL can force domain-specific tokens to distribute uniformly across experts, limiting specialization and harming MoE performance.
Method
Global-batch LBL synchronizes expert-selection frequencies across micro-batches and uses the aggregated statistics to compute the loss.
Results
Global-batch LBL improves pre-training perplexity and benchmark performance while producing clearer domain-level expert specialization.
Takeaways & Limitations
Balancing over diverse global-batch data relaxes micro-batch routing constraints and yields more performant and interpretable MoE models.
Takeaways & Limitations
The study focuses on language-model pre-training, does not evaluate fine-tuning or vision and multimodality, and reports some added latency.
Abstract
from arXiv · showhide
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.
1 Introduction
MoE training must balance expert utilization, but micro-batch LBL can impose overly strict sequence-level routing constraints that limit specialization and performance. The paper proposes global-batch LBL, which improves performance and produces clearer domain-level expert specialization.
- MoE routers require balanced expert utilization for effective updates and efficient parallel training.Imbalanced routing can cause parameter redundancy and overload expert-parallel processing.
- Micro-batches contain few sequences, so micro-batch LBL is nearly sequence-level and can force domain-specific tokens across all experts.This constraint may hurt model performance and inhibit specialization.
- Global-batch LBL synchronizes expert-selection frequencies across micro-batches and computes balancing over more diverse data.The approach relaxes the tight routing constraint imposed by micro-batch balancing.
- Global-batch balancing improves model performance and reveals domain-level, interpretable expert specialization.The paper reports approximately 0.1 improvement in pre-training PPL and 2 points in benchmark scores, while specialization emerges clearly only with global-batch LBL.
2 Preliminary
LBL regularizes MoE routing by penalizing concentration on a few experts, but existing parallel training frameworks compute it within individual parallel groups. This micro-batch constraint can prevent domain-specific routing and expert specialization.
- LBL encourages balanced expert utilization and prevents collapse by penalizing excessive routing to particular experts.It uses each expert’s routed-token fraction and total routing probability.
- Minimizing LBL distributes tokens more evenly among experts so each receives a fair share of training updates.This maintains balanced expert utilization during training.
- Existing frameworks compute routing statistics within each parallel group and average the resulting losses across groups.This procedure defines the paper’s micro-batch LBL setting.
- Micro-batch LBL can force domain-specific tokens to distribute uniformly across experts, hindering domain-specialized allocation.Micro-batches are often packed from one domain, whereas global batches combine micro-batches from different domains.
3 Method
The method converts micro-batch LBL into global-batch LBL by synchronizing expert-selection frequencies across parallel groups. When gradient accumulation prevents direct global-batch statistics, a buffer approximates the same computation across accumulation steps.
- Global-batch LBL synchronizes expert-selection frequencies across parallel groups before computing the loss.The synchronized statistics make the result equivalent to aggregating token statistics across the global batch.
- Communicating only the expert-number-dimensional frequency vector avoids transmitting token-by-expert matrices and expert-selection scores.This reduces the communication burden of global-batch balancing.
- Algorithm 1 initializes per-expert counts, adds synchronized counts at each accumulation step, computes current frequencies, and clears the buffer after optimization.This procedure approximates global-batch LBL when available micro-batches do not fill the target global batch.
- The approximate method buffers synchronized expert-selection counts across gradient-accumulation steps and resets the buffer after each optimizer step.The buffer maintains counts for each expert while the accumulated gradients are formed.
4 Experiments
Experiments evaluate global-batch load balancing across MoE model sizes, training scales, downstream benchmarks, and domain-specific held-out data. Increasing the balance batch size improves performance and produces stronger domain-level expert specialization.
- Experimental Setup: The experiments cover 3.4B, 15B, and 43B total-parameter MoE models trained on multilingual, math, and general-knowledge data.Evaluation uses four zero-shot benchmarks and held-out perplexity sets spanning English, Chinese, code, math, literature, law, and other domains.
- Global Load Balance: As Balance BSZ increases, all reported metrics consistently improve for the 3.4A0.6B model.A buffer approximating the global batch also performs significantly better than Balance BSZ 128 and approaches Balance BSZ 512.
- Expert Specialization: With micro-batch balance, expert-selection frequencies are mostly similar across domains, and no SFT-Code or SFT-Math expert exceeds 0.15.This pattern corresponds to limited domain-level specialization under default load balancing.
- Expert Specialization: With global-batch balance, many SFT-Math experts exceed 0.2 selection frequency, indicating more pronounced domain specialization.The comparison examines topK expert selections across SFT-Code, SFT-Math, and EN-Literature.
5 Analysis
The analysis shows that global-batch balancing is a looser constraint than micro-batch balancing, improves language-modeling performance, and promotes domain-linked expert specialization. It also examines token-distribution effects and the computational trade-off of global-batch balancing.
- Ablation Study on Token Numbers and Token Distributional Bias: Shuffle LBL_micro achieves performance similar to LBL_global and significantly outperforms LBL_micro, separating token-distribution effects from token-count effects.The ablation supports the claim that global-batch balancing helps because it relaxes sequence-level distributional constraints, not merely because it estimates frequencies from more tokens.
- LBL_global is a looser constraint than LBL_micro: Global-batch balancing requires even token distribution across the global batch, whereas micro-batch balancing requires uniform distribution within each micro-batch.Switching from micro-batch to global-batch balancing rapidly lowers LBL toward the global-batch trajectory; switching in the opposite direction sharply raises it before gradual decline.
- LBL_global is a looser constraint than LBL_micro: Global-batch balancing yields language-modeling loss over 0.02 lower than micro-batch balancing, while switching methods during training remains inferior to training with global-batch balancing from scratch.Earlier switches from micro-batch to global-batch balancing perform better, but even switching at 10k steps remains notably worse than starting with global-batch balancing.
- The computation cost and efficiency of global-batch balance: Global-batch balancing is 1.64 s/iteration versus 1.55 s/iteration for micro-batch balancing, making it 5.8% slower in the 43A6.6B setting.About 1% of the slowdown comes from communication, while the remainder mainly reflects local expert-load imbalance under the dropless strategy.
- The computation cost and efficiency of global-batch balance: Adding a small micro-batch balancing loss raises speed to 1.59 s/iteration, or 2.6% slower than baseline, with only a minimal performance decrease.The added micro-batch loss uses 1% of the global-batch loss weight and is introduced at the 20k step.
- Global batch balance brings interpretable specialization: Global-batch models have higher topK score sums across layers, while micro-batch models have lower sums with little domain difference.The results associate concentrated expert scores with stronger specialization; aux-loss-free global balancing lies between LBL and micro-batch balancing.
- Global batch balance brings interpretable specialization: Chinese domains share several high-frequency experts, whereas SFT-Code overlaps little with Chinese domains and SFT-EN has fewer highly activated individual experts.The distribution indicates domain-dependent expert usage under global-batch balancing, especially among related Chinese domains.
6 Related works
Related work addresses expert load balancing through routing mechanisms, bias-based adjustments, and global balancing, but differs in motivation, implementation setting, and isolation of global-balance effects. Prior work also documents limited domain-level specialization under tight micro-batch balancing.
- Load Balancing: Expert Choice Routing balances load by allowing each expert to select tokens according to its capacity, rather than relying on the same load-balancing formulation.The cited related work presents this as a routing alternative to conventional topK routing with load constraints.
- Load Balancing: Bias-based balancing adjusts an expert-selection bias from frequency without changing routing scores, but does not emphasize whether frequency is computed per micro-batch or global batch.The passage contrasts this approach with the paper’s focus on the granularity used to compute selection frequency.
- Load Balancing: GRIN introduces global load-balance adaptations but does not isolate global load balancing’s effects or provide the paper’s motivation and properties for using it.The comparison concerns both GRIN’s framework context and its limited analysis of global balancing as an independent intervention.
- Expert Specialization: Under tight micro-batch balancing, most MoE models, including multimodal MoEs, have not exhibited domain-level specialization.This motivates examining whether a less restrictive balancing strategy can recover the specialization expected from the divide-and-conquer design of MoEs.
7 Conclusion
The work replaces overly restrictive micro-batch balance with global-batch balance to improve MoE performance and expert specialization across models of various sizes.
- Global-batch balance relaxes the overly tight routing constraint imposed by micro-batch LBL.The proposed synchronization and buffering methods shift balancing from micro-batches to global batches.
- The methods improve model performance and foster domain-level, interpretable expert specialization.The authors report validation across models of various sizes.
Limitations
The study focuses on micro-batch LBL during LLM pre-training and leaves several settings and validation questions open.
- The analysis is limited to LLM pre-training and does not examine fine-tuning, vision, or multimodality.
- Specialization is assessed mainly through selection frequencies across domains rather than more rigorous validation.
- Relaxing micro-batch LBL can introduce latency.