Duivels in het Detail: Over het Implementeren van Load Balancing Loss voor Training van Gespecialiseerde Mixture-of-Expert Modellen

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

January 21, 2025

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

cs.AI

Samenvatting

Deze paper herziet de implementatie van Load-balancing Loss (LBL) bij het trainen van Mixture-of-Experts (MoE) modellen. Specifiek wordt LBL voor MoE gedefinieerd als N_E sum_{i=1}^{N_E} f_i p_i, waarbij N_E het totale aantal experts is, f_i de frequentie van expert i die geselecteerd wordt vertegenwoordigt, en p_i de gemiddelde gate score van expert i aanduidt. Bestaande MoE trainingskaders maken doorgaans gebruik van de parallelle trainingsstrategie zodat f_i en de LBL worden berekend binnen een micro-batch en vervolgens gemiddeld worden over parallelle groepen. In essentie bevat een micro-batch voor het trainen van LLMs op miljardenschaal meestal zeer weinig sequenties. Daarom bevindt de micro-batch LBL zich bijna op sequentieniveau, en wordt de router gedwongen om de token gelijkmatig binnen elke sequentie te verdelen. Onder deze strikte beperking worden zelfs tokens van een domeinspecifieke sequentie (bijv. code) uniform gerouteerd naar alle experts, waardoor expertspecialisatie wordt belemmerd. In dit werk stellen we voor om LBL te berekenen met behulp van een global-batch om deze beperking te versoepelen. Omdat een global-batch veel diversere sequenties bevat dan een micro-batch, zal dit een evenwichtige belasting op corpusniveau aanmoedigen. Specifiek introduceren we een extra communicatiestap om f_i te synchroniseren over micro-batches en gebruiken dit vervolgens om de LBL te berekenen. Uit experimenten met het trainen van MoE-gebaseerde LLMs (tot 42,8B totale parameters en 400B tokens) blijkt verrassend genoeg dat de global-batch LBL-strategie uitstekende prestatiewinst oplevert, zowel in perplexiteit bij voorafgaande training als bij taken nadien. Onze analyse onthult dat de global-batch LBL ook aanzienlijk de domeinspecialisatie van MoE-experts verbetert.

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.