[论文解读] HetuMoE: An Efficient Trillion-scale Mixture-of-Expert Distributed Training System

Xiaonan Nie, Pinxue Zhao+3 more|arXiv (Cornell University)|Mar 28, 2022

Domain Adaptation and Few-Shot Learning被引用 11

一句话总结

HetuMoE 是一个基于 Hetu 构建的高性能、分布式混合专家（MoE）训练系统，支持多种门控策略，并在通用 GPU 集群上实现了训练加速。它引入了分层 AllToAll 通信和优化内核，在批量大小为 32 的 Switch 门控下，相较于 DeepSpeed-MoE 最高实现 8.1 倍的加速。

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.

研究动机与目标

• 为解决现有 MoE 训练系统在通用 GPU 集群上的低效与可扩展性受限问题。
• 在一个统一框架中支持多种主流门控策略，如 Switch、GShard、Top-k 等。
• 通过优化低带宽网络中的 AllToAll 通信，减少分布式 MoE 训练中的通信瓶颈。
• 在不依赖 NVLink 或 InfiniBand 等高速互连技术的情况下，实现最先进的训练性能。
• 实现在标准 GPU 集群上高效、可扩展且实用的万亿参数 MoE 模型训练。

提出的方法

• 提出一种分层 AllToAll 通信模式，结合节点内与节点间通信，聚合小规模消息以提升带宽利用率。
• 为关键 MoE 操作（包括 Top-k 门控、布局转换和 AllToAll）实现优化的 GPU 内核，并针对 k=1 和 k=2 的情况引入算法改进。
• 采用两级数据布局转换：首先在节点内按专家分组 token，然后在节点间执行 AllToAll，以减少单次通信的消息大小。
• 利用 Hetu 的分布式深度学习框架，原生支持 MoE 训练，仅需极少代码修改即可实现高性能。
• 支持多种门控策略，包括 Switch、GShard、M6、BASE Layer、Hash Layer、SAM 和 Dense-to-Sparse，实现广泛的模型兼容性。
• 端到端优化 MoE 训练流水线：门控计算、布局转换、GPU 间通信、专家处理以及反向布局转换。

实验结果

研究问题

• RQ1 一个 MoE 训练系统是否能在不依赖 NVLink 或 InfiniBand 等高速互连技术的通用 GPU 集群上实现高性能？
• RQ2 分层 AllToAll 通信在多节点 MoE 训练中减少通信开销方面的有效性如何？
• RQ3 针对 Top-k 门控和布局转换的内核级优化，能在多大程度上提升 MoE 训练效率？
• RQ4 在不同门控策略和批量大小下，HetuMoE 与 DeepSpeed-MoE、FastMoE 和 Tutel 等先进系统相比性能如何？
• RQ5 统一系统能否高效支持广泛的 MoE 门控策略，同时保持高训练吞吐量？

主要发现

• HetuMoE 在各种配置和门控策略下，相比现有 MoE 系统至少实现 15% 的加速。
• 在批量大小为 32 的 Switch 门控下，HetuMoE 的训练速度相比 DeepSpeed-MoE 最高提升 8.1 倍。
• 与原始 AllToAll 相比，分层 AllToAll 优化在 4×8 GPU 设置下实现 1.66 倍加速，在 8×8 GPU 设置下实现 2 倍加速。
• HetuMoE 中优化的 Top-k 内核相比 PyTorch 实现，在 k=1 和 k=2 情况下平均性能提升 25%。
• HetuMoE 的布局转换内核相比现有最先进实现，性能提升超过 26%。
• HetuMoE 有效支持多种 MoE 门控策略，包括 Switch、GShard、M6、BASE Layer、Hash Layer、SAM 和 Dense-to-Sparse，实现广泛的模型兼容性。