Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts
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.
66d0746981e0683bca3fb7b0
Lean Wang
claimed_verified
2024-09-02T07:19:01.270Z
650c509472afb1e60e6151ae
/avatars/c16ab5053a586819dc2b965303215ff7.svg
Lean Wang
user
AdaHousman
66d0746981e0683bca3fb7b1
Huazuo Gao
admin_assigned
2024-08-29T16:47:01.712Z
64e370be59aa5366642ac329
/avatars/0fa1eb6ac6c1aeff3e65bc86a6617f64.svg
Huazuo Gao
user
gaohuazuo
66d0746981e0683bca3fb7b2
Chenggang Zhao
admin_assigned
2024-08-29T16:47:07.557Z
66053b1f9e3555d648b21c3d
/avatars/c8b33e7f702c4edb17add47f0eafe5e6.svg
Chenggang Zhao
user
LyricZ
66d0746981e0683bca3fb7b3
Xu Sun
66d0746981e0683bca3fb7b4
Damai Dai
extracted_confirmed
2024-08-30T02:45:15.707Z
659389f8de82e1ef7b9a8b13
/avatars/896ed9f4cdbd317493b303d070b7e12a.svg
Damai Dai
user
DeepSeekDDM
Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts
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.
Loss-Free Balancing improves performance and load balance in Mixture-of-Experts models by dynamically adjusting expert biases without introducing auxiliary loss gradients.
Mixture-of-Experts expert load routing collapse computational overhead auxiliary loss interference gradients top-K routing expert-wise bias recent load model performance