FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models

Abstract

Recent large language models (LLMs) have tended to leverage sparsity to reduce computations, employing the sparsely activated mixture-of-experts (MoE) technique. MoE introduces four modules, including token routing, token communication, expert computation, and expert parallelism, that impact model quality and training efficiency. To enable versatile usage of MoE models, we introduce FSMoE, a flexible training system optimizing task scheduling with three novel techniques: 1) Unified abstraction and online profiling of MoE modules for task scheduling across various MoE implementations. 2) Co-scheduling intra-node and inter-node communications with computations to minimize communication overheads. 3) To support near-optimal task scheduling, we design an adaptive gradient partitioning method for gradient aggregation and a schedule to adaptively pipeline communications and computations. We conduct extensive experiments with configured MoE layers and real-world MoE models on two GPU clusters. Experimental results show that 1) our FSMoE supports four popular types of MoE routing functions and is more efficient than existing implementations (with up to a 1.42$\times$ speedup), and 2) FSMoE outperforms the state-of-the-art MoE training systems (DeepSpeed-MoE and Tutel) by 1.18$\times$-1.22$\times$ on 1458 MoE layers and 1.19$\times$-3.01$\times$ on real-world MoE models based on GPT-2 and Mixtral using a popular routing function.

References in corpus (7)

• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
• BASE Layers: Simplifying Training of Large, Sparse Models
• A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training
• FlexMoE: Scaling Large-scale Sparse Pre-trained Model Training via Dynamic Device Placement
• T3: Transparent Tracking & Triggering for Fine-grained Overlap of Compute & Collectives