GRIN: Градиентно-информированный MoE

GRIN: GRadient-INformed MoE

September 18, 2024

Авторы: 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

Аннотация

Модели смеси экспертов (MoE) масштабируются более эффективно, чем плотные модели, благодаря разреженным вычислениям через маршрутизацию экспертов, выборочно активирующую только небольшой поднабор экспертных модулей. Однако разреженные вычисления представляют вызов для традиционных методов обучения, поскольку дискретная маршрутизация экспертов затрудняет стандартное обратное распространение и, следовательно, градиентно-ориентированную оптимизацию, которые являются основой глубокого обучения. Для более эффективного использования масштабирующей силы MoE мы представляем GRIN (GRadient-INformed MoE training), который включает разреженную оценку градиента для маршрутизации экспертов и настраивает параллелизм модели для предотвращения потери токенов. Применяя GRIN к авторегрессионному языковому моделированию, мы разрабатываем модель MoE размером 16times3.8B с топ-2. Наша модель, с активированными параметрами всего 6.6B, превосходит плотную модель размером 7B и соответствует производительности плотной модели размером 14B, обученной на тех же данных. Обширные оценки по различным задачам демонстрируют потенциал GRIN для значительного улучшения эффективности MoE, достигая 79.4 на MMLU, 83.7 на HellaSwag, 74.4 на HumanEval и 58.9 на 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.