[Paper Review] FastMoE: A Fast Mixture-of-Expert Training System

Jiaao He, Jiezhong Qiu+4 more|arXiv (Cornell University)|Mar 24, 2021

Topic Modeling16 references39 citations

TL;DR

FastMoE is a PyTorch-based distributed Mixture-of-Experts (MoE) training system that enables placing multiple experts across GPUs and nodes, with optimized CUDA kernels and a flexible Megatron-LM integration, achieving scalable MoE training on commodity hardware.

ABSTRACT

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.

Motivation & Objective

• Motivate scalable MoE training on publicly available hardware (GPUs) beyond TPU/Mesh TensorFlow stacks.
• Provide an easy-to-use, flexible, and efficient MoE framework integrated with PyTorch and Megatron-LM.
• Enable distributing a large pool of experts across multiple GPUs and nodes to scale model capacity.

Proposed method

• Introduce a flexible FMoE interface that replicates any neural network as an expert and supports custom expert_fn for advanced MoE behavior.
• Provide a high-performance FFN (FMoETransformerMLP) optimized for Transformer-based MoE with a dedicated FMoELinear module for parallel expert execution.
• Implement plugin-style integration with Megatron-LM via a two-line transformation (fmoefy) to swap FFNs with MoE networks.
• Support distributing experts across multiple workers and nodes (model parallel MoE) with hidden global data exchange operations and a heterogeneous synchronization module.
• Use a custom scatter/gather memory layout to batch inputs by expert to maximize GPU utilization and reduce data movement.

Experimental results

Research questions

• RQ1 Can FastMoE enable scalable MoE training on multi-GPU and multi-node clusters using commodity GPUs?
• RQ2 Does an optimized MoE implementation with a dedicated FFN and parallel expert execution outperform naive PyTorch MoE implementations on single GPUs?
• RQ3 How does FastMoE perform in end-to-end training when integrated with Megatron-LM compared to non-MoE baselines?
• RQ4 What are the scalability limits and bottlenecks (e.g., communication) when increasing the number of GPUs or experts?

Key findings

• FastMoE achieves faster forward/backward times than a baseline PyTorch MoE implementation on a single GPU for similar tasks.
• Across 8 NVIDIA Tesla V100 GPUs on a cluster, FastMoE scales the throughput from about 10 TFLOPs to 25 TFLOPs, showing sub-linear scalability due to communication overhead.
• End-to-end training of a 12-layer GPT model with 96 experts per layer (12 experts per GPU) across 8 GPUs demonstrates MoE gains in model capacity despite slower per-iteration speed, achieving lower loss given equal training time.
• Using FastMoE with Megatron-LM, an MoE GPT model with 96 experts per layer trained for 70 hours shows MoE gains in model performance over a non-MoE baseline with similar compute.
• FastMoE enables easy conversion to MoE by replacing FFNs, and demonstrates practical integration with Megatron-LM via a two-line code change (fmoefy).
• The system highlights open-sourcing availability and plans to add load-balancing mechanisms for future improvements.