Elamur Research⌕
☰
home / papers / 2022 / Tutel: Adaptive Mixture-of-Experts at Scale

Tutel: Adaptive Mixture-of-Experts at Scale

Paper · 2022 · Changho Hwang, Wei Cui, Yifan Xiong et al.

Summary

Mixture-of-Experts (MoE) layers send each token to a few experts chosen by a gating function, so the load on each expert shifts from one iteration to the next. The authors measured shifts of up to 4.38× during one training run. Earlier frameworks such as Fairseq and DeepSpeed-style stacks fix their parallelism and pipelining in advance, and that fixed choice is often the wrong one for the current load. Tutel addresses this with one distribution layout for both expert weights and input data. Because that layout fits every candidate strategy, the system can change parallelism and pipelining degree at each iteration without moving tensors or altering the math (it stays consistent with GShard). The system also adds Flexible All-to-All, a two-dimensional hierarchical (2DH) All-to-All, and fast sparse encode/decode kernels. On A100 clusters a single MoE layer runs 4.96× faster on 16 GPUs and 5.75× faster on 2,048 GPUs than the prior state of the art. The authors used Tutel to build SwinV2-MoE, which trained 1.55× faster and ran inference 2.11× faster than Fairseq. It also beat dense Swin V2 models on ImageNet and COCO. The work matters because it makes large sparse models practical to train and serve.

Key points

  • Adaptive parallelism switching alone accounts for a 1.35×–14.57× speedup of a single MoE layer. With 128 GPUs, the MoE layer runs up to 3.11× faster than Fairseq.
  • Tutel is open source (github.com/microsoft/tutel) and has been integrated into Fairseq and DeepSpeed.
  • Strong load-balancing loss hurts accuracy. On ImageNet-22K with SwinV2-S, top-1 falls to 34.71% at weight 1.0, against 37.78% at weight 0.01.
  • SwinV2-MoE-S with 32 experts (top-1 routing, f=1.0) compared with dense SwinV2-S:
    • ImageNet-22K top-1: 37.4% vs 35.5%.
    • ImageNet-1K 5-shot: 75.9% vs 70.3%.
    • Compute is the same, at 6.76 GFLOPs.
  • In the experiments, 32 and 64 experts worked best. Batch prioritized routing (BPR) was crucial when the capacity factor was low. Top-1 routing gave the best balance of speed and accuracy.

Interesting ideas

  • I1. Choose a data layout that keeps every strategy open. Distributing weights and inputs in one shared layout makes switching execution strategies free at runtime. The same principle could help any system whose best configuration depends on the incoming data.
  • I2. Size buffers to what each step actually needs. Tutel picks the smallest capacity factor that drops no tokens, instead of a fixed worst-case upper bound. This removes wasted computation and the accuracy loss from dropped tokens, and it applies to any pipeline with uneven batch content.
  • I3. Freeze the expert layers when fine-tuning on a downstream task. On COCO, keeping the MoE layers fixed during fine-tuning gave +0.4 box AP over the dense model, while tuning them lost 1.7 AP. This is a cheap safeguard when transferring sparse models to smaller target datasets.
  • I4. Use different capacity factors for training and inference. Lowering the capacity factor at inference trades a small accuracy drop for higher throughput, with no retraining. This is useful when the serving budget is tighter than the training budget.
  • I5. Consider a cosine router with a temperature floor. Tutel offers a router that normalizes the input projection and the expert embeddings, with a learnable temperature clamped at 0.01. It matched the linear router in image classification. It may route more stably when feature scale or dimension changes, and concurrent work cited by the authors reports gains on cross-lingual language tasks.

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 4The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksComputer Vision and Pattern Recognition · 3,486 worksHardware and Architecture · 167 works

Cite

Cite as Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose and Prabhat Ram. 2022. Tutel: Adaptive Mixture-of-Experts at Scale. MLSys 2023.

The paper itself, not this page. BibTeX from the library's record of the work. arXiv 2206.03382

@inproceedings{hwang2022tutel,
  title = {{Tutel: Adaptive Mixture-of-Experts at Scale}},
  author = {Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose and Prabhat Ram},
  booktitle = {MLSys 2023},
  year = {2022},
  eprint = {2206.03382},
  archivePrefix = {arXiv},
  url = {https://arxiv.org/abs/2206.03382}
}

References

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

  1. Azure, M. NDm A100 v4-series - Azure Virtual Machines. https://docs.microsoft.com/en-us/azure/virtual-machines/ ndm-a100-v4-series, 2023. [Online; accessed Apr 2023]. ↗
  2. Bruck, J., Ho, C.-T., Kipnis, S., Upfal, E., and Weathersby, D. Efficient algorithms for all-to-all communications in multiport message-passing systems. IEEE Transactions on parallel and distributed systems, 8(11):1143–1156, 1997. find
  3. Chi, Z., Dong, L., Huang, S., Dai, D., Ma, S., Patra, B., Singhal, S., Bajaj, P., Song, X., and Wei, F. On the representation collapse of sparse mixture of experts. CoRR, abs/2204.09179, 2022. — [arXiv:2204.09179](https://arxiv.org/abs/2204.09179) find
  4. Clark, A., de Las Casas, D., Guy, A., Mensch, A., Paganini, M., Hoffmann, J., Damoc, B., Hechtman, B. A., Cai, T., Borgeaud, S., van den Driessche, G., Rutherford, E., Hennigan, T., Johnson, M. J., Cassirer, A., Jones, C., Buchatskaya, E., Budden, D., Sifre, L., Osindero, S., Vinyals, O., Ranzato, M… find
  5. Cowan, M., Maleki, S., Musuvathi, M., Saarikivi, O., and Xiong, Y. MSCCLang: Microsoft collective communication language. In Proceedings of the ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASP-LOS), 2023. find
  6. Du, N., Huang, Y., Dai, A. M., Tong, S., Lepikhin, D., Xu, Y., Krikun, M., Zhou, Y., Yu, A. W., Firat, O., Zoph, B., Fedus, L., Bosma, M. P., Zhou, Z., Wang, T., Wang, Y. E., Webster, K., Pellat, M., Robinson, K., Meier-Hellstern, K. S., Duke, T., Dixon, L., Zhang, K., Le, Q. V., Wu, Y., Chen, Z., a… — GLaM: Efficient Scaling of Language Models with Mixture-of-Experts
  7. Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1–39, 2022. find
  8. He, J., Zhai, J., Antunes, T., Wang, H., Luo, F., Shi, S., and Li, Q. Fastermoe: Modeling and optimizing training of large-scale dynamic pre-trained models. In Proceedings of the ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming (PPoPP), 2022. find
  9. Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, D., Chen, M. X., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., and Chen, Z. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In Proceedings of the Advances in Neural Information Processing Systems (NeurIPS), 2019. find
  10. Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. CoRR, abs/2001.08361, 2020.
  11. Kim, J., Dally, W. J., Scott, S., and Abts, D. Technology-driven, highly-scalable dragonfly topology. In Proceedings of the International Symposium on Computer Architecture (ISCA). IEEE, 2008. find
  12. Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z. Gshard: Scaling giant models with conditional computation and automatic sharding. In Proceedings of the International Conference on Learning Representations (ICLR), 2021. find
  13. Lewis, M., Bhosale, S., Dettmers, T., Goyal, N., and Zettlemoyer, L. BASE layers: Simplifying training of large, sparse models. In Meila, M. and Zhang, T. (eds.), Proceedings of the International Conference on Machine Learning (ICML), 2021.
  14. Lin, J., Yang, A., Bai, J., Zhou, C., Jiang, L., Jia, X., Wang, A., Zhang, J., Li, Y., Lin, W., Zhou, J., and Yang, H. M6-10T: A sharing-delinking paradigm for efficient multi-trillion parameter pretraining. CoRR, abs/2110.03888, 2021. find
  15. Lin, T.-Y., Maire, M., Belongie, S., Hays, J., Perona, P., Ramanan, D., Dollár, P., and Zitnick, C. L. Microsoft coco: Common objects in context. In European conference on computer vision, pp. 740–755. Springer, 2014.
  16. Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
  17. Liu, Z., Hu, H., Lin, Y., Yao, Z., Xie, Z., Wei, Y., Ning, J., Cao, Y., Zhang, Z., Dong, L., Wei, F., and Guo, B. Swin transformer v2: Scaling up capacity and resolution. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022. find
  18. Mellanox. RDMA and SHARP Plugins for NCCL Library. https://github.com/Mellanox/nccl-rdma-sharp-plugins, 2023. [Online; accessed Apr 2023]. ↗
  19. Microsoft. DeepSpeed. https://www.deepspeed.ai/, 2023. [Online; accessed Apr 2023]. ↗
  20. NVIDIA. What is LL128 Protocol? https://github.com/NVIDIA/nccl/issues/281, 2020a. [Online; accessed Apr 2023]. ↗
  21. NVIDIA. NVIDIA A100 Tensor Core GPU Architecture – Unprecedented Acceleration at Every Scale. Whitepaper, 2020b.
  22. NVIDIA. NVIDIA Collective Communications Library (NCCL). https://github.com/NVIDIA/nccl/tree/v2.10.3-1, 2023a. [Online; accessed Apr 2023]. ↗
  23. NVIDIA. Point-to-point communication – NCCL 2.10.3 documentation. https://docs.nvidia.com/deeplearning/nccl/archives/nccl_2103/user-guide/docs/usage/p2p.html, 2023b. [Online; accessed Apr 2023]. ↗
  24. NVIDIA. NCCL Tests. https://github.com/NVIDIA/nccl-tests, 2023c. [Online; accessed Apr 2023]. ↗
  25. NVIDIA. NVLink & NVSwitch: Fastest HPC Data Center Platform. https://www.nvidia.com/en-us/data-center/nvlink/, 2023d. [Online; accessed Apr 2023]. ↗
  26. Ott, M., Edunov, S., Baevski, A., Fan, A., Gross, S., Ng, N., Grangier, D., and Auli, M. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019: Demonstrations, 2019.
  27. Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al. Pytorch: An imperative style, high-performance deep learning library. Proceedings of the Advances in Neural Information Processing Systems (NeurIPS), 2019.
  28. Pjesivac-Grbovic, J. Towards automatic and adaptive optimizations of mpi collective operations. 2007. find
  29. Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. Zero: memory optimizations toward training trillion parameter models. In Cuicchi, C., Qualters, I., and Kramer, W. T. (eds.), Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC). IEEE/AC… find
  30. Rajbhandari, S., Li, C., Yao, Z., Zhang, M., Aminabadi, R. Y., Awan, A. A., Rasley, J., and He, Y. Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation AI scale. In Proceedings of the International Conference on Machine Learning (ICML), 2022.
  31. Riquelme, C., Puigcerver, J., Mustafa, B., Neumann, M., Jenatton, R., Pinto, A. S., Keysers, D., and Houlsby, N. Scaling vision with sparse mixture of experts. In Proceedings of the Neural Information Processing Systems (NeurIPS), 2021.
  32. Roller, S., Sukhbaatar, S., Szlam, A., and Weston, J. Hash layers for large sparse models. In Ranzato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W. (eds.), Proceedings of the Neural Information Processing Systems (NeurIPS), 2021. find
  33. Sergeev, A. and Del Balso, M. Horovod: fast and easy distributed deep learning in tensorflow. arXiv preprint arXiv:1802.05799, 2018.
  34. Sharir, O., Peleg, B., and Shoham, Y. The cost of training NLP models: A concise overview. CoRR, abs/2004.08900, 2020. — [arXiv:2004.08900](https://arxiv.org/abs/2004.08900) find
  35. Shazeer, N., 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 Proceedings of the International Conference on Learning Representations (ICLR), 2017. find
  36. Snir, M., Gropp, W., Otto, S., Huss-Lederman, S., Dongarra, J., and Walker, D. MPI–the Complete Reference: the MPI core, volume 1. MIT press, 1998. find
  37. Thakur, R. and Choudhary, A. All-to-all communication on meshes with wormhole routing. In Proceedings of 8th International Parallel Processing Symposium, pp. 561– 565. IEEE, 1994. find
  38. Wikipedia. Ternary Search. https://en.wikipedia.org/wiki/Ternary_search, 2023. [Online; accessed Apr 2023]. ↗
  39. Yang, A., Lin, J., Men, R., Zhou, C., Jiang, L., Jia, X., Wang, A., Zhang, J., Wang, J., Li, Y., Zhang, D., Lin, W., Qu, L., Zhou, J., and Yang, H. Exploring sparse expert models and beyond. CoRR, abs/2105.15082, 2021. find
  40. Zheng, L., Li, Z., Zhang, H., Zhuang, Y., Chen, Z., Huang, Y., Wang, Y., Xu, Y., Zhuo, D., Xing, E. P., Gonzalez, J. E., and Stoica, I. Alpa: Automating inter- and Intra-Operator parallelism for distributed deep learning. In Proceedings of the USENIX Symposium on Operating Systems Design and Impleme…

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/hwang-2022-tutel-adaptive-mixture-experts-scale/