Global-Batch Load Balance: Um Almoço Quase Grátis para Melhorar o Treinamento de MoE LLM

A arquitetura Mixture-of-Experts (MoE) é uma técnica popular para escalar parâmetros de modelos, mas frameworks de treinamento existentes, como o Megatron-core, usam perda de balanceamento no nível de micro-batches, o que pode prejudicar o desempenho quando os micro-batches contêm dados homogêneos (por exemplo, apenas código). A Alibaba Qwen introduz a perda de balanceamento por lote global, que sincroniza a frequência de seleção de especialistas em todos os grupos paralelos e calcula a perda globalmente. Experimentos com três configurações MoE (3,4B/0,6B ativados, 15B/2,54B ativados, 43B/6,6B ativados) e até 400 bilhões de tokens mostram que o balanceamento por lote global supera consistentemente o balanceamento por micro-batch em várias tarefas. Ele também leva a uma significativa especialização de domínio dos especialistas, que está ausente em modelos balanceados por micro-batch. Adicionar uma pequena perda de balanceamento por micro-batch (peso 0,01) sobre a perda global melhora a velocidade de treinamento sem afetar a eficácia. O método é quase gratuito, pois a sincronização envolve apenas um vetor de tamanho igual ao número de especialistas. O trabalho foi publicado no arXiv preprint 2501.11873.