GRIN: GRadient-INformed MoE
Authors: name: Liyuan Liu; name: Young Jin Kim; name: Shuohang Wang; name: Chen Liang; name: Yelong Shen; name: Hao Cheng; name: Xiaodong Liu; name: Masahiro Tanaka; name: Xiaoxia Wu; name: Wenxiang Hu; name: Vishrav Chaudhary; name: Zeqi Lin; name: Chenruidong Zhang; name: Jilong Xue; name: Hany Awadalla; name: Jianfeng Gao; name: Weizhu Chen
Published: 2024-09-18 17:00:20

Abstract: 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.