Dämonen im Detail: Zur Implementierung des Lastenausgleichsverlusts für das Training von spezialisierten Mixture-of-Expert-Modellen

Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models

January 21, 2025

Autoren: Zihan Qiu, Zeyu Huang, Bo Zheng, Kaiyue Wen, Zekun Wang, Rui Men, Ivan Titov, Dayiheng Liu, Jingren Zhou, Junyang Lin

cs.AI

Zusammenfassung

Dieses Papier überprüft die Implementierung des Load-Balancing Loss (LBL) beim Training von Mixture-of-Experts (MoEs)-Modellen. Spezifisch wird LBL für MoEs definiert als N_E sum_{i=1}^{N_E} f_i p_i, wobei N_E die Gesamtanzahl der Experten ist, f_i die Häufigkeit repräsentiert, mit der Experte i ausgewählt wird, und p_i den durchschnittlichen Gate-Score des Experten i bezeichnet. Bestehende MoE-Trainings-Frameworks verwenden in der Regel die parallele Trainingsstrategie, so dass f_i und das LBL innerhalb eines Mikro-Batches berechnet und dann über parallele Gruppen gemittelt werden. Im Wesentlichen enthält ein Mikro-Batch für das Training von LLMs im Milliarden-Maßstab normalerweise sehr wenige Sequenzen. Daher ist das Mikro-Batch-LBL fast auf Sequenzebene, und der Router wird dazu gebracht, das Token gleichmäßig innerhalb jeder Sequenz zu verteilen. Unter dieser strengen Einschränkung werden sogar Tokens aus einer domänenspezifischen Sequenz (z. B. Code) gleichmäßig an alle Experten geroutet, was die Expertenspezialisierung hemmt. In dieser Arbeit schlagen wir vor, das LBL unter Verwendung eines Global-Batches zu berechnen, um diese Einschränkung zu lockern. Da ein Global-Batch viel vielfältigere Sequenzen als ein Mikro-Batch enthält, wird dies die Lastenverteilung auf Korpus-Ebene fördern. Spezifisch führen wir einen zusätzlichen Kommunikationsschritt ein, um f_i über Mikro-Batches zu synchronisieren und verwenden es dann, um das LBL zu berechnen. Durch Experimente zum Training von MoEs-basierten LLMs (bis zu 42,8B Gesamtparameter und 400B Tokens) stellen wir überraschenderweise fest, dass die Strategie des Global-Batch-LBL hervorragende Leistungssteigerungen sowohl bei der Vor-Training-Perplexität als auch bei nachgelagerten Aufgaben bringt. Unsere Analyse zeigt, dass das Global-Batch-LBL auch die Domänenspezialisierung der MoE-Experten erheblich verbessert.

English

This paper revisits the implementation of Load-balancing Loss (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 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 (e.g., code) are uniformly routed to all experts, thereby inhibiting expert specialization. In this work, we propose calculating LBL using a 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 42.8B total parameters and 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.