FlexMoE: Scaling Large-scale Sparse Pre-trained Model Training via Dynamic Device Placement

Abstract

With the increasing data volume, there is a trend of using large-scale pre-trained models to store the knowledge into an enormous number of model parameters. The training of these models is composed of lots of dense algebras, requiring a huge amount of hardware resources. Recently, sparsely-gated Mixture-of-Experts (MoEs) are becoming more popular and have demonstrated impressive pretraining scalability in various downstream tasks. However, such a sparse conditional computation may not be effective as expected in practical systems due to the routing imbalance and fluctuation problems. Generally, MoEs are becoming a new data analytics paradigm in the data life cycle and suffering from unique challenges at scales, complexities, and granularities never before possible. In this paper, we propose a novel DNN training framework, FlexMoE, which systematically and transparently address the inefficiency caused by dynamic dataflow. We first present an empirical analysis on the problems and opportunities of training MoE models, which motivates us to overcome the routing imbalance and fluctuation problems by a dynamic expert management and device placement mechanism. Then we introduce a novel scheduling module over the existing DNN runtime to monitor the data flow, make the scheduling plans, and dynamically adjust the model-to-hardware mapping guided by the real-time data traffic. A simple but efficient heuristic algorithm is exploited to dynamically optimize the device placement during training. We have conducted experiments on both NLP models (e.g., BERT and GPT) and vision models (e.g., Swin). And results show FlexMoE can achieve superior performance compared with existing systems on real-world workloads -- FlexMoE outperforms DeepSpeed by 1.70x on average and up to 2.10x, and outperforms FasterMoE by 1.30x on average and up to 1.45x.

arXiv commentAccepted by SIGMOD 2023

References in corpus (18)

• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
• Scaling Laws for Neural Language Models
• SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• GLaM: Efficient Scaling of Language Models with Mixture-of-Experts
• PyTorch Distributed: Experiences on Accelerating Data Parallel Training
• BASE Layers: Simplifying Training of Large, Sparse Models
• Galvatron: Efficient Transformer Training over Multiple GPUs Using Automatic Parallelism
• HET: Scaling out Huge Embedding Model Training via Cache-enabled Distributed Framework
• DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale
• ST-MoE: Designing Stable and Transferable Sparse Expert Models
• Hash Layers For Large Sparse Models
• An Experimental Evaluation of Large Scale GBDT Systems
• Towards Communication-efficient Vertical Federated Learning Training via Cache-enabled Local Updates
• Tutel: Adaptive Mixture-of-Experts at Scale
• HetuMoE: An Efficient Trillion-scale Mixture-of-Expert Distributed Training System

Cited by in corpus (7)

• A Survey on Mixture of Experts in Large Language Models
• Towards Efficient Generative Large Language Model Serving: A Survey from Algorithms to Systems
• Towards Efficient Control Flow Handling in Spatial Architecture via Architecting the Control Flow Plane
• FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models
• Mixture of Experts for Decentralized Generative AI and Reinforcement Learning in Wireless Networks: A Comprehensive Survey
• Complex QA and language models hybrid architectures, Survey
• Input Domain Aware MoE: Decoupling Routing Decisions from Task Optimization in Mixture of Experts