FastMoE: A Fast Mixture-of-Expert Training System

By Jiaao He, Jiezhong Qiu, Aohan Zeng, Zhilin Yang, Jidong Zhai, Jie Tang. 2021

148 citations recorded by the source corpus.

Mixture-of-Expert (MoE) presents a strong potential in enlarging the size of language model to trillions of parameters. However, training trillion-scale MoE requires algorithm and system co-design for a well-tuned high performance distributed training system. Unfortunately, the only existing platform that meets the requirements strongly depends on Google's hardware (TPU) and software (Mesh Tensorflow) stack, and is not open and available to the public, especially GPU and PyTorch communities. In this paper, we present FastMoE, a distributed MoE training system based on PyTorch with common accelerators. The system provides a hierarchical interface for both flexible model design and easy adaption to different applications, such as Transformer-XL and Megatron-LM. Different from direct implementation of MoE models using PyTorch, the training speed is highly optimized in FastMoE by sophisticated high-performance acceleration skills. The system supports placing different experts on multiple GPUs across multiple nodes, enabling enlarging the number of experts linearly against the number of GPUs. The source of FastMoE is available at https://github.com/laekov/fastmoe under Apache-2 license.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

Citations in the corpus

Pre-Trained Models: Past, Present and Future (2021, methodology): “Some efficient open-source toolkits (He et al., 2021) are also developed to train large-scale MoE-based models.”
A Survey of Large Language Models (2023, unknown): “• FastMoE [128] is a specialized training library for MoE (i.”
CoST: Contrastive Learning of Disentangled Seasonal-Trend Representations for Time Series Forecasting (2022, methodology): “Furthermore, the extra computation time of CoST as compared to TS2Vec is due to the sequential computation of each expert in the mixture of AR expert component, it can be further accelerated by parallel methods (He et al., 2021).”
Residual Mixture of Experts (2022, methodology): “For MoE implementation, we use FastMoE [23] as basic code.”
A Comprehensive Overview of Large Language Models (2023, unknown): “Some commonly used libraries for LLMs training are: 1) Transformer [64], 2) DeepSpeed [65], 3) Megatraon-LM [62], 4) JAX [66], 5) Colossal-AI [67], 6) BMTrain [63], 7) FastMoE [68], and frameworks are 1) MindSpore [69], 2) PyTorch [70], 3) Tensorflow [71], 4) MXNet [72].”
DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale (2022, methodology): “FastMoE [32] is a research software developed to show how MoE models can be trained under data and expert (model) parallelism.”
Optimizing Mixture of Experts using Dynamic Recompilations (2022, methodology): “FastMoE (He et al., 2021) is a distributed MoE training system based on PyTorch with support for multiple GPUs.”
A Survey on Mixture of Experts in Large Language Models (2024, unknown): “Therefore, recent studies like DeepSpeed-MoE [64], FastMoE [104], HetuMoE [109] and Tutel [105] have focused on the development of tailored GPU kernels to enhance the operation efﬁciency.”
Speechmoe2: Mixture-of-Experts Model with Improved Routing (2021, methodology): “Besides, we use FastMoE [14] in our sys-tem to support scaling up the MoE model size by training across multiple GPUs on multiple nodes.”
CPM-2: Large-scale Cost-effective Pre-trained Language Models (2021, background): “FastMoE (He et al., 2021) is a novel high-performance MoE implementation on top of PyTorch.”
FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models (2022, methodology): “We implement FasterMoE based on FastMoE [7].”
Mixture-of-Experts in Remote Sensing: A Survey (2026, unknown): “Works such as FastMoE [45], DeepSpeed-MoE [118], Tutel [55], HetuMoE [108], FasterMoE [46] and MegaBlocks [32] focus on efficient distributed training and inference, including all-to-all communication scheduling, CUDA kernel optimizations, expert placement, and parallelism strategies.”