GRIN : Mélange d'Experts Informé par le Gradient

GRIN: GRadient-INformed MoE

September 18, 2024

Auteurs: 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

Résumé

Les modèles de Mélange d'Experts (MoE) sont plus efficaces à l'échelle que les modèles denses en raison de la computation clairsemée à travers le routage des experts, activant sélectivement seulement un petit sous-ensemble de modules experts. Cependant, la computation clairsemée pose des défis aux pratiques d'entraînement traditionnelles, car le routage discret des experts entrave la rétropropagation standard et donc l'optimisation basée sur les gradients, qui sont la pierre angulaire de l'apprentissage profond. Pour mieux exploiter la puissance d'échelle de MoE, nous introduisons GRIN (Entraînement MoE InforMé par les Gradients), qui intègre une estimation clairsemée des gradients pour le routage des experts et configure le parallélisme du modèle pour éviter la suppression de jetons. En appliquant GRIN à la modélisation de langage autorégressive, nous développons un modèle MoE top-2 de 16 fois 3,8 milliards. Notre modèle, avec seulement 6,6 milliards de paramètres activés, surpasse un modèle dense de 7 milliards et égale les performances d'un modèle dense de 14 milliards entraîné sur les mêmes données. Des évaluations approfondies sur diverses tâches démontrent le potentiel de GRIN à améliorer significativement l'efficacité de MoE, atteignant 79,4 sur MMLU, 83,7 sur HellaSwag, 74,4 sur HumanEval et 58,9 sur 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.