전문가 집합을 위한 보조 손실 없는 부하 분산 전략

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

August 28, 2024

저자: Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, Damai Dai

cs.AI

초록

Mixture-of-Experts (MoE) 모델의 경우, 불균형한 전문가 부하는 라우팅 붕괴나 증가된 계산 오버헤드로 이어질 수 있습니다. 기존 방법들은 주로 보조 손실을 사용하여 부하 균형을 촉진하지만, 큰 보조 손실은 훈련 중 무시할 수 없는 간섭 그래디언트를 도입하여 모델 성능을 손상시킬 수 있습니다. 훈련 중 원치 않는 그래디언트를 생성하지 않으면서 부하 균형을 제어하기 위해, 우리는 손실 없는 균형을 제안합니다. 이는 보조 손실이 없는 부하 균형 전략을 특징으로 합니다. 구체적으로, 상위 K개의 라우팅 결정 이전에, 손실 없는 균형은 각 전문가의 라우팅 점수에 전문가별 편향을 먼저 적용할 것입니다. 최근 부하에 따라 각 전문가의 편향을 동적으로 업데이트함으로써, 손실 없는 균형은 전문가 부하의 균형 재분배를 일관되게 유지할 수 있습니다. 게다가, 손실 없는 균형은 어떠한 간섭 그래디언트도 생성하지 않기 때문에 MoE 훈련으로 얻은 모델 성능 상한선을 높일 수도 있습니다. 우리는 최대 3B 매개변수를 가진 MoE 모델에서 200B 토큰까지 훈련된 손실 없는 균형의 성능을 검증합니다. 실험 결과는 손실 없는 균형이 전통적인 보조 손실 제어 부하 균형 전략과 비교하여 더 나은 성능과 부하 균형을 달성한다는 것을 보여줍니다.

English

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.