2022

MegaBlocks: Efficient Sparse Training with Mixture-of-Experts

Gale, Trevor, Narayanan, Deepak, Young, Cliff et al.

Understand

We present MegaBlocks, a system for efficient Mixture-of-Experts (MoE) training on GPUs.

Our system is motivated by the limitations of current frameworks, which restrict the dynamic routing in MoE layers to satisfy the constraints of existing software and hardware.
These formulations force a tradeoff between model quality and hardware efficiency, as users must choose between dropping tokens from the computation or wasting computation and memory on padding.
To address these limitations, we reformulate MoE computation in terms of block-sparse operations and develop new block-sparse GPU kernels that efficiently handle the dynamism present in MoEs.

Built on

Parallel sparse matrix-vector and matrix-transpose-vector multiplication using compressed sparse blocks
Buluç, A., Fineman, J. T., Frigo, M., Gilbert, J. R., and Leiserson, C. E · 2009
Earlier work this paper cites.
Learning both weights and connections for efficient neural network
Han, S., Pool, J., Tran, J., and Dally, W. J · 2015
Earlier work this paper cites.
Tensor-matrix products with a compressed sparse tensor
Smith, S. and Karypis, G · 2015
Earlier work this paper cites.
Block-sparse GPU kernels
Gray, S., Radford, A., and Kingma, D. P · 2017
Earlier work this paper cites.
In-datacenter performance analysis of a tensor processing unit
Jouppi, N. P., Young, C., Patil, N., Patterson, D., Agrawal, G., Bajwa, R., Bates, S., Bhatia, S., Boden, N., Borchers, A., Boyle, R., Cantin, P.-l., Chao, C., Clark, C., Coriell, J., Daley, M., Dau, M., Dean, J., Gelb, B., Ghaemmaghami, T. V., Gottipati, R., Gulland, W., Hagmann, R., Ho, C. R., Hogberg, D., Hu, J., Hundt, R., Hurt, D., Ibarz, J., Jaffey, A., Jaworski, A., Kaplan, A., Khaitan, H., Killebrew, D., Koch, A., Kumar, N., Lacy, S., Laudon, J., Law, J., Le, D., Leary, C., Liu, Z., Lucke, K., Lundin, A., MacKean, G., Maggiore, A., Mahony, M., Miller, K., Nagarajan, R., Narayanaswami, R., Ni, R., Nix, K., Norrie, T., Omernick, M., Penukonda, N., Phelps, A., Ross, J., Ross, M., Salek, A., Samadiani, E., Severn, C., Sizikov, G., Snelham, M., Souter, J., Steinberg, D., Swing, A., Tan, M., Thorson, G., Tian, B., Toma, H., Tuttle, E., Vasudevan, V., Walter, R., Wang, W., Wilcox, E., and Yoon, D. H · 2017
Earlier work this paper cites.
Block-sparse recurrent neural networks
Original
Narang, S., Undersander, E., and Diamos, G. F · 2017
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.
Attention is all you need
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I · 2017
Earlier work this paper cites.
Efficient neural audio synthesis
Kalchbrenner, N., Elsen, E., Simonyan, K., Noury, S., Casagrande, N., Lockhart, E., Stimberg, F., van den Oord, A., Dieleman, S., and Kavukcuoglu, K · 2018
Earlier work this paper cites.
Hicoo: Hierarchical storage of sparse tensors
Li, J., Sun, J., and Vuduc, R · 2018
Earlier work this paper cites.
Mixed precision training
Micikevicius, P., Narang, S., Alben, J., Diamos, G. F., Elsen, E., García, D., Ginsburg, B., Houston, M., Kuchaiev, O., Venkatesh, G., and Wu, H · 2018
Earlier work this paper cites.
Generating long sequences with sparse transformers
Original
Child, R., Gray, S., Radford, A., and Sutskever, I · 2019
Earlier work this paper cites.
The state of sparsity in deep neural networks
Original
Gale, T., Elsen, E., and Hooker, S · 2019
Earlier work this paper cites.

Similar

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., Desmaison, A., Köpf, A., Yang, E. Z., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S · 2019
Cited alongside, same era.
Language models are unsupervised multitask learners
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I · 2019
Cited alongside, same era.
Megatron-lm: Training multi-billion parameter language models using model parallelism
Original
Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B · 2019
Cited alongside, same era.
Triton: an intermediate language and compiler for tiled neural network computations
Tillet, P., Kung, H., and Cox, D. D · 2019
Cited alongside, same era.
Balanced sparsity for efficient DNN inference on GPU
Yao, Z., Cao, S., Xiao, W., Zhang, C., and Nie, L · 2019
Cited alongside, same era.
Language models are few-shot learners
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D · 2020
Cited alongside, same era.
Fast sparse convnets
Elsen, E., Dukhan, M., Gale, T., and Simonyan, K · 2020
Cited alongside, same era.
Sparse GPU kernels for deep learning
Gale, T., Zaharia, M., Young, C., and Elsen, E · 2020
Cited alongside, same era.
The Pile: An 800GB dataset of diverse text for language modeling
Original
Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C · 2020
Cited alongside, same era.
Gshard: Scaling giant models with conditional computation and automatic sharding
Original
Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z · 2020
Cited alongside, same era.
NVIDIA A100 Tensor Core GPU Architecture
NVIDIA · 2020
Cited alongside, same era.
Efficient large scale language modeling with mixtures of experts
Original
Artetxe, M., Bhosale, S., Goyal, N., Mihaylov, T., Ott, M., Shleifer, S., Lin, X. V., Du, J., Iyer, S., Pasunuru, R., Anantharaman, G., Li, X., Chen, S., Akin, H., Baines, M., Martin, L., Zhou, X., Koura, P. S., O’Horo, B., Wang, J., Zettlemoyer, L., Diab, M. T., Kozareva, Z., and Stoyanov, V · 2021
Cited alongside, same era.

Then

GLaM: Efficient scaling of language models with mixture-of-experts, 2021
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., Zhou, Z., Wang, T., Wang, Y. E., Webster, K., Pellat, M., Robinson, K., Meier-Hellstern, K., Duke, T., Dixon, L., Zhang, K., Le, Q. V., Wu, Y., Chen, Z., and Cui, C · 2021
Later among the works it cites.
BASE layers: Simplifying training of large, sparse models
Lewis, M., Bhosale, S., Dettmers, T., Goyal, N., and Zettlemoyer, L · 2021
Later among the works it cites.
Efficient large-scale language model training on GPU clusters using megatron-lm
Narayanan, D., Shoeybi, M., Casper, J., LeGresley, P., Patwary, M., Korthikanti, V., Vainbrand, D., Kashinkunti, P., Bernauer, J., Catanzaro, B., Phanishayee, A., and Zaharia, M · 2021
Later among the works it cites.
Accelerating matrix multiplication with block sparse format and nvidia tensor cores
NVIDIA · 2021
Later among the works it cites.
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
Later among the works it cites.
Hash layers for large sparse models
Roller, S., Sukhbaatar, S., Szlam, A., and Weston, J · 2021
Later among the works it cites.
Unified scaling laws for routed language models
Original
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., Millican, K., Cassirer, A., Jones, C., Buchatskaya, E., Budden, D., Sifre, L., Osindero, S., Vinyals, O., Rae, J. W., Elsen, E., Kavukcuoglu, K., and Simonyan, K · 2022
Closest in time.
Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity
Fedus, W., Zoph, B., and Shazeer, N · 2022
Closest in time.
Fastermoe: Modeling and optimizing training of large-scale dynamic pre-trained models
He, J., Zhai, J., Antunes, T., Wang, H., Luo, F., Shi, S., and Li, Q · 2022
Closest in time.
Tutel: Adaptive mixture-of-experts at scale, 2022
Hwang, C., Cui, W., Xiong, Y., Yang, Z., Liu, Z., Hu, H., Wang, Z., Salas, R., Jose, J., Ram, P., Chau, J., Cheng, P., Yang, F., Yang, M., and Xiong, Y · 2022
Closest in time.
Reducing Activation Recomputation in Large Transformer Models
Original
Korthikanti, V., Casper, J., Lym, S., McAfee, L., Andersch, M., Shoeybi, M., and Catanzaro, B · 2022
Closest in time.
Mosaic LLMs (part 2): GPT-3 quality for <$500k
MosaicML · 2022
Closest in time.
Mixture-of-experts with expert choice routing
Original
Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V. Y., Dai, A. M., Chen, Z., Le, Q., and Laudon, J · 2022
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…