2022

Tutel: Adaptive Mixture-of-Experts at Scale

Hwang, Changho, Cui, Wei, Xiong, Yifan et al.

Understand

Sparsely-gated mixture-of-experts (MoE) has been widely adopted to scale deep learning models to trillion-plus parameters with fixed computational cost.

The algorithmic performance of MoE relies on its token routing mechanism that forwards each input token to the right sub-models or experts.
While token routing dynamically determines the amount of expert workload at runtime, existing systems suffer inefficient computation due to their static execution, namely static parallelism and pipelining, which does not adapt to the dynamic workload.
We present Flex, a highly scalable stack design and implementation for MoE with dynamically adaptive parallelism and pipelining.

Built on

All-to-all communication on meshes with wormhole routing
Thakur, R. and Choudhary, A · 1994
Earlier work this paper cites.
Efficient algorithms for all-to-all communications in multiport message-passing systems
Bruck, J., Ho, C.-T., Kipnis, S., Upfal, E., and Weathersby, D · 1997
Earlier work this paper cites.
MPI–the Complete Reference: the MPI core , volume 1
Snir, M., Gropp, W., Otto, S., Huss-Lederman, S., Dongarra, J., and Walker, D · 1998
Earlier work this paper cites.
Towards automatic and adaptive optimizations of mpi collective operations
Pjesivac-Grbovic, J · 2007
Earlier work this paper cites.
Technology-driven, highly-scalable dragonfly topology
Kim, J., Dally, W. J., Scott, S., and Abts, D · 2008
Earlier work this paper cites.
Microsoft coco: Common objects in context
Lin, T.-Y., Maire, M., Belongie, S., Hays, J., Perona, P., Ramanan, D., Dollár, P., and Zitnick, C. L · 2014
Earlier work this paper cites.
Outrageously large neural networks: The sparsely-gated mixture-of-experts layer
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q. V., Hinton, G. E., and Dean, J · 2017
Earlier work this paper cites.
Horovod: fast and easy distributed deep learning in tensorflow
Original
Sergeev, A. and Del Balso, M · 2018
Earlier work this paper cites.
Gpipe: Efficient training of giant neural networks using pipeline parallelism
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 · 2019
Earlier work this paper cites.
fairseq: A fast, extensible toolkit for sequence modeling
Ott, M., Edunov, S., Baevski, A., Fan, A., Gross, S., Ng, N., Grangier, D., and Auli, M · 2019
Earlier work this paper cites.
Pytorch: An imperative style, high-performance deep learning library
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al · 2019
Earlier work this paper cites.
Scaling laws for neural language models
Original
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D · 2020
Earlier work this paper cites.
Zero: memory optimizations toward training trillion parameter models
Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y · 2020
Earlier work this paper cites.

Similar

The cost of training NLP models: A concise overview
Original
Sharir, O., Peleg, B., and Shoham, Y · 2020
Cited alongside, same era.
Gshard: Scaling giant models with conditional computation and automatic sharding
Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z · 2021
Cited alongside, same era.
BASE layers: Simplifying training of large, sparse models
Lewis, M., Bhosale, S., Dettmers, T., Goyal, N., and Zettlemoyer, L · 2021
Cited alongside, same era.
M6-10T: A sharing-delinking paradigm for efficient multi-trillion parameter pretraining
Original
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 · 2021
Cited alongside, same era.
Swin transformer: Hierarchical vision transformer using shifted windows
Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B · 2021
Cited alongside, same era.
Scaling vision with sparse mixture of experts
Riquelme, C., Puigcerver, J., Mustafa, B., Neumann, M., Jenatton, R., Pinto, A. S., Keysers, D., and Houlsby, N · 2021
Cited alongside, same era.
Hash layers for large sparse models
Roller, S., Sukhbaatar, S., Szlam, A., and Weston, J · 2021
Cited alongside, same era.
Exploring sparse expert models and beyond
Original
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 · 2021
Cited alongside, same era.
On the representation collapse of sparse mixture of experts
Original
Chi, Z., Dong, L., Huang, S., Dai, D., Ma, S., Patra, B., Singhal, S., Bajaj, P., Song, X., and Wei, F · 2022
Cited alongside, same era.
Unified scaling laws for routed language models
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., Rae, J. W., Elsen, E., Kavukcuoglu, K., and Simonyan, K · 2022
Cited alongside, same era.
Glam: Efficient scaling of language models with mixture-of-experts
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., and Cui, C · 2022
Cited alongside, same era.
Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity
Fedus, W., Zoph, B., and Shazeer, N · 2022
Cited alongside, same era.

Then

Swin transformer v2: Scaling up capacity and resolution
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 · 2022
Closest in time.
Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation AI scale
Rajbhandari, S., Li, C., Yao, Z., Zhang, M., Aminabadi, R. Y., Awan, A. A., Rasley, J., and He, Y · 2022
Closest in time.
Alpa: Automating inter- and Intra-Operator parallelism for distributed deep learning
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 · 2022
Closest in time.
NDm A100 v4-series - Azure Virtual Machines
Azure, M · 2023
Closest in time.
MSCCLang: Microsoft collective communication language
Cowan, M., Maleki, S., Musuvathi, M., Saarikivi, O., and Xiong, Y · 2023
Closest in time.
RDMA and SHARP Plugins for NCCL Library
Mellanox · 2023
Closest in time.
DeepSpeed
Microsoft · 2023
Closest in time.
What is LL128 Protocol?
NVIDIA · 2023
Closest in time.
NVIDIA Collective Communications Library (NCCL)
NVIDIA · 2023
Closest in time.
Point-to-point communication – NCCL 2.10.3 documentation
NVIDIA · 2023
Closest in time.
NCCL Tests
NVIDIA · 2023
Closest in time.
NVLink & NVSwitch: Fastest HPC Data Center Platform
NVIDIA · 2023
Closest in time.
Ternary Search
Wikipedia · 2023
Closest in time.

Beyond the bibliography

alphaXiv searches the wider corpus for related work and actual follow-ups.

Open on alphaXiv

alphaXiv is searching for related work…

Reading the bibliography…