[论文解读] Tutel: Adaptive Mixture-of-Experts at Scale

Chang Ho Hwang, Wei Cui+13 more|arXiv (Cornell University)|Jun 7, 2022

Advanced Neural Network Applications被引用 31

一句话总结

Tutel 引入了一个零开销切换的自适应并行和流水线框架，针对稀疏 MoE 模型，在保持对现有MoE形式的算法等价性的同时，在大规模场景下实现显著的加速。

ABSTRACT

Sparsely-gated mixture-of-experts (MoE) has been widely adopted to scale deep learning models to trillion-plus parameters with fixed computational cost. The algorithmic performance of MoE relies on its token routing mechanism that forwards each input token to the right sub-models or experts. While token routing dynamically determines the amount of expert workload at runtime, existing systems suffer inefficient computation due to their static execution, namely static parallelism and pipelining, which does not adapt to the dynamic workload. We present Flex, a highly scalable stack design and implementation for MoE with dynamically adaptive parallelism and pipelining. Flex designs an identical layout for distributing MoE model parameters and input data, which can be leveraged by all possible parallelism or pipelining methods without any mathematical inequivalence or tensor migration overhead. This enables adaptive parallelism/pipelining optimization at zero cost during runtime. Based on this key design, Flex also implements various MoE acceleration techniques. Aggregating all techniques, Flex finally delivers huge speedup at any scale -- 4.96x and 5.75x speedup of a single MoE layer over 16 and 2,048 A100 GPUs, respectively, over the previous state-of-the-art. Our evaluation shows that Flex efficiently and effectively runs a real-world MoE-based model named SwinV2-MoE, built upon Swin Transformer V2, a state-of-the-art computer vision architecture. On efficiency, Flex accelerates SwinV2-MoE, achieving up to 1.55x and 2.11x speedup in training and inference over Fairseq, respectively. On effectiveness, the SwinV2-MoE model achieves superior accuracy in both pre-training and down-stream computer vision tasks such as COCO object detection than the counterpart dense model, indicating the readiness of Flex for end-to-end real-world model training and inference.

研究动机与目标

• 解释由于令牌路由和专家工作负载不均导致的 MoE 工作负载在海量级别上具有动态性和挑战性。
• 提出一个零开销自适应并行切换机制，以统一并切换 DP、EP+DP+MP 及相关配置。
• 开发自适应流水线和全对全(All-to-All)策略（包括 2DH），以共同优化计算和通信。
• 提供一个实用的、与框架无关的 MoE 系统（Tutel），具备动态容量因子和动态 top-k 路由支持。
• 演示在大规模 GPU 集群上使用 SwinV2-MoE 的可扩展性与真实世界的有效性。

提出的方法

• 分析 MoE 层的动态工作负载特征及其对并行性选择的影响。
• 推导出在各种配置下实现最佳性能所需的最小而完整的并行化策略集合（DP 与 EP+DP+MP）。
• 引入零开销可切换的并行性，通过重新构建数据布局使切换的开销为 O(1)。
• 开发自适应流水线，基于工作负载共同选择流水线度数与全对全算法（Linear 与 2DH）。
• 构建跨容量区间的最优并行性/流水线设置的字典（哈希映射），以实现快速运行时决策。
• 实现带稀疏计算的快速编码/解码，以及灵活的 All-to-All，以降低内存和通信开销。

实验结果

研究问题

• RQ1 MoE 工作负载如何随训练动态和容量因子而变化，以及这对静态与动态并行性的含义？
• RQ2 在保持相同数据布局和算法正确性的前提下，是否能够实现并行策略之间的零成本切换？
• RQ3 在动态 MoE 工作负载下，哪种数据并行、模型并行和专家并行的组合可以实现最佳吞吐量，以及如何在运行时进行适应？
• RQ4 如何联合优化自适应流水线和 All-to-All 算法（Linear 与 2DH），以在大规模下最大化 MoE 模型的吞吐量？
• RQ5 在实际的 MoE 模型，如 SwinV2-MoE 的训练和推理中，Tutel 能带来哪些实际的性能提升？

主要发现

• Tutel 相较于先前的最先进方法，在单个 MoE 层上对 16 GPUs 实现最高 4.96x 的加速，在 2048 GPUs 上实现 5.75x。
• 端到端的 SwinV2-MoE 训练在训练阶段最高达到 1.55x 的加速，在推理阶段达到 2.11x。
• 自适应并行切换（通过 DP 和 EP+DP+MP）提供的吞吐提升取决于容量因子 f，在多种设置中优于静态策略。
• 自适应流水线在 243 种 MoE 模型设置下的平均提升范围为 9% 至超过 100%，最坏情况下可提升高达 599%。
• 灵活的 All-to-All 和快速编码/解码技术降低内存使用并提升吞吐量，使在大型 GPU 集群上实现可扩展的稀疏 MoE 成为可能。
• 在下游计算机视觉任务（如 COCO）上，使用 Tutel 的 SwinV2-MoE 比密集基线具备更高的准确度，表明该方法的实际有效性。