Dettagli Cruciali: Sull'Implementazione della Perdita di Bilanciamento del Carico per l'Addestramento di Modelli Specializzati a Mischia di Esperti

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

January 21, 2025

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

cs.AI

Abstract

Questo articolo rivede l'implementazione del Bilanciamento del Carico della Perdita (LBL) durante l'addestramento dei modelli Mixture-of-Experts (MoEs). In particolare, LBL per MoEs è definito come la somma N_E sum_{i=1}^{N_E} f_i p_i, dove N_E è il numero totale di esperti, f_i rappresenta la frequenza dell'esperto i selezionato e p_i indica il punteggio medio di gating dell'esperto i. I framework di addestramento MoE esistenti di solito adottano la strategia di addestramento parallelo in modo che f_i e il LBL siano calcolati all'interno di un micro-batch e poi mediati tra gruppi paralleli. In sostanza, un micro-batch per l'addestramento di LLM su scala miliardaria contiene normalmente pochissime sequenze. Quindi, il LBL del micro-batch è quasi a livello di sequenza, e il router è spinto a distribuire uniformemente il token all'interno di ciascuna sequenza. In virtù di questo vincolo rigoroso, anche i token da una sequenza specifica del dominio (ad esempio, codice) vengono instradati uniformemente a tutti gli esperti, inibendo così la specializzazione dell'esperto. In questo lavoro, proponiamo di calcolare il LBL utilizzando un global-batch per allentare questo vincolo. Poiché un global-batch contiene sequenze molto più diverse rispetto a un micro-batch, ciò incoraggerà un equilibrio del carico a livello di corpus. In particolare, introduciamo un passaggio di comunicazione aggiuntivo per sincronizzare f_i tra micro-batch e quindi utilizzarlo per calcolare il LBL. Attraverso esperimenti sull'addestramento di LLM basati su MoEs (fino a 42,8 miliardi di parametri totali e 400 miliardi di token), scopriamo sorprendentemente che la strategia del LBL del global-batch porta a eccellenti miglioramenti delle prestazioni sia nella perplessità di pre-addestramento che nei compiti successivi. La nostra analisi rivela che il LBL del global-batch migliora notevolmente anche la specializzazione di dominio degli esperti MoE.

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.