GRIN: MoE Informado por Gradiente

GRIN: GRadient-INformed MoE

September 18, 2024

Autores: Liyuan Liu, Young Jin Kim, Shuohang Wang, Chen Liang, Yelong Shen, Hao Cheng, Xiaodong Liu, Masahiro Tanaka, Xiaoxia Wu, Wenxiang Hu, Vishrav Chaudhary, Zeqi Lin, Chenruidong Zhang, Jilong Xue, Hany Awadalla, Jianfeng Gao, Weizhu Chen

cs.AI

Resumo

Os modelos de Mixture-of-Experts (MoE) escalam de forma mais eficaz do que os modelos densos devido à computação esparsa por meio do roteamento de especialistas, ativando seletivamente apenas um pequeno subconjunto de módulos de especialistas. No entanto, os desafios da computação esparsa para as práticas de treinamento tradicionais, uma vez que o roteamento discreto de especialistas dificulta a retropropagação padrão e, portanto, a otimização baseada em gradientes, que são a pedra angular do aprendizado profundo. Para melhor explorar o poder de escalabilidade do MoE, introduzimos o GRIN (GRadient-INformed MoE training), que incorpora a estimativa esparsa de gradientes para o roteamento de especialistas e configura o paralelismo de modelos para evitar a eliminação de tokens. Aplicando o GRIN à modelagem de linguagem autoregressiva, desenvolvemos um modelo MoE top-2 de 16 vezes 3,8B. Nosso modelo, com apenas 6,6B de parâmetros ativados, supera um modelo denso de 7B e iguala o desempenho de um modelo denso de 14B treinado nos mesmos dados. Avaliações extensas em diversas tarefas demonstram o potencial do GRIN para melhorar significativamente a eficácia do MoE, alcançando 79,4 em MMLU, 83,7 em HellaSwag, 74,4 em HumanEval e 58,9 em MATH.

English

Mixture-of-Experts (MoE) models scale more effectively than dense models due to sparse computation through expert routing, selectively activating only a small subset of expert modules. However, sparse computation challenges traditional training practices, as discrete expert routing hinders standard backpropagation and thus gradient-based optimization, which are the cornerstone of deep learning. To better pursue the scaling power of MoE, we introduce GRIN (GRadient-INformed MoE training), which incorporates sparse gradient estimation for expert routing and configures model parallelism to avoid token dropping. Applying GRIN to autoregressive language modeling, we develop a top-2 16times3.8B MoE model. Our model, with only 6.6B activated parameters, outperforms a 7B dense model and matches the performance of a 14B dense model trained on the same data. Extensive evaluations across diverse tasks demonstrate the potential of GRIN to significantly enhance MoE efficacy, achieving 79.4 on MMLU, 83.7 on HellaSwag, 74.4 on HumanEval, and 58.9 on MATH.