Proceedings

LocMoE: A Low-overhead MoE for Large Language Model Training

Download BibTeX Cite

Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence. Main Track. Pages 6377-6387.

https://doi.org/10.24963/ijcai.2024/705

Abstract

The Mixtures-of-Experts (MoE) model is a widespread distributed and integrated learning method for large language models (LLM), which is favored due to its ability to sparsify and expand models efficiently. However, the performance of MoE is limited by load imbalance and high latency of All-to-All communication, along with relatively redundant computation owing to large expert capacity. Load imbalance may result from existing routing policies that consistently tend to select certain experts. The frequent inter-node communication in the All-to-All procedure also significantly prolongs the training time. To alleviate the above performance problems, we propose a novel routing strategy that combines load balance and locality by converting partial inter-node communication to that of intra-node. Notably, we elucidate that there is a minimum threshold for expert capacity, calculated through the maximal angular deviation between the gating weights of the experts and the assigned tokens. We port these modifications on the PanGu-Σ model based on the MindSpore framework with multi-level routing and conduct experiments on Ascend clusters. The experiment results demonstrate that the proposed LocMoE reduces training time per epoch by 12.68% to 22.24% compared to classical routers, such as hash router and switch router, without impacting the model accuracy.

  1. [Bi et al., 2023] Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian. Accurate medium- range global weather forecasting with 3d neural networks. Nature, 619(7970):533-538, 2023.
  2. [Brown et al., 2020] Tom Brown, Benjamin Mann, Nick Ry- der, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing sys- tems, 33:1877-1901, 2020.
  3. [Chi et al., 2022] Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Xian-Ling Mao, et al. On the repre- sentation collapse of sparse mixture of experts. Advances in Neural Information Processing Systems, 35:34600- 34613, 2022.
  4. [Chowdhury et al., 2023] Mohammed Nowaz Rabbani Chowdhury, Shuai Zhang, Meng Wang, Sijia Liu, and Pin-Yu Chen. Patch-level routing in mixture-of-experts is provably sample-efficient for convolutional neural networks. arXiv preprint arXiv:2306.04073, 2023.
  5. [Clark et al., 2022] Aidan Clark, Diego De Las Casas, Aure- lia Guy, Arthur Mensch, Michela Paganini, Jordan Hoff- mann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Se- bastian Borgeaud, et al. Unified scaling laws for routed language models. In International Conference on Machine Learning, pages 4057-4086. PMLR, 2022.
  6. [Dai et al., 2022] Damai Dai, Li Dong, Shuming Ma, Bo Zheng, Zhifang Sui, Baobao Chang, and Furu Wei. Sta- blemoe: Stable routing strategy for mixture of experts. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 7085-7095, 2022.
  7. [Fedus et al., 2022] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parame- ter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1):5232-5270, 2022.
  8. [Guo et al., 2020] Yuanyuan Guo, Yifan Xia, Jing Wang, Hui Yu, and Rung-Ching Chen. Real-time facial affective computing on mobile devices. Sensors, 20(3):870, 2020.
  9. [He et al., 2021] Jiaao He, Jiezhong Qiu, Aohan Zeng, Zhilin Yang, Jidong Zhai, and Jie Tang. Fastmoe: A fast mixture-of-expert training system. arXiv preprint arXiv:2103.13262, 2021.
  10. [Jacobs et al., 1991] Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mix- tures of local experts. Neural computation, 3(1):79-87, 1991.
  11. [Kenton and Toutanova, 2019] Jacob Devlin Ming- Wei Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for lan- guage understanding. In Proceedings of NAACL-HLT, pages 4171-4186, 2019.
  12. [Kudugunta et al., 2021] Sneha Kudugunta, Yanping Huang, Ankur Bapna, Maxim Krikun, Dmitry Lepikhin, Minh- Thang Luong, and Orhan Firat. Beyond distillation: Task- level mixture-of-experts for efficient inference. In Find- ings of the Association for Computational Linguistics: EMNLP 2021, pages 3577-3599, 2021.
  13. [Kumar et al., 2022] N Ashok Kumar, A Kavitha, P Venka- tramana, and Durgesh Nandan. Architecture design: Network-on-chip. In VLSI Architecture for Signal, Speech, and Image Processing, pages 147-165. Apple Academic Press, 2022.
  14. [Lepikhin et al., 2020] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computa- tion and automatic sharding. In International Conference on Learning Representations, 2020.
  15. [Lewis et al., 2021] Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base lay- ers: Simplifying training of large, sparse models. In Inter- national Conference on Machine Learning, pages 6265- 6274. PMLR, 2021.
  16. [Li et al., 2022a] Jianjun Li, Yu Han, Ming Zhang, Gang Li, and Baohua Zhang. Multi-scale residual network model combined with global average pooling for action recog- nition. Multimedia Tools and Applications, pages 1-19, 2022.
  17. [Li et al., 2022b] Yuan Li, Ke Wang, Hao Zheng, Ahmed Louri, and Avinash Karanth. Ascend: A scalable and energy-efficient deep neural network accelerator with pho- tonic interconnects. IEEE Transactions on Circuits and Systems I: Regular Papers, 69(7):2730-2741, 2022.
  18. [Liao et al., 2021] Heng Liao, Jiajin Tu, Jing Xia, Hu Liu, Xiping Zhou, Honghui Yuan, and Yuxing Hu. As- cend: a scalable and unified architecture for ubiqui- tous deep neural network computing: Industry track pa- per. In 2021 IEEE International Symposium on High- Performance Computer Architecture (HPCA), pages 789- 801. IEEE, 2021.
  19. [Mi et al., 2022] Fei Mi, Yitong Li, Yulong Zeng, Jingyan Zhou, Yasheng Wang, Chuanfei Xu, Lifeng Shang, Xin Jiang, Shiqi Zhao, and Qun Liu. Pangu-bot: Efficient generative dialogue pre-training from pre-trained language model. arXiv preprint arXiv:2203.17090, 2022.
  20. [Nie et al., 2022] Xiaonan Nie, Pinxue Zhao, Xupeng Miao, Tong Zhao, and Bin Cui. Hetumoe: An efficient trillion- scale mixture-of-expert distributed training system. arXiv preprint arXiv:2203.14685, 2022.
  21. [Nie et al., 2023] Xiaonan Nie, Xupeng Miao, Zilong Wang, Zichao Yang, Jilong Xue, Lingxiao Ma, Gang Cao, and Bin Cui. Flexmoe: Scaling large-scale sparse pre-trained model training via dynamic device placement. Proceed- ings of the ACM on Management of Data, 1(1):1-19, 2023.
  22. [Puigcerver et al., 2020] Joan Puigcerver, Carlos Riquelme Ruiz, Basil Mustafa, Cedric Renggli, Andr 'e Susano Pinto, Sylvain Gelly, Daniel Keysers, and Neil Houlsby. Scalable transfer learning with expert models. In Interna- tional Conference on Learning Representations, 2020.
  23. [Rajbhandari et al., 2022] Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale. In Inter- national Conference on Machine Learning, pages 18332- 18346. PMLR, 2022.
  24. [Ren et al., 2023] Xiaozhe Ren, Pingyi Zhou, Xinfan Meng, Xinjing Huang, Yadao Wang, Weichao Wang, Pengfei Li, Xiaoda Zhang, Alexander Podolskiy, Grigory Arshinov, et al. Pangu-σ: Towards trillion parameter language model with sparse heterogeneous computing. arXiv e-prints, pages arXiv-2303, 2023.
  25. [Roller et al., 2021] Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, et al. Hash layers for large sparse mod- els. Advances in Neural Information Processing Systems, 34:17555-17566, 2021.
  26. [Shazeer et al., 2016] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hin- ton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In Interna- tional Conference on Learning Representations, 2016.
  27. [Shen et al., 2023] Bo Shen, Jiaxin Zhang, Taihong Chen, Daoguang Zan, Bing Geng, An Fu, Muhan Zeng, Ailun Yu, Jichuan Ji, Jingyang Zhao, et al. Pangu-coder2: Boost- ing large language models for code with ranking feedback. arXiv preprint arXiv:2307.14936, 2023.
  28. [Tong et al., 2021] Zhihao Tong, Ning Du, Xiaobo Song, and Xiaoli Wang. Study on mindspore deep learn- ing framework. In 2021 17th International Conference on Computational Intelligence and Security (CIS), pages 183-186. IEEE, 2021.
  29. [Touvron et al., 2023] Hugo Touvron, Thibaut Lavril, Gau- tier Izacard, Xavier Martinet, Marie-Anne Lachaux, Tim- othée Lacroix, Baptiste Rozière, Naman Goyal, Eric Ham- bro, Faisal Azhar, et al. Llama: Open and efficient founda- tion language models. arXiv preprint arXiv:2302.13971, 2023.
  30. [Vaswani et al., 2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
  31. [Wang et al., 2023] Yunhe Wang, Hanting Chen, Yehui Tang, Tianyu Guo, Kai Han, Ying Nie, Xutao Wang, Hailin Hu, Zheyuan Bai, Yun Wang, et al. Pangu-π: Enhancing language model architectures via nonlinearity compensation. arXiv preprint arXiv:2312.17276, 2023.
  32. [Xia et al., 2021] Jing Xia, Chuanning Cheng, Xiping Zhou, Yuxing Hu, and Peter Chun. Kunpeng 920: The first 7-nm chiplet-based 64-core arm soc for cloud services. IEEE Micro, 41(5):67-75, 2021.
  33. [Zeng et al., 2021] Wei Zeng, Xiaozhe Ren, Teng Su, Hui Wang, Yi Liao, Zhiwei Wang, Xin Jiang, ZhenZhang Yang, Kaisheng Wang, Xiaoda Zhang, et al. Pangu- α: Large-scale autoregressive pretrained chinese language models with auto-parallel computation. arXiv preprint arXiv:2104.12369, 2021.
  34. [Zuo et al., 2021] Simiao Zuo, Xiaodong Liu, Jian Jiao, Young Jin Kim, Hany Hassan, Ruofei Zhang, Jianfeng Gao, and Tuo Zhao. Taming sparsely activated transformer with stochastic experts. In International Conference on Learning Representations, 2021.

References were extracted automatically from the PDF and may contain errors.

Welcome Future Conferences AI Journal Division IJCAI Proceedings About IJCAI IJCAI Awards