GRIN: GRadient-INformierte MoE

GRIN: GRadient-INformed MoE

September 18, 2024

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

Zusammenfassung

Mixture-of-Experts (MoE)-Modelle skalieren effektiver als dichte Modelle aufgrund der spärlichen Berechnung durch Experten-Routing, das selektiv nur eine kleine Teilmenge von Expertenmodulen aktiviert. Allerdings stellen spärliche Berechnungen traditionelle Trainingspraktiken vor Herausforderungen, da diskretes Experten-Routing herkömmliches Backpropagation behindert und somit die auf Gradienten basierende Optimierung, die das Fundament des Deep Learning bildet. Um die Skalierbarkeit von MoE besser zu verfolgen, führen wir GRIN (GRadient-INformed MoE-Training) ein, das spärliche Gradientenschätzung für das Experten-Routing integriert und die Modellparallelität konfiguriert, um das Fallenlassen von Tokens zu vermeiden. Durch die Anwendung von GRIN auf autoregressive Sprachmodellierung entwickeln wir ein Top-2 16mal3,8 Milliarden MoE-Modell. Unser Modell, mit nur 6,6 Milliarden aktivierten Parametern, übertrifft ein 7 Milliarden dichtes Modell und erreicht die Leistung eines 14 Milliarden dichten Modells, das auf denselben Daten trainiert wurde. Umfangreiche Bewertungen über verschiedene Aufgaben hinweg zeigen das Potenzial von GRIN, die Wirksamkeit von MoE signifikant zu steigern, wobei 79,4 bei MMLU, 83,7 bei HellaSwag, 74,4 bei HumanEval und 58,9 bei MATH erreicht werden.

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.