GRIN: Mezclador de Expertos Informado por Gradiente

GRIN: GRadient-INformed MoE

September 18, 2024

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

Resumen

Los modelos de Mezcla de Expertos (MoE) escalan de manera más efectiva que los modelos densos debido a la computación dispersa a través del enrutamiento de expertos, activando selectivamente solo un pequeño subconjunto de módulos expertos. Sin embargo, la computación dispersa desafía las prácticas de entrenamiento tradicionales, ya que el enrutamiento discreto de expertos obstaculiza la retropropagación estándar y, por lo tanto, la optimización basada en gradientes, que son la piedra angular del aprendizaje profundo. Para perseguir mejor el poder de escalado de MoE, presentamos GRIN (entrenamiento de Mezcla de Expertos informado por Gradiente), que incorpora estimación dispersa de gradientes para el enrutamiento de expertos y configura el paralelismo del modelo para evitar la eliminación de tokens. Aplicando GRIN a la modelización de lenguaje autoregresivo, desarrollamos un modelo MoE top-2 de 16 veces 3.8 mil millones. Nuestro modelo, con solo 6.6 mil millones de parámetros activados, supera a un modelo denso de 7 mil millones y coincide con el rendimiento de un modelo denso de 14 mil millones entrenado con los mismos datos. Evaluaciones extensas en diversas tareas demuestran el potencial de GRIN para mejorar significativamente la eficacia de MoE, logrando 79.4 en MMLU, 83.7 en HellaSwag, 74.4 en HumanEval y 58.9 en 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.