GRIN: GRadiënt-Geïnformeerde MoE

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

Samenvatting

Mixture-of-Experts (MoE) modellen schalen effectiever dan dichte modellen vanwege schaarse berekeningen via expert routing, waarbij slechts een klein subset van expertmodules selectief geactiveerd wordt. Echter, schaarse berekeningen vormen een uitdaging voor traditionele trainingsmethoden, aangezien discrete expert routing standaard backpropagation belemmert en daarmee gradientgebaseerde optimalisatie, die de hoeksteen vormen van diep leren. Om de schaalbaarheid van MoE beter te benutten, introduceren we GRIN (GRadient-INformed MoE training), dat schattingen van schaarse gradienten voor expert routing incorporeert en modelparallelisme configureert om token dropping te vermijden. Door GRIN toe te passen op autoregressieve taalmodellering, ontwikkelen we een top-2 16 keer 3,8 miljard MoE-model. Ons model, met slechts 6,6 miljard geactiveerde parameters, presteert beter dan een 7 miljard dicht model en evenaart de prestaties van een 14 miljard dicht model dat op dezelfde data is getraind. Uitgebreide evaluaties over diverse taken tonen het potentieel van GRIN aan om de effectiviteit van MoE aanzienlijk te verbeteren, met scores van 79,4 voor MMLU, 83,7 voor HellaSwag, 74,4 voor HumanEval en 58,9 voor 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.