An Almost-Free Method to Improve MoE LLM Training Performance: Global-Batch Load Balancing

The conventional Load Balancing Loss used when training Mixture-of-Experts (MoE) architectures is typically computed at the micro-batch level. However, when a micro-batch consists solely of data from a specific domain, forcing the model to use all experts equally hinders the experts' specialization.

To address this, the proposed Global-batch load balance synchronizes expert selection frequencies across all parallel groups and aggregates the loss across the entire micro-batch. Since the expert selection frequency takes the form of a single vector, the synchronization cost across parallel groups is very low, making it close to an 'almost free lunch' in terms of efficiency.

Experimental results across MoE models of various sizes—3.4B, 15B, and 43B—showed that global-batch load balancing achieved the following outcomes:

• Achieved superior performance compared to the existing method across all settings (models, data, tasks)
• Confirmed a clear domain specialization phenomenon where experts are activated according to specific domains
• Enabled stable training even with larger batch sizes

To prevent degradation in training efficiency, adding a small amount of micro-batch load balancing loss to the global-batch loss can maintain model performance while increasing computational efficiency.