Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts

TL;DR

This paper introduces Loss-Free Balancing, a novel load balancing strategy for Mixture-of-Experts models that maintains expert load balance without auxiliary loss, improving performance and efficiency.

Contribution

The paper proposes a new load balancing method for MoE models that avoids auxiliary loss and interference gradients, enhancing model performance and load distribution.

Findings

Achieves better load balance than traditional methods.

Improves model performance without auxiliary loss interference.

Validated on models with up to 3B parameters and 200B tokens.

Abstract

For Mixture-of-Experts (MoE) models, an unbalanced expert load will lead to routing collapse or increased computational overhead. Existing methods commonly employ an auxiliary loss to encourage load balance, but a large auxiliary loss will introduce non-negligible interference gradients into training and thus impair the model performance. In order to control load balance while not producing undesired gradients during training, we propose Loss-Free Balancing, featured by an auxiliary-loss-free load balancing strategy. To be specific, before the top-K routing decision, Loss-Free Balancing will first apply an expert-wise bias to the routing scores of each expert. By dynamically updating the bias of each expert according to its recent load, Loss-Free Balancing can consistently maintain a balanced distribution of expert load. In addition, since Loss-Free Balancing does not produce any interference gradients, it also elevates the upper bound of model performance gained from MoE training. We validate the performance of Loss-Free Balancing on MoE models with up to 3B parameters trained on up to 200B tokens. Experimental results show that Loss-Free Balancing achieves both better performance and better load balance compared with traditional auxiliary-loss-controlled load balancing strategies.