HetuMoE: An Efficient Trillion-scale Mixture-of-Expert Distributed Training System

TL;DR

This work tackles the efficiency bottlenecks of training enormous sparse Mixture-of-Experts models by introducing HetuMoE, a high-performance distributed MoE system built on the Hetu framework. HetuMoE supports a broad set of gating strategies and introduces Hierarchical AllToAll to optimize communication on commodity GPU clusters, delivering substantial end-to-end speedups. Key contributions include multiple gating implementations (Topk, kTop1, Hierarchical Topk, BASE Layer, Hash Layer, Dense-To-Sparse) and optimization of gate computation, data layout, and AllToAll, achieving at least 15% speedup over baselines and up to 8.1x over DeepSpeed-MoE under certain settings. The approach enables more scalable, efficient training of trillion-parameter MoE models in practical cluster environments, broadening applicability of MoE architectures in industry and academia.

Abstract

As giant dense models advance quality but require large amounts of GPU budgets for training, the sparsely gated Mixture-of-Experts (MoE), a kind of conditional computation architecture, is proposed to scale models while keeping their computation constant. Specifically, the input tokens are routed by the gate network and only activates part of the expert network. Existing MoE training systems only support part of mainstream MoE models (e.g. Top k) training under expensive high-bandwidth GPU clusters. In this paper, we present HetuMoE, a high-performance large-scale sparse MoE training system built on Hetu. HetuMoE provides multiple gating strategies and efficient GPU kernel implementations. To further improve the training efficiency on commodity GPU clusters (e.g, with only 1 NiC), we introduce the hierarchical AllToAll communication that combines hierarchical networks and aggregating messages. Compared with existing state-of-the-art MoE systems, HetuMoE obtains at least 15% speedup. Specifically, HetuMoE outperforms DeepSpeed-MoE up to 8.1x under the switch gate with a batch size of 32. Our code is available at: https://github.com/PKU-DAIR/Hetu.

Figures (8)

• Figure 1: Time consumption of MoE layer of DeepSpeed-MoE under 8 A100 GPUs in a single node
• Figure 2: HetuMoE compared with other MoE systems
• Figure 3: Topk kernel performance comparison with PyTorch
• Figure 4: Data layout transformation illustration and performance comparison
• Figure 5: Illustration of NCCL AllToAll.