Demônios nos Detalhes: Sobre a Implementação da Perda de Balanceamento de Carga para Treinamento de Modelos Especializados de Mistura de Especialistas

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

January 21, 2025

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

cs.AI

Resumo

Este artigo revisita a implementação do Balanceamento de Carga de Perda (LBL) ao treinar modelos de Mistura de Especialistas (MoEs). Especificamente, o LBL para MoEs é definido como a soma de N_E sum_{i=1}^{N_E} f_i p_i, onde N_E é o número total de especialistas, f_i representa a frequência do especialista i ser selecionado, e p_i denota a pontuação média de portão do especialista i. Os frameworks de treinamento de MoE existentes geralmente empregam a estratégia de treinamento paralelo para que f_i e o LBL sejam calculados dentro de um micro lote e depois sejam médios em grupos paralelos. Em essência, um micro lote para treinar LLMs em escala de bilhões normalmente contém muito poucas sequências. Assim, o LBL do micro lote está quase no nível da sequência, e o roteador é pressionado a distribuir o token uniformemente dentro de cada sequência. Sob essa restrição rigorosa, até mesmo tokens de uma sequência específica de domínio (por exemplo, código) são roteados uniformemente para todos os especialistas, inibindo assim a especialização do especialista. Neste trabalho, propomos calcular o LBL usando um lote global para afrouxar essa restrição. Como um lote global contém sequências muito mais diversas do que um micro lote, isso incentivará o equilíbrio de carga no nível do corpus. Especificamente, introduzimos uma etapa de comunicação adicional para sincronizar f_i entre micro lotes e então usá-lo para calcular o LBL. Através de experimentos no treinamento de LLMs baseados em MoEs (até 42,8B de parâmetros totais e 400B de tokens), descobrimos surpreendentemente que a estratégia de LBL do lote global resulta em excelentes ganhos de desempenho tanto na perplexidade de pré-treinamento quanto nas tarefas subsequentes. Nossa análise revela que o LBL do lote global também melhora significativamente a especialização de domínio dos especialistas 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.