[论文解读] FastMoE: A Fast Mixture-of-Expert Training System

Jiaao He, Jiezhong Qiu+4 more|arXiv (Cornell University)|Mar 24, 2021

Topic Modeling参考文献 16被引用 39

一句话总结

FastMoE 是一个基于 PyTorch 的分布式 Mixture-of-Experts (MoE) 训练系统，能够在多张 GPU 和节点上放置多个专家，具备优化的 CUDA 内核和灵活的 Megatron-LM 集成，在商用硬件上实现可扩展的 MoE 训练。

ABSTRACT

Mixture-of-Expert (MoE) presents a strong potential in enlarging the size of language model to trillions of parameters. However, training trillion-scale MoE requires algorithm and system co-design for a well-tuned high performance distributed training system. Unfortunately, the only existing platform that meets the requirements strongly depends on Google's hardware (TPU) and software (Mesh Tensorflow) stack, and is not open and available to the public, especially GPU and PyTorch communities. In this paper, we present FastMoE, a distributed MoE training system based on PyTorch with common accelerators. The system provides a hierarchical interface for both flexible model design and easy adaption to different applications, such as Transformer-XL and Megatron-LM. Different from direct implementation of MoE models using PyTorch, the training speed is highly optimized in FastMoE by sophisticated high-performance acceleration skills. The system supports placing different experts on multiple GPUs across multiple nodes, enabling enlarging the number of experts linearly against the number of GPUs. The source of FastMoE is available at https://github.com/laekov/fastmoe under Apache-2 license.

研究动机与目标

• 在公开可用的硬件（GPU）上推动可扩展的 MoE 训练，超越 TPU/Mesh TensorFlow 堆栈。
• 提供一个易于使用、灵活且高效的 MoE 框架，集成于 PyTorch 与 Megatron-LM。
• 实现将大量专家分布在多张 GPU 和节点上，以扩展模型容量。

提出的方法

• 引入一个灵活的 FMoE 接口，将任意神经网络复制为一个专家，并支持自定义 expert_fn 以实现高级 MoE 行为。
• 提供为基于 Transformer 的 MoE 进行了性能优化的 FFN (FMoETransformerMLP)，并具备用于并行专家执行的专用 FMoELinear 模块。
• 通过两行变换（fmoefy）实施与 Megatron-LM 的插件式集成，将 FFN 替换为 MoE 网络。
• 支持将专家分布在多个工作进程和节点上（模型并行 MoE），具备隐藏的全局数据交换操作和异构同步模块。
• 使用自定义的 scatter/gather 内存布局按专家分批输入，以最大化 GPU 利用率并减少数据传输。

实验结果

研究问题

• RQ1 FastMoE 能否在多 GPU、多节点集群上使用普通过货 GPU 实现可扩展的 MoE 训练？
• RQ2 具有专用 FFN 和并行专家执行的优化 MoE 实现，是否优于在单个 GPU 上的朴素 PyTorch MoE 实现？
• RQ3 与非 MoE 基线相比，在与 Megatron-LM 集成的端到端训练中，FastMoE 的表现如何？
• RQ4 在增加 GPU 或专家数量时，扩展性极限和瓶颈是什么（例如通信）？

主要发现

• FastMoE 在单个 GPU 上针对类似任务实现比基线 PyTorch MoE 的前向/后向更快。
• 在由 8 个 NVIDIA Tesla V100 GPU 构成的集群上，FastMoE 的吞吐量从约 10 TFLOPs 增长到 25 TFLOPs，因通信开销而呈现亚线性可扩展性。
• 在 8 GPU 上，12 层 GPT 模型的端到端训练，每层 96 个专家（每 GPU 12 个专家），在 MoE 下展示了模型容量的提升，尽管单轮迭代速度较慢，在同等训练时间内获得更低的损失。
• 使用 FastMoE 与 Megatron-LM，一种每层 96 个专家的 MoE GPT 模型训练 70 小时，显示出相对于具有类似计算量的非 MoE 基线的模型性能提升。
• FastMoE 通过替换 FFN 实现轻松转换为 MoE，并通过两行代码变更（fmoefy）展示与 Megatron-LM 的实际集成。
• 该系统强调开源可用性，并计划为未来改进加入负载均衡机制。