Who Says Elephants Can’t Run: Bringing Large Scale M o E Models into Cloud Scale Production (2022)

Abstract: Mixture of Experts (MoE) models with conditional execution of sparsely activated layers has enabled training models with a much larger number of parameters. As a result, these models have achieved significantly better quality on various natural language processing tasks including machine translation. However, it remains challenging to deploy such models in real-life scenarios due to the large memory requirements and inefficient inference. In this work, we introduce a highly efficient inference framework with several optimization approaches to accelerate the computation of sparse models and cut down the memory consumption significantly. While we achieve up to 26x speed-up in terms of throughput, we also reduce the model size almost to one eighth of the original 32-bit float model by quantizing expert weights into 4-bit integers. As a result, we are able to deploy 136x larger models with 27% less cost and significantly better quality with large scale MoE model deployment compared to the existing solutions. This enables a paradigm shift in deploying large scale multilingual MoE transformers models instead of distilling into dozens of smaller models per language or task.

## introduction transformer models are getting larger and better on a continuous basis. the largest transformer models scale up to hundreds of billions of parameters, @xcite resulting in high training and inference costs. this makes it difficult to deploy such models in any real-life scenario with reasonable latency and throughput. mixture of experts (moe) models offer a more cost-effective method to scaling model sizes by using sparsely activated computations. more specifically, feed forward layers can be easily enlarged by replicating the original weights e times where e is the number of experts. each of these replicas is referred to as an expert, and tokens get routed to these experts depending on a gating function. transformer models have a much larger number of parameters when utilizing these moe layers. however, the number of flops remains comparable to their dense counterparts thanks to sub-linear scaling in computation costs @xcite . recently, the mixture of experts (moe) architecture has been successfully utilized to scale massive large scale multilingual models @xcite ), nlu tasks @xcite and multilingual multitask models @xcite . moe offers the benefits of scaling the model to gain better accuracy without paying the huge compute cost of massive dense models. however, large scale moe models bring their own set of unique challenges to get efficient training and inference methods. most of the previous work focused on improving training efficiency and throughput @xcite . in this work, we focus on optimizing moe models inference and latency since it is crucial to harvest the benefits of such models in real-life scenarios. production-scale multilingual machine translation systems: in this work, we explore deploying moe models for large scale multilingual machine translation systems to benefit from large language models, while maintaining reasonable serving cost. multilingual large scale systems are already very attractive due to multiple aspects. first, they benefit modeling since they allow better accuracy, especially through transfer learning across languages. additionally, they improve deployment and serving since we can replace dozens of models with a single model that is able to serve many languages at the same time. nevertheless, we need the infer-ence to be highly optimized to make inference costefficient. despite these benefits, shipping such multilingual models brings a new challenge, because they usually require a much larger model capacity in terms of the number of parameters and the computation. the moe model architecture could be a promising solution given its sub-linear or constant flops increase in terms of the number of model parameters. but, the large memory consumption issue still remains. in this work, we show how to enable deploying a single moe model that can serve many languages replacing dozens of traditional models while improving accuracy and maintaining latency, throughput and cost efficiency. we set the goal for this work to match latency and throughput of a distilled small model deployed on cpu while achieving better serving cost. it is worth noting that while the optimizations presented here are applied to moe encoder-decoder architecture for multilingual machine translation task, they are applicable to other architectures and tasks without any loss of generality. given the recent success of moe models on wide set of nlu and nlg tasks @xcite , we believe the optimization presented in this work will be equally enabling to other tasks as it is for machine translation. ## results and discussion all experiments in this section are run on a single nvidia pcie v100 running inside a docker container running ubuntu 20.04 and cuda 11.6. all code is compiled with nvcc and gcc/g++ 9.3. we run our experiments considering an encoderdecoder moe model with 32 experts with tupe @xcite , similar to the setup in @xcite but with a vocabulary size of 128k. all throughput metrics measure the time to translate 1000 tokenized english sentences (∼ 40k tokens) to german (en-de) or vice-versa (de-en) and record the total number of input tokens translated per second. bleu metrics are reported on the same data set. ## conclusions and future work this paper describes how to make large moe models cost-efficient on a single gpu in a real-world inference environment. the final implementation achieves a speedup of up to 26x over pytorch baseline. our gpu moe implementation allows serving much larger and higher-quality models compared to dense models on cpus without increasing the cost of serving. we consider two main avenues for future work. we are currently working on improving our fused gemm + dequantize kernel to enable the use of fully vectorized 16 byte loads on the weight matrix. in addition, we plan to explore deploying even larger models with distributed inference in the future in a cost-efficient way.