Source-linked AI summary
FastMoE: A Fast Mixture-of-Expert Training System
Jiaao He, Jiezhong Qiu, Aohan Zeng, Zhilin Yang, Jidong Zhai, Jie Tang
TL;DR
Large-scale MoE can expand language models to trillions of parameters, but scalable training previously depended on Google’s private TPU and Mesh TensorFlow stack. FastMoE provides an open PyTorch system that optimizes MoE execution across GPUs and nodes, with experiments showing faster single-GPU execution, reasonable multi-node scalability, and promising end-to-end model-training performance. Its current limitation is the lack of expert load-balancing support.
Problem
Scalable MoE training requires coordinated algorithm and system design, while the available scalable platform depends on Google’s private TPU and Mesh TensorFlow stack.
Method
FastMoE is an open PyTorch MoE system using optimized kernels, distributed expert placement, hidden communication, and GPU multi-node execution.
Results
FastMoE is faster than a pure-PyTorch single-GPU baseline, scales reasonably across nodes, and demonstrates real-model performance advantages with distributed MoE training.
Takeaways & Limitations
FastMoE supports enlarging MoE model capacity with additional GPUs while providing practical distributed training for open GPU and PyTorch environments.
Takeaways & Limitations
FastMoE lacks functionality for load balancing among experts, including load-balance monitoring and load-balance loss support.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Scaling model size has driven major language-model advances, with MoE enabling trillion-parameter models through sparse expert activation. FastMoE addresses the resulting distributed-training gap with an open, efficient, scalable PyTorch system for GPUs and multi-node clusters.
- Motivation: Model scaling grew from BERT’s 340 million parameters to Switch Transformer’s 1.6 trillion, with MoE as the main contributor to this expansion.The cited progression includes T5 at 11 billion, GPT-3 at 175 billion, and GShard at 600 billion parameters.
- Motivation: MoE activates only a tiny minority of experts per input, enlarging model size by orders of magnitude without significantly increasing FLOPs.Scaling to thousands of experts introduces imbalanced all-to-all communication challenges that traditional PyTorch and TensorFlow libraries do not directly support.
- Problem: Before FastMoE, scalable MoE training depended on Google’s private TPU and Mesh TensorFlow stack, leaving a need for publicly available GPU and PyTorch support.The passage contrasts this system with naive single-GPU PyTorch implementations.
- FastMoE: FastMoE provides easy-to-use, flexible, efficient, and scalable MoE training with customizable gates and experts, optimized Transformer FFNs, and multi-GPU multi-node execution.It also aims to support Megatron-LM and hide distributed communication from model developers.
- FastMoE: FastMoE distributes experts across GPUs while retaining batch- or tensor-dimension parallelism elsewhere, allowing model size to scale with the number of GPUs.NCCL supports execution across GPUs and nodes, while communication details are hidden from model developers.
- Results: FastMoE is faster than a pure-PyTorch single-GPU baseline, scales reasonably across Infiniband-connected nodes, and trains a GPT model with 96 experts per layer.The experiment reports promising end-to-end training speed and benefits from the enlarged MoE model size compared with a non-MoE model using the same computation.
2 MIXTURE-OF-EXPERTS (MOE)
An MoE layer uses a gate to select a small set of experts whose weighted outputs form the layer output. Distributed MoE systems enlarge expert populations, but existing implementations impose hardware, software, or scalability limitations.
- MoE Architecture: An MoE layer contains multiple experts, each an arbitrary neural network receiving the same input and producing outputs in the same vector space.A gate scores experts, and a selection policy chooses which experts process each input.
- MoE Architecture: Top-k gating selects the k experts with highest scores and combines their outputs using the scores as weights.Because the scores weight the outputs, gradients can propagate through the gate network.
- MoE Architecture: The formal forward pass computes gate scores, selects expert indices, applies the selected experts to x, accumulates score-weighted outputs, and returns y.The algorithm requires a pool of n experts, a gate G, and the number k of selected experts.
- Existing Systems: GShard trains distributed MoE models on up to 2048 TPUs, placing one expert per layer on each TPU and producing layers with 2048× more parameters than non-MoE layers.Switch Transformer later enlarged the model to 1.6 trillion parameters, but the system was not publicly available.
- Existing Systems: Tensor2tensor provides an MoE Transformer implementation through Mesh TensorFlow, but its GPU support is weak and its Transformer FFN code is cumbersome to modify.The implementation requires more than 100 lines of TensorFlow with complicated einsum operators.
- Existing Systems: Existing PyTorch MoE implementations offer straightforward coding but lack multi-dimensional parallel-training tools and do not support multi-GPU training.This limits their suitability for training larger MoE models.
3 FastMoE : SYSTEM DESIGN
FastMoE provides a flexible interface for defining and integrating MoE models while hiding distributed communication and supporting experts distributed across workers and nodes. Its design also addresses heterogeneous synchronization and configurable expert placement.
- Flexible model design: FastMoE supports arbitrary expert networks through an interface that decouples expert implementation from the MoE architecture.Developers provide a neural network module constructor, while FastMoE replicates the module into expert instances.
- Flexible model design: Multiple experts can reside on one worker, allowing expert counts independent of the number of data-parallel workers.This provides a more flexible configuration space than assigning one expert per worker.
- Application integration: FastMoE offers optimized Transformer FFNs and plugin-style integration with Megatron-LM through a two-line model transformation.The integration replaces Transformer FFNs with MoE networks while preserving interface compatibility.
- Distributed model capacity: FastMoE distributes experts across multiple workers and nodes while hiding cross-worker data exchange from model developers.Developers write code for a single expert, and FastMoE gathers input data across workers for it.
- Distributed model capacity: Global operations count expert assignments, exchange input-size metadata, calculate buffer offsets, and then exchange the data directly.The incoming and outgoing sample statistics can be reused throughout a training iteration.
- Heterogeneous synchronization: FastMoE uses parameter communication-group tags to determine whether gradients synchronize globally, within data-parallel groups, or not at all.The tags accommodate parameters replicated across workers, divided across model-parallel groups, or belonging to unique experts.
4 OPTIMIZATIONS TO ACHIEVE HIGH-PERFORMANCE
FastMoE improves MoE execution by organizing samples for efficient matrix multiplication and by using specialized memory movement and concurrent expert execution. These optimizations address low utilization from sample-wise computation and small or imbalanced expert batches.
- Performance challenge: Simple PyTorch MoE implementations can achieve less than 5% the peak performance of GPUs.Sample-by-sample computation splits fully connected layers into low-throughput matrix-vector operations.
- Performance challenge: High MoE performance requires batching samples so matrix multiplications operate on sufficiently large dimensions.The paper connects this requirement to hardware tiling techniques that reach high throughput at larger batch sizes.
- Data reordering: FastMoE batches inputs assigned to the same expert and uses a dedicated CUDA kernel for scatter and gather memory movement.Scatter creates contiguous expert inputs, while gather restores the original sample order for later network components.
- Expert execution: Load imbalance can leave some experts with very few samples, while placing multiple experts on one worker lowers their average local batch sizes.These effects reduce the opportunity for efficient individual expert execution.
- Expert execution: FastMoE uses a customized stream manager to execute multiple experts simultaneously and recover potential throughput from small local batches.The concurrent execution strategy targets the reduced batch sizes caused by multiple experts sharing a worker.
5 EVALUATION
FastMoE is evaluated for single-GPU speed, multi-GPU and multi-node scalability, and end-to-end GPT training. It outperforms a pure-PyTorch baseline, scales sub-linearly across nodes, and achieves lower loss at matched training iterations or time despite slower raw training speed.
- FastMoE is compared with a pure-PyTorch MoE implementation on a single GPU and evaluated for distributed scalability and end-to-end Transformer training.
- 5.2 TRAINING SPEED ON A SINGLE GPU: FastMoE outperforms the baseline in overall iteration time because its latency remains stable as the number of experts grows.The baseline’s forward computation takes increasingly more time, while FastMoE benefits from its customized stream manager.
- 5.3 CROSS-GPU AND CROSS-NODE SCALABILITY: 10 TFLOPs to 25 TFLOPs: FastMoE throughput increases as GPUs grow from 2 to 8, but scaling remains sub-linear across nodes.At 2 GPUs, performance is half that of a single GPU, indicating communication-bounded execution.
- 5.3 CROSS-GPU AND CROSS-NODE SCALABILITY: FastMoE supports training large MoE models across multiple GPUs and nodes, though throughput remains open to further optimization.
- 5.4 END-TO-END PERFORMANCE GAIN USING FastMoE: The baseline trains about 3× faster, but the FastMoE MoE model reaches lower loss with the same training iterations and within the same training time.The comparison uses a 12-layer GPT model with 96 experts per layer on 8 GPUs.
6 SUMMARY AND FUTURE WORK
FastMoE is an open-source PyTorch system for efficient MoE training on GPUs. It supports multi-node execution with reasonable scalability and shows a real-model performance advantage, while load balancing and other usability features remain unfinished.
- FastMoE is an open-source PyTorch system with multi-level interfaces for training Mixture-of-Experts models.
- FastMoE exploits GPU performance and runs across GPUs on multiple nodes with reasonable scalability, enabling larger model sizes.
- End-to-end training experiments observe a real-model performance advantage using FastMoE.
- Future work: FastMoE lacks expert load-balancing functionality and is still adding utilities such as model loading and saving.