Elamur Research⌕
☰
home / papers / 2024 / Skywork-MoE: A Deep Dive into Training Techniques for Mixture-of-Experts Language Models

Skywork-MoE: A Deep Dive into Training Techniques for Mixture-of-Experts Language Models

Paper · 2024 · Tianwen Wei, Bo Zhu, Liang Zhao et al.

Summary

Training sparse Mixture-of-Experts (MoE) LLMs involves several practical choices whose effects are poorly understood. Two examples are whether to start from a dense checkpoint or from random weights, and how to keep experts both distinct and evenly loaded. This technical report from Kunlun's Skywork Team addresses these choices while building Skywork-MoE. The model has 146B total parameters, 22B of which are active per token. It uses 16 experts with top-2 routing and is initialized from Skywork-13B.

Small-scale experiments with a 0.3B dense model and 8-expert MoEs show that upcycling is not always the better start. With a larger training budget, a randomly initialized MoE overtakes the upcycled variants. The authors also propose two techniques:

  • Gating logit normalization, which makes routing decisions sharper.
  • Adaptive auxiliary loss coefficients, which tune the load-balancing weight separately for each layer.

The final model was trained mainly on a filtered part of the SkyPile corpus plus synthetic data. It beats Deepseek-67B and Llama2-70B on most benchmarks but stays behind Llama3-70B and larger MoEs. The work offers practical guidance for planning MoE training budgets.

Key points

  • Upcycling vs. scratch:
    • With a 100B-token budget, the from-scratch MoE roughly matched the one upcycled from the 100B-token checkpoint.
    • With 300B tokens, the from-scratch MoE outperformed every upcycled version.
    • Upcycled runs with a too-low learning rate did worst.
  • Expert similarity: average similarity between experts falls during upcycled training, so the authors suggest it as a useful monitoring signal. The worst run kept the highest similarity, while from-scratch experts stayed near zero similarity throughout.
  • Gating logit normalization: this yielded a moderately lower loss and fewer dropped tokens. Without it, the gap between the top gate probabilities shrinks toward a ratio of one, which the authors consider harmful.
  • Adaptive auxiliary loss coefficients: the per-layer coefficient follows each layer's token drop rate, rising when imbalance grows.
  • Data and results:
    • The training mix is about 7:2:1 English:Chinese:code.
    • Benchmark scores are CEVAL 82.2, CMMLU 79.5, MMLU 77.4, GSM8K 76.1, MATH 31.9 and HumanEval 43.9.

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

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,320 works

Cite

Cite as Tianwen Wei, Bo Zhu, Liang Zhao, Cheng Cheng, Biye Li, Weiwei L\"u, Peng Cheng, Jianhao Zhang, Xiaoyu Zhang, Liang Zeng, Xiaokun Wang, Yutuan Ma, Rui Hu, Shuicheng Yan, Han Fang and Yahui Zhou. 2024. Skywork-MoE: A Deep Dive into Training Techniques for Mixture-of-Experts Language Models. arXiv:2406.06563.

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

@misc{wei2024skywork,
  title = {{Skywork-MoE: A Deep Dive into Training Techniques for Mixture-of-Experts Language Models}},
  author = {Tianwen Wei, Bo Zhu, Liang Zhao, Cheng Cheng, Biye Li, Weiwei L"u, Peng Cheng, Jianhao Zhang, Xiaoyu Zhang, Liang Zeng, Xiaokun Wang, Yutuan Ma, Rui Hu, Shuicheng Yan, Han Fang and Yahui Zhou},
  year = {2024},
  doi = {10.48550/arxiv.2406.06563},
  eprint = {2406.06563},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2406.06563}
}

References

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

  1. Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. 2024. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language mo… find
  2. Anthropic. 2024. The claude 3 model family: Opus, sonnet, haiku. find
  3. Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Ott, Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pasunuru, Giri Anantharaman, Xian Li, Shuohui Chen, Halil Akin, Mandeep Baines, Louis Martin, Xing Zhou, Punit Singh Koura, Brian O’Horo, Jeff Wang, Luke Zettle… find
  4. Databricks. 2024. Dbrx-instruct. find
  5. DeepSeek-AI. 2024a. Deepseek llm: Scaling open-source language models with longtermism. find
  6. DeepSeek-AI. 2024b. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. find
  7. Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, Harsha Nori, Hamid Palangi, Marco Tulio Ribeiro, and Yi Zhang. 2023. Sparks of artificial general intelligence: Early experiments with gpt-4.
  8. Nan Du, Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma, Zongwei Zhou, Tao Wang, Yu Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathleen Meier-Hellstern, Toju Duke, Lucas… — GLaM: Efficient Scaling of Language Models with Mixture-of-Experts
  9. Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhai… — Evaluating Large Language Models Trained on Code
  10. David Eigen, Marc’Aurelio Ranzato, and Ilya Sutskever. 2014. Learning factored representations in a deep mixture of experts. find
  11. William Fedus, Barret Zoph, and Noam Shazeer. 2022. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. find
  12. Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2021a. Measuring massive multitask language understanding.
  13. Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. 2021b. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03874.
  14. Aidan Clark, Diego de las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, George van den Driessche, Eliza Rutherford, Tom Hennigan, Matthew Johnson, Katie Millican, Albin Cassirer, Chris Jones, Elena Buchatskaya, Dav… find
  15. Yuzhen Huang, Yuzhuo Bai, Zhihao Zhu, Junlei Zhang, Jinghan Zhang, Tangjun Su, Junteng Liu, Chuancheng Lv, Yikai Zhang, Jiayi Lei, Yao Fu, Maosong Sun, and Junxian He. 2023. Ceval: A multi-level multi-discipline chinese evaluation suite for foundation models. arXiv preprint arXiv:2305.08322.
  16. Robert Jacobs, Michael Jordan, Steven Nowlan, and Geoffrey Hinton. 1991. Adaptive mixture of local expert. Neural Computation, 3:78–88. find
  17. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. 2021. Training verifiers to solve math word problems.
  18. Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. 2023. Sparse upcycling: Training mixture-of-experts from dense checkpoints. find
  19. Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2020. Gshard: Scaling giant models with conditional computation and automatic sharding. find
  20. Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. 2022. Roformer: Enhanced transformer with rotary position embedding.
  21. Gemini Team. 2024. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context.
  22. Haonan Li, Yixuan Zhang, Fajri Koto, Yifei Yang, Hai Zhao, Yeyun Gong, Nan Duan, and Timothy Baldwin. 2023. Cmmlu: Measuring massive multitask language understanding in chinese.
  23. Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023a. Llama: Open and efficient foundation language models.
  24. Meta-AI. 2024. Llama 3 model card. find
  25. Mistral-AI. 2023. Mixtral-8x7B-v0.1. find
  26. Mistral-AI. 2024. Mixtral-8x22B-v0.1. find
  27. Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Anand Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. 2021. Efficient large-scale language model training on gpu clusters using megatr… find
  28. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. 2023b. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288.
  29. OpenAI. 2023. GPT-4 technical report.
  30. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS’17, page 6000–6010, Red Hook, NY, USA. Cur…
  31. Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. 2022. Train…
  32. Tianwen Wei, Liang Zhao, Lichang Zhang, Bo Zhu, Lijie Wang, Haihua Yang, Biye Li, Cheng Cheng, Weiwei Lü, Rui Hu, Chenxia Li, Liu Yang, Xilin Luo, Xuejie Wu, Lunan Liu, Wenjun Cheng, Peng Cheng, Jianhao Zhang, Xiaoyu Zhang, Lei Lin, Xiaokun Wang, Yutuan Ma, Chuanhai Dong, Yanqi Sun, Yifu Chen, Yongy… — Skywork: A More Open Bilingual Foundation Model
  33. Qwen Team. 2023. QWEN technical report. https://github.com/QwenLM/Qwen . ↗
  34. Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. 2022. Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale.
  35. Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. Zero: Memory optimizations toward training trillion parameter models. find
  36. Biao Zhang and Rico Sennrich. 2019. Root Mean Square Layer Normalization. In Advances in Neural Information Processing Systems 32, Vancouver, Canada. find
  37. Noam Shazeer. 2020. Glu variants improve transformer. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/wei-2024-skywork-moe-deep-dive/