Elamur Research⌕
☰
home / papers / 2024 / GRIN: GRadient-INformed MoE

GRIN: GRadient-INformed MoE

Paper · 2024 · Liyuan Liu, Young Jin Kim, Shuohang Wang et al.

Summary

Mixture-of-Experts (MoE) models activate only a few expert subnetworks per token. However, the discrete routing step has no true gradient, so standard training substitutes the gating gradient. Large-scale MoE training also usually relies on expert parallelism, which forces capacity limits and discards overflow tokens. The Microsoft authors address both problems with GRIN.

First, they introduce SparseMixer-v2. During training it treats expert selection as random sampling and estimates the routing gradient with a straight-through-style scheme based on Heun's third-order method. Second, they distribute training with pipeline and tensor parallelism instead of expert parallelism, so no token is ever dropped.

With this recipe they train GRIN MoE, a 16×3.8B decoder that routes each token to 2 of 16 experts. On the same pretraining data, it beats a 7B dense model and roughly equals a 14B dense one, with particularly strong coding and math results. The work shows that better routing gradients help sparse models use their capacity more fully.

Key points

  • Architecture: 32-layer transformer with an MoE block in every layer (16 SwiGLU experts, top-2 routing), grouped-query and sliding-window attention, and RoPE. It has 42B total and 6.6B active parameters. Weights are released as microsoft/GRIN-MoE, and a related variant is released as Phi-3.5-MoE-instruct.
  • Benchmark scores: MMLU 79.4, HellaSwag 83.7, HumanEval 74.4, MBPP 80.3, MATH 58.9.
  • Comparison with dense models: GRIN MoE's average across the benchmark suite is 79.58. Dense baselines pretrained on an identical corpus average 75.74 (7B) and 78.46 (14B).
  • Controlled ablation: two 16×0.9B models were trained on 2.5T tokens. GShard led for roughly the first 0.5T tokens, after which SparseMixer-v2 overtook it on MMLU.
  • Throughput and routing behavior: on 64 H100 GPUs, the 42B MoE reaches about 86.6% of the per-GPU throughput of a dense model with equal active size. Its routing patterns differ noticeably across MMLU subjects, which suggests the experts have specialized.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 works

Cite

Cite as 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 and Weizhu Chen. 2024. GRIN: GRadient-INformed MoE. arXiv:2409.12136.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2409.12136 · arXiv 2409.12136

@misc{liu2024grin,
  title = {{GRIN: GRadient-INformed MoE}},
  author = {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 and Weizhu Chen},
  year = {2024},
  doi = {10.48550/arxiv.2409.12136},
  eprint = {2409.12136},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2409.12136}
}

References

55 references, 29 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Abdin, M., Aneja, J., Awadalla, H., Awadallah, A., Awan, A. A., Bach, N., Bahree, A., Bakhtiari, A., Bao, J., Behl, H., Benhaim, A., Bilenko, M., Bjorck, J., Bubeck, S., Cai, M., Cai, Q., Chaudhary, V., Chen, D., Chen, D., Chen, W., Chen, Y.-C., Chen, Y.-L., Cheng, H., Chopra, P., Dai, X., Dixon, M.… — Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone
  2. Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebron, F., and Sanghai, S. GQA: Training generalized multi-query transformer models from multi-head checkpoints. In Bouamor, H., Pino, J., and Bali, K. (eds.), EMNLP, 2023.
  3. Austin, J., Odena, A., Nye, M., Bosma, M., Michalewski, H., Dohan, D., Jiang, E., Cai, C., Terry, M., Le, Q., and Sutton, C. Program synthesis with large language models. arXiv:2108.07732, 2021.
  4. Bengio, Y., Léonard, N., and Courville, A. C. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv:1308.3432, 2013. find
  5. Bisk, Y., Zellers, R., Gao, J., and Choi, Y. Piqa: Reasoning about physical commonsense in natural language. In AAAI, 2019.
  6. Chen, M., Tworek, J., Jun, H., Yuan, Q., de Oliveira Pinto, H. P., Kaplan, J., Edwards, H., Burda, Y., Joseph, N., Brockman, G., Ray, A., Puri, R., Krueger, G., Petrov, M., Khlaaf, H., Sastry, G., Mishkin, P., Chan, B., Gray, S., Ryder, N., Pavlov, M., Power, A., Kaiser, L., Bavarian, M., Winter, C.… — Evaluating Large Language Models Trained on Code
  7. Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. arXiv:1904.10509, 2019.
  8. Clark, C., Lee, K., Chang, M.-W., Kwiatkowski, T., Collins, M., and Toutanova, K. Boolq: Exploring the surprising difficulty of natural yes/no questions. In NAACL-HLT, 2019.
  9. Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv:1803.05457, 2018.
  10. Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., Hesse, C., and Schulman, J. Training verifiers to solve math word problems. arXiv:2110.14168, 2021.
  11. Dao, T. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv:2307.08691, 2023.
  12. Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. JMLR, 2022. find
  13. Gale, T., Narayanan, D., Young, C., and Zaharia, M. A. Megablocks: Efficient sparse training with mixture-of-experts. In MLSys, 2023. find
  14. Glorot, X. and Bengio, Y. Understanding the difficulty of training deep feedforward neural networks. In AISTATS, 2010. find
  15. Haider, E., Perez-Becker, D., Portet, T., Madan, P., Garg, A., Majercak, D., Wen, W., Kim, D., Yang, Z., Zhang, J., Sharma, H., Bullwinkel, B., Pouliot, M., Minnich, A., Chawla, S., Herrera, S., Warreth, S., Engler, M., Lopez, G., Chikanov, N., Dheekonda, R. S. R., Jagdagdorj, B.-E., Lutz, R., Lunde… — Phi-3 Safety Post-Training: Aligning Language Models with a "Break-Fix" Cycle
  16. He, K., Zhang, X., Ren, S., and Sun, J. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In ICCV, 2015. find
  17. Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D., and Steinhardt, J. Measuring mathematical problem solving with the math dataset. In NeurIPS, 2021.
  18. Hsieh, C.-Y., Li, C.-L., Yeh, C.-K., Nakhost, H., Fujii, Y., Ratner, A., Krishna, R., Lee, C.-Y., and Pfister, T. Distilling step-by-step! outperforming larger language models with less training data and smaller model sizes. In ACL Findings, 2023.
  19. Jacobs, R. A., Jordan, M. I., Nowlan, S. J., and Hinton, G. E. Adaptive mixtures of local experts. Neural Computation, 1991. find
  20. Jain, N., yeh Chiang, P., Wen, Y., Kirchenbauer, J., Chu, H.-M., Somepalli, G., Bartoldson, B. R., Kailkhura, B., Schwarzschild, A., Saha, A., Goldblum, M., Geiping, J., and Goldstein, T. NEFTune: Noisy embeddings improve instruction finetuning. In ICLR, 2024. find
  21. Jiang, A. Q., Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., de las Casas, D., Bressand, F., Lengyel, G., Lample, G., Saulnier, L., Lavaud, L. R., Lachaux, M.-A., Stock, P., Scao, T. L., Lavril, T., Wang, T., Lacroix, T., and Sayed, W. E. Mistral 7b. arXiv:2310.06825, 2023a. find
  22. Jiang, A. Q., Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., de las Casas, D., Bressand, F., Lengyel, G., Lample, G., Saulnier, L., Lavaud, L. R., Lachaux, M.-A., Stock, P., Scao, T. L., Lavril, T., Wang, T., Lacroix, T., and Sayed, W. E. Mistral 7b. arXiv:2310.06825, 2023b. find
  23. Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., Casas, D. d. l., Hanna, E. B., Bressand, F., et al. Mixtral of experts. arXiv:2401.04088, 2024. find
  24. Jin, D., Pan, E., Oufattole, N., Weng, W.-H., Fang, H., and Szolovits, P. What disease does this patient have? a large-scale open domain question answering dataset from medical exams. arXiv:2009.13081, 2020.
  25. Jordan, M. I. and Jacobs, R. A. Hierarchical mixtures of experts and the em algorithm. Neural Computation, 1994. find
  26. Joshi, M., Choi, E., Weld, D. S., and Zettlemoyer, L. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In ACL, 2017.
  27. Kim, Y. J., Awan, A. A., Muzio, A., Salinas, A. F. C., Lu, L., Hendy, A., Rajbhandari, S., He, Y., and Awadalla, H. H. Scalable and efficient moe training for multitask multilingual models. arXiv:2109.10465, 2021. — [arXiv:2109.10465](https://arxiv.org/abs/2109.10465) find
  28. Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N. M., and Chen, Z. Gshard: Scaling giant models with conditional computation and automatic sharding. In ICLR, 2021. find
  29. Lewis, M., Bhosale, S., Dettmers, T., Goyal, N., and Zettlemoyer, L. Base layers: Simplifying training of large, sparse models. In ICML, 2021. find
  30. Lin, S., Hilton, J., and Evans, O. Truthfulqa: Measuring how models mimic human falsehoods. In ACL, 2021.
  31. Liu, L., Liu, X., Gao, J., Chen, W., and Han, J. Understanding the difficulty of training transformers. In EMNLP, 2020.
  32. Liu, L., Dong, C., Liu, X., Yu, B., and Gao, J. Bridging discrete and backpropagation: Straight-through and beyond. In NeurIPS 2023, 2023a. find
  33. Liu, L., Gao, J., and Chen, W. Sparse backpropagation for moe training. In WANT@NeurIPS 2023, 2023b. find
  34. Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can a suit of armor conduct electricity? a new dataset for open book question answering. In EMNLP, 2018. find
  35. Mukherjee, S., Mitra, A., Jawahar, G., Agarwal, S., Palangi, H., and Awadallah, A. Orca: Progressive learning from complex explanation traces of gpt-4. arXiv:2306.02707, 2023.
  36. Narayanan, D., Shoeybi, M., Casper, J., LeGresley, P., Patwary, M., Korthikanti, V. A., Vainbrand, D., Kashinkunti, P., Bernauer, J., Catanzaro, B., Phanishayee, A., and Zaharia, M. A. Efficient large-scale language model training on gpu clusters using megatron-lm. In SC21, 2021. find
  37. Nie, Y., Williams, A., Dinan, E., Bansal, M., Weston, J., and Kiela, D. Adversarial nli: A new benchmark for natural language understanding. In ACL, 2020.
  38. OpenAI. Gpt-4 technical report. arXiv:2303.08774, 2023.
  39. Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Ermon, S., and Finn, C. Direct preference optimization: Your language model is secretly a reward model. NeurIPS, 2024.
  40. Rosenblatt, F. The perceptron, a perceiving and recognizing automaton. Cornell Aeronautical Laboratory, 1957. find
  41. Sakaguchi, K., Le Bras, R., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale. In AAAI, 2020. find
  42. Shazeer, N. M., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q. V., Hinton, G. E., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In ICLR, 2017. find
  43. Srivastava, A., Rastogi, A., Rao, A., Shoeb, A. A. M., Abid, A., Fisch, A., Brown, A. R., Santoro, A., Gupta, A., Garriga-Alonso, A., Kluska, A., Lewkowycz, A., Agarwal, A., Power, A., Ray, A., Warstadt, A., Kocurek, A. W., Safaya, A., Tazarv, A., Xiang, A., Parrish, A., Nie, A., Hussain, A., Askell… — Beyond the Imitation Game: Quantifying and extrapolating the capabilities of language models
  44. Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomput., 2024.
  45. Suzgun, M., Scales, N., Schärli, N., Gehrmann, S., Tay, Y., Chung, H. W., Chowdhery, A., Le, Q. V., Chi, E. H., Zhou, D., and Wei, J. Challenging big-bench tasks and whether chain-of-thought can solve them. In ACL Findings, 2023.
  46. Talmor, A., Herzig, J., Lourie, N., and Berant, J. Commonsenseqa: A question answering challenge targeting commonsense knowledge. In NAACL-HLT, 2019.
  47. Team, G., Mesnard, T., Hardin, C., Dadashi, R., Bhupatiraju, S., Pathak, S., Sifre, L., Rivière, M., Kale, M. S., Love, J., et al. Gemma: Open models based on gemini research and technology. arXiv:2403.08295, 2024. find
  48. Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L., Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshorn, A., Hossei… — Llama 2: Open Foundation and Fine-Tuned Chat Models
  49. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. In NIPS, 2017.
  50. Wei, T., Zhu, B., Zhao, L., Cheng, C., Li, B., Lü, W., Cheng, P., Zhang, J., Zhang, X., Zeng, L., Wang, X., Ma, Y., Hu, R., Yan, S., Fang, H., and Zhou, Y. Skywork-moe: A deep dive into training techniques for mixture-of-experts language models. arXiv:2406.06563, 2024.
  51. White, C., Dooley, S., Roberts, M., Pal, A., Feuer, B., Jain, S., Shwartz-Ziv, R., Jain, N., Saifullah, K., Naidu, S., et al. Livebench: A challenging, contamination-free llm benchmark. arXiv:2406.19314, 2024. find
  52. Yang, G., Hu, J. E., Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J. W., Chen, W., and Gao, J. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. In NeurIPS, 2021. find
  53. Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can a machine really finish your sentence? In ACL, 2019. find
  54. Zhong, W., Cui, R., Guo, Y., Liang, Y., Lu, S., Wang, Y., Saied, A., Chen, W., and Duan, N. Agieval: A human-centric benchmark for evaluating foundation models. In NAACL-HLT Findings, 2024.
  55. Zoph, B., Bello, I., Kumar, S., Du, N., Huang, Y., Dean, J., Shazeer, N., and Fedus, W. St-moe: Designing stable and transferable sparse expert models. arXiv:2202.08906, 2022. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/liu-2024-grin-gradient-informed-moe/