GRIN: MoE Informato dal Gradiente

GRIN: GRadient-INformed MoE

September 18, 2024

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

Abstract

I modelli Mixture-of-Experts (MoE) scalano in modo più efficace rispetto ai modelli densi grazie alla computazione sparsa attraverso il routing degli esperti, attivando selettivamente solo un piccolo sottoinsieme dei moduli esperti. Tuttavia, la computazione sparsa sfida le pratiche di addestramento tradizionali, poiché il routing esperto discreto ostacola la backpropagation standard e quindi l'ottimizzazione basata sui gradienti, che sono il fondamento del deep learning. Per perseguire meglio la potenza di scalabilità di MoE, introduciamo GRIN (GRadient-INformed MoE training), che incorpora stime sparse del gradiente per il routing esperto e configura il parallelismo del modello per evitare la perdita di token. Applicando GRIN alla modellazione del linguaggio autoregressiva, sviluppiamo un modello MoE top-2 16 volte 3.8B. Il nostro modello, con soli 6.6B di parametri attivati, supera un modello denso da 7B e eguaglia le prestazioni di un modello denso da 14B addestrato sugli stessi dati. Valutazioni approfondite su diverse attività dimostrano il potenziale di GRIN nel migliorare significativamente l'efficacia di MoE, raggiungendo 79.4 su MMLU, 83.7 su HellaSwag, 74.4 su HumanEval e 58.9 su 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.