{"id":798607,"date":"2021-11-22T09:40:53","date_gmt":"2021-11-22T17:40:53","guid":{"rendered":"https:\/\/www.microsoft.com\/en-us\/research\/?p=798607"},"modified":"2021-11-22T09:40:54","modified_gmt":"2021-11-22T17:40:54","slug":"tutel-an-efficient-mixture-of-experts-implementation-for-large-dnn-model-training","status":"publish","type":"post","link":"https:\/\/www.microsoft.com\/en-us\/research\/blog\/tutel-an-efficient-mixture-of-experts-implementation-for-large-dnn-model-training\/","title":{"rendered":"Tutel: An efficient mixture-of-experts implementation for large DNN model training"},"content":{"rendered":"\n
\"A<\/figure>\n\n\n\n

Mixture of experts (MoE) is a deep learning model architecture in which computational cost is sublinear to the number of parameters, making scaling easier. Nowadays, MoE is the only approach demonstrated to scale deep learning models to trillion-plus parameters, paving the way for models capable of learning even more information and powering computer vision, speech recognition, natural language processing, and machine translation systems, among others, that can help people and organizations in new ways.<\/p>\n\n\n\n

Today, we’re proud to announce Tutel, a high-performance MoE library (opens in new tab)<\/span><\/a> to facilitate the development of large-scale DNN models; Tutel is highly optimized for the new Azure NDm A100 v4 series, now generally available (opens in new tab)<\/span><\/a>. With Tutel\u2019s diverse and flexible MoE algorithmic support, developers across AI domains can execute MoE more easily and efficiently. For a single MoE layer, Tutel achieves an 8.49x speedup on an NDm A100 v4 node with 8 GPUs and a 2.75x speedup on 64 NDm A100 v4 nodes with 512 A100 GPUs (all experiments in this blog are tested on Azure NDm A100 v4 nodes with 8 x 80 GB NVIDIA A100 and an 8 x 200 gigabits per second InfiniBand network), respectively, compared with state-of-the-art MoE implementations such as that in Meta\u2019s Facebook AI Research Sequence-to-Sequence Toolkit (fairseq) (opens in new tab)<\/span><\/a> in PyTorch. For end-to-end performance, Tutel\u2014benefiting from an optimization for all-to-all communication\u2014achieves a more than 40 percent speedup with 64 NDm A100 v4 nodes for Meta\u2019s (Facebook is now Meta) 1.1 trillion\u2013parameter MoE language model. Tutel provides great compatibility with rich features to ensure the great performance when working on the Azure NDm A100 v4 cluster. Tutel is open source and has been integrated into fairseq.<\/p>\n\n\n\n

Tutel MoE optimizations<\/h2>\n\n\n\n

Complementary to other high-level MoE solutions like fairseq and FastMoE, Tutel mainly focuses on the optimizations of MoE-specific computation and the all-to-all communication, as well as other diverse and flexible algorithmic MoE supports. Tutel has a concise interface, making it easy to integrate into other MoE solutions. Alternatively, developers can use the Tutel interface to incorporate standalone MoE layers into their own DNN models from scratch and benefit from the highly optimized state-of-the-art MoE features directly.<\/p>\n\n\n\n

MoE-specific optimization for computation<\/h3>\n\n\n\n

Because of the lack of efficient implementations, MoE-based DNN models rely on a naive combination of multiple off-the-shelf DNN operators provided by deep learning frameworks such as PyTorch and TensorFlow to compose the MoE computation. Such a practice incurs significant performance overheads thanks to redundant computation. Tutel designs and implements multiple highly optimized GPU kernels to provide operators for MoE-specific calculation. For example, Tutel reduces the time complexity of dispatching \u201cgating output\u201d from O(N^3) to O(N^2), which significantly improves the data dispatching efficiency. Tutel also implements a fast cumsum-minus-one operator, achieving a 24x speedup compared with the fairseq implementation. Tutel also leverages NVRTC (opens in new tab)<\/span><\/a>, a runtime compilation library for CUDA C++, to further optimize the customized MoE kernel just-in-time. Figure 1 shows the comparison results of Tutel with fairseq on the Azure NDm A100 v4 platform, where\u2014as mentioned above\u2014a single MoE layer with Tutel achieves an 8.49x speedup on 8 A100 GPUs and a 2.75x speedup on 512 A100 GPUs.<\/p>\n\n\n\n

Underlying all-to-all communication optimization on Azure NDm A100 v4 clusters<\/h3>\n\n\n\n

Tutel also optimizes the all-to-all collective communication for large-scale MoE training on Azure NDm A100 v4 clusters, including CPU-GPU binding and adaptive routing (AR) tuning. A proper CPU-GPU binding on a multi-non-uniform memory access (NUMA) system, especially on the NDm A100 v4 nodes, is very critical for all-to-all performance. Unfortunately, existing machine learning frameworks have not provided an efficient all-to-all communication library, resulting in performance regression for large-scale distributed training. Tutel optimizes the binding automatically and provides an elegant interface for user fine-tuning. Furthermore, Tutel leverages multipath technology, namely AR, on NDm A100 v4 clusters. For the all-to-all communication in MoE, the total data traffic size of the communication for each GPU doesn\u2019t change, but the data size between each GPU pair becomes smaller with the increasing number of GPUs. Smaller data size incurs a larger overhead in the all-to-all communication, leading to poorer MoE training performance. By taking advantage of AR technology available on NDm A100 v4 clusters, Tutel improves communication efficiency for groups of small messages and provides high-performance all-to-all communication on NDm A100 v4 systems. Benefiting from CPU-GPU binding and AR tuning, Tutel achieves a 2.56x to 5.93x all-to-all speedup with 512 A100 GPUs for message sizes hundreds of MiB large, which are typically used in MoE training, as illustrated in Figure 2.<\/p>\n\n\n\n

\"On
Figure 1 (left): Compared to fairseq, for a single MoE layer, Tutel achieves an 8.49x speedup on an NDm A100 v4 node with 8 GPUs and a 2.75x speedup on 64 NDm A100 v4 nodes with 512 A100 GPUs. The detailed setting is as follows: batch_size = 32, sequence_length = 1,024, Top_K = 2, model_dim = 2,048, and hidden_size = 2,048. Figure 2 (right): The all-to-all bandwidth for different message sizes with 64 NDm A100 v4 nodes (512 A100 GPUs) before and after applying Tutel. Tutel achieves a 2.56x to 5.93x all-to-all speedup with 512 A100 GPUs for message sizes hundreds of MiB large.<\/figcaption><\/figure><\/div>\n\n\n\n

Diverse and flexible MoE algorithms support<\/h3>\n\n\n\n

Tutel provides diverse and flexible support for state-of-the-art MoE algorithms, including support for:<\/p>\n\n\n\n