Elamur Research⌕
☰
home / papers / 2022 / Accelerating Distributed MoE Training and Inference with Lina

Accelerating Distributed MoE Training and Inference with Lina

Paper · 2022 · Jiamin Li, Yimin Jiang, Yibo Zhu et al.

Summary

Mixture-of-Experts (MoE) models keep per-token compute low by sending each token to only one or two expert FFNs. In distributed setups that mix data and expert parallelism, however, much of the runtime goes to all-to-all exchanges while GPUs sit nearly idle. The authors show the cause differs by phase. In training, all-to-all competes for bandwidth with the allreduce of non-expert gradients during the backward pass. In inference, uneven expert popularity overloads some devices and unbalances transfers.

Lina, built on DeepSpeed MoE and PyTorch, handles both cases:

  • Training: it cuts tensors into small micro-ops and schedules them by priority, so all-to-all gets the whole link. Allreduce runs in the gaps, and expert computation overlaps with communication.
  • Inference: it predicts the next layer's expert load from tokens' earlier routing choices and allocates resources in advance. It corrects the plan only when the real load diverges strongly.

On up to 16 A100 GPUs, training steps run 1.57x faster on average than DeepSpeed, and inference latency drops clearly. The work shows that communication, more than FLOPs, limits practical MoE scaling.

Key points

  • Profiling on Transformer-XL with DeepSpeed MoE showed how costly all-to-all is:
    • It takes about 34.1% of step time on average and 74.9% of one MoE layer's forward pass.
    • GPU SM efficiency during all-to-all is only about 3.7%.
  • The testbed had up to 16 A100 GPUs (40GB) connected by 100Gbps InfiniBand. The baseline was DeepSpeed MoE.
  • Training results: all-to-all ran at least 2.21x faster, and the overall training step was on average 1.57x faster.
  • Inference results: median latency fell by 1.45x and 95th-percentile latency by 1.63x.
  • Popularity estimation accuracy per MoE layer was 58.41% (Transformer-XL) and 54.16% (BERT-Large). Across four extra tasks it reached at least 62.3%, with tail latency within 1.04x of an Ideal scheduler.

Interesting ideas

  • I1. Prioritize the blocking collective. When a blocking operation and a non-blocking one share the network, give the blocking one strict priority. Splitting work into small chunks makes this possible without starving the other traffic. The idea applies to any pipeline where one transfer stalls computation and another can overlap with it.
  • I2. Chunk tensors so communication and compute overlap. Breaking a tensor into micro-ops lets the experts start computing on chunks that have already arrived. This shortens idle periods and is a generic trick for communication-bound layers.
  • I3. Predict routing from earlier layers. Tokens follow recognizable expert paths across layers, so their recent choices predict which experts the next layer will need. This lets resources be planned before the gate actually decides. Using a short history (default 3 layers) worked well, and a longer one gave little extra speed.
  • I4. Two-phase scheduling. Act on a cheap estimate first, and fine-tune only when actual demand departs far from it. This keeps overhead low while protecting against bad predictions, a pattern useful for any load balancer driven by forecasts.
  • I5. Diagnose per phase before optimizing. The same bottleneck (all-to-all) had different root causes in training and in inference, which called for two separate fixes. Profiling each phase separately is a good evaluation habit.

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

 Institutions

🎓Chinese University of Hong KongHong Kong · China159 works · 16,154 citations🏢ByteDanceBeijing · China47 works · 4,339 citations🎓City University of Hong KongHong Kong · China59 works · 3,480 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksComputer Vision and Pattern Recognition · 3,486 worksInformation Systems · 1,730 worksHardware and Architecture · 167 works

Cite

Cite as Jiamin Li, Yimin Jiang, Yibo Zhu, Cong Wang and Hong Xu. 2022. Accelerating Distributed MoE Training and Inference with Lina. arXiv:2210.17223.

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

@misc{li2022accelerating,
  title = {{Accelerating Distributed MoE Training and Inference with Lina}},
  author = {Jiamin Li, Yimin Jiang, Yibo Zhu, Cong Wang and Hong Xu},
  year = {2022},
  doi = {10.48550/arxiv.2210.17223},
  eprint = {2210.17223},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2210.17223}
}

References

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

  1. [1] ChatGPT: Optimizing Language Models for Dialogue. — Efficient Large Scale Language Modeling with Mixtures of Experts https://openai.com/blog/chatgpt/. [2] DeepSpeed. DeepSpeed. https://github.com/microsoft/[3] Enwik8. http://prize.hutter1.net/. [4] NCCL. https://github.com/NVIDIA/nccl. https://[5] PyTorch Distributed Data Parallel. pytorch.org/docs/stable/note… ↗
  2. [10] Mandeep Baines, Shruti Bhosale, Vittorio Caggiano, Naman Goyal, Siddharth Goyal, Myle Ott, Benjamin Lefaudeux, Vitaliy Liptchinsky, Mike Rabbat, Sam Sheiffer, Anjali Sridhar, and Min Xu. Fairscale: A general purpose modular pytorch library for high performance and large scale training. https://… find
  3. [12] Paul Barham, Aakanksha Chowdhery, Jeff Dean, Sanjay Ghemawat, Steven Hand, Daniel Hurt, Michael Isard, Hyeontaek Lim, Ruoming Pang, Sudip Roy, Brennan Saeta, Parker Schuh, Ryan Sepassi, Laurent Shafey, Chandu Thekkath, and Yonghui Wu. Pathways: Asynchronous distributed dataflow for ml. In Proc.… find
  4. [13] Emmanuel Bengio, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup. Conditional computation in neural networks for faster models. arXiv preprint arXiv:1511.06297, 2015. — [arXiv:1511.06297](https://arxiv.org/abs/1511.06297) find
  5. [14] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013. find
  6. [19] Minsik Cho, Ulrich Finkler, David Kung, and Hillery Hunter. Blueconnect: Decomposing all-reduce for deep learning on heterogeneous network hierarchy. In Proc. MLSys, 2019. find
  7. [20] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov. Transformer-XL: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901.02860, 2019. find
  8. [21] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
  9. [22] 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 P Bosma, Zongwei Zhou, Tao Wang, Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathleen Meier-Hellstern, Toju Duke, Lu… — GLaM: Efficient Scaling of Language Models with Mixture-of-Experts
  10. [23] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.03961, 2021. find
  11. [24] Sayed Hadi Hashemi, Sangeetha Abdu Jyothi, and Roy Campbell. Tictac: Accelerating distributed deep learning with communication scheduling. In A. Talwalkar, V. Smith, and M. Zaharia, editors, Proc. MLSys, 2019. find
  12. [15] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.
  13. [25] Jiaao He, Jidong Zhai, Tiago Antunes, Haojie Wang, Fuwen Luo, Shangfeng Shi, and Qin Li. FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models. In Proc. ACM SIGPLAN PPoPP, pages 120–134, 2022. find
  14. [16] Tianlong Chen, Zhenyu Zhang, AJAY KUMAR JAISWAL, Shiwei Liu, and Zhangyang Wang. Sparse moe as the new dropout: Scaling dense and self-slimmable transformers. In Proc. ICLR, 2023. find
  15. [26] Xianyan Jia, Le Jiang, Ang Wang, Jie Zhang, Xinyuan Li, Wencong Xiao, Yong Li, Zhen Zheng, Xiaoyong Liu, Wei Lin, et al. Whale: Scaling deep learning model training to the trillions. arXiv preprint arXiv:2011.09208, 2020. find
  16. [17] Yangrui Chen, Yanghua Peng, Yixin Bao, Chuan Wu, Yibo Zhu, and Chuanxiong Guo. Elastic parameter server load distribution in deep learning clusters. In Proc. ACM SoCC, 2020. find
  17. [27] Zhihao Jia, Matei Zaharia, and Alex Aiken. Beyond data and model parallelism for deep neural networks. Proc. MLSys, 2019.
  18. [18] 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 representation collapse of sparse mixture of experts. Proc. NeurIPS, 2022. find
  19. [28] Yimin Jiang, Yibo Zhu, Chang Lan, Bairen Yi, Yong Cui, and Chuanxiong Guo. A unified architecture for accelerating distributed DNN training in heterogeneous GPU/CPU clusters. In Proc. USENIX OSDI, pages 463– 479, 2020. find
  20. [29] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
  21. [40] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 2020.
  22. [30] Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upcycling: Training mixture-of-experts from dense checkpoints. In Proc. ICLR, 2023.
  23. [41] 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. arXiv preprint arXiv:2201.05596, 2022.
  24. [31] 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 computation and automatic sharding. arXiv preprint arXiv:2006.16668, 2020. find
  25. [32] Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In Proc. USENIX ICML, 2021.
  26. [33] Hanxue Liang, Zhiwen Fan, Rishov Sarkar, Ziyu Jiang, Tianlong Chen, Kai Zou, Yu Cheng, Cong Hao, and Zhangyang Wang. M³vit: Mixture-of-experts vision transformer for efficient multi-task learning with model-accelerator co-design. In Proc. NeurIPS, 2022. find
  27. [34] Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proc. ACL, 2011.
  28. [35] Ibrahim Naji. TSATC: Twitter Sentiment Analysis Training Corpus. In thinknook, 2012. find
  29. [36] Kay Ousterhout, Aurojit Panda, Joshua Rosen, Shivaram Venkataraman, Reynold Xin, Sylvia Ratnasamy, Scott Shenker, and Ion Stoica. The Case for Tiny Tasks in Compute Clusters. In Proc. USENIX HotOS, 2013. find
  30. [37] Yanghua Peng, Yibo Zhu, Yangrui Chen, Yixin Bao, Bairen Yi, Chang Lan, Chuan Wu, and Chuanxiong Guo. A generic communication scheduler for distributed dnn training acceleration. In Proc. ACM SOSP, 2019. find
  31. [38] Andrey Proskurin. DeepSpeed: Advancing MoE inference and training to power next-generation AI scale. https://www.microsoft.com/en-us/research/ blog/deepspeed-advancing-moe-inference-and-training-to-power-next-generation-ai-scale. [39] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei… ↗
  32. [42] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. ZeRO-Offload: Democratizing Billion-Scale Model Training. In Proc. USENIX ATC, 2021.
  33. [43] Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, Ryan Sepassi, and Blake Hechtman. Mesh-TensorFlow: Deep learning for supercomputers. In Proc. ACM NeurIPS, 2018.
  34. [44] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538, 2017. find
  35. [45] Liang Shen, Zhihua Wu, WeiBao Gong, Hongxiang Hao, Yangfan Bai, HuaChao Wu, Xinxuan Wu, Haoyi Xiong, Dianhai Yu, and Yanjun Ma. Se-moe: A scalable and efficient mixture-of-experts distributed training and inference system. arXiv preprint arXiv:2205.10034, 2022. find
  36. [46] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019. find
  37. [47] 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.
  38. [48] Guanhua Wang, Shivaram Venkataraman, Amar Phanishayee, Nikhil Devanur, Jorgen Thelin, and Ion Stoica. Blink: Fast and generic collectives for distributed ml. In Proc. MLSys, 2020. find
  39. [49] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin… — Transformers: State-of-the-Art Natural Language Processing
  40. [50] Zhewei Yao, Xiaoxia Wu, Conglong Li, Connor Holmes, Minjia Zhang, Cheng Li, and Yuxiong He. Random-ltd: Random and layerwise token dropping brings efficient training for large-scale transformers. arXiv preprint arXiv:2211.11586, 2022. — [arXiv:2211.11586](https://arxiv.org/abs/2211.11586) find
  41. [51] Hao Zhang, Zeyu Zheng, Shizhen Xu, Wei Dai, Qirong Ho, Xiaodan Liang, Zhiting Hu, Jinliang Wei, Pengtao Xie, and Eric P Xing. Poseidon: An efficient communication architecture for distributed deep learning on GPU clusters. In Proc. USENIX ATC, 2017. find
  42. [52] Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Eric P Xing, et al. Alpa: Automating inter-and intra-operator parallelism for distributed deep learning. In Proc. USENIX OSDI, 2022. — Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning
  43. [53] Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. Mixture-of-experts with expert choice routing. Proc. NeurIPS, 2022. find
  44. [54] Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. Designing Effective Sparse Expert Models. arXiv preprint arXiv:2202.08906, 2022. find
  45. [55] Simiao Zuo, Xiaodong Liu, Jian Jiao, Young Jin Kim, Hany Hassan, Ruofei Zhang, Tuo Zhao, and Jianfeng Gao. Taming sparsely activated transformer with stochastic experts. arXiv preprint arXiv:2110.04260, 2021. — [arXiv:2110.04260](https://arxiv.org/abs/2110.04260) find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/li-2022-accelerating-distributed-moe-training-inference/