ES-MoE: overcoming the scalability challenges in mixture-of-experts models

Abstract

Mixture-of-Experts (MoE) models have recently emerged as a powerful technique for enhancing the scalability and performance of neural networks, primarily by leveraging learnable gating networks to allocate input tokens to different expert models. However, training MoE models on GPUs presents unique challenges, including insufficient GPU memory capacity for a large number of experts and computational inefficiency due to token load imbalance. To address these issues, we introduce Expert Server MoE (ES-MoE), a novel method that offloads all expert parameters and their optimizer states to CPUs. This approach not only mitigates the memory constraints of GPU-based training but also enhances training throughput by creating a unified pool of experts that allows for more efficient scheduling. Furthermore, ES-MoE employs pipelined expert optimization to minimize the iteration latency, effectively circumventing the issue of extended CPU optimization time. We validate our approach using GPT-based MoE architectures, demonstrating that ES-MoE scales up to 16 times better, and improves throughput up to 4.55x over the existing frameworks.