[论文解读] MoESys: A Distributed and Efficient Mixture-of-Experts Training and Inference System for Internet Services

Dianhai Yu, Liang Shen+6 more|arXiv (Cornell University)|May 20, 2022

Advanced Neural Network Applications被引用 12

一句话总结

SE-MoE 是一种分布式、高效的混合专家（MoE）训练与推理系统，通过二维预取、分层存储融合以及基于 CPU-GPU 内存环的推理机制，显著提升了负载均衡性、通信效率和内存利用率。与 DeepSpeed 相比，SE-MoE 在标准 MoE 工作负载下实现了 33% 的更高训练吞吐量和 13% 的更高推理吞吐量；在负载不均衡的 MoE 工作负载下，吞吐量提升达 64%，内存占用降低 18%。该系统成功在 48 张 A100 GPU 上用 8 天时间训练完成一个 120 亿参数的 UFO 模型。

ABSTRACT

While modern internet services, such as chatbots, search engines, and online advertising, demand the use of large-scale deep neural networks (DNNs), distributed training and inference over heterogeneous computing systems are desired to facilitate these DNN models. Mixture-of-Experts (MoE) is one the most common strategies to lower the cost of training subject to the overall size of models/data through gating and parallelism in a divide-and-conquer fashion. While DeepSpeed has made efforts in carrying out large-scale MoE training over heterogeneous infrastructures, the efficiency of training and inference could be further improved from several system aspects, including load balancing, communication/computation efficiency, and memory footprint limits. In this work, we present a novel MoESys that boosts efficiency in both large-scale training and inference. Specifically, in the training procedure, the proposed MoESys adopts an Elastic MoE training strategy with 2D prefetch and Fusion communication over Hierarchical storage, so as to enjoy efficient parallelisms. For scalable inference in a single node, especially when the model size is larger than GPU memory, MoESys builds the CPU-GPU memory jointly into a ring of sections to load the model, and executes the computation tasks across the memory sections in a round-robin manner for efficient inference. We carried out extensive experiments to evaluate MoESys, where MoESys successfully trains a Unified Feature Optimization (UFO) model with a Sparsely-Gated Mixture-of-Experts model of 12B parameters in 8 days on 48 A100 GPU cards. The comparison against the state-of-the-art shows that MoESys outperformed DeepSpeed with 33% higher throughput (tokens per second) in training and 13% higher throughput in inference in general. Particularly, under unbalanced MoE Tasks, e.g., UFO, MoESys achieved 64% higher throughput with 18% lower memory footprints.

研究动机与目标

• 解决由于负载不均、通信瓶颈和高内存占用导致的分布式 MoE 训练与推理效率低下问题。
• 实现在异构 GPU-CPU 集群上的可扩展、高吞吐量 MoE 训练与推理。
• 在部署超出 GPU 内存容量的 MoE 模型时，减轻内存压力并提升系统利用率。
• 通过创新的系统级技术优化 MoE 系统中的通信与计算负载。
• 在互联网服务中实现大规模 MoE 模型的生产级部署，且资源开销最小化。

提出的方法

• 提出弹性 MoE 训练机制，结合二维预取，提升设备间数据并行与模型并行的效率。
• 采用分层存储上的融合通信机制，减少节点间通信开销，提升带宽利用率。
• 利用 CPU-GPU 内存环结构，实现对超出 GPU 内存容量的 MoE 模型的高效、轮询式推理。
• 在数据并行中应用嵌入层分区技术，降低 GPU 内存消耗，提升训练吞吐量。
• 结合分层存储与内存管理策略，最小化数据移动及主机到设备（D2H）/设备到主机（H2D）的开销。
• 集成优化的路由与负载均衡策略，缓解 MoE 训练过程中的专家不均衡问题。

实验结果

研究问题

• RQ1 如何在分布式异构计算环境中提升 MoE 训练与推理的效率？
• RQ2 哪些系统级优化能够有效缓解 MoE 训练中的通信与内存瓶颈？
• RQ3 如何高效执行超出 GPU 内存容量的 MoE 模型？
• RQ4 通过系统级设计，能在多大程度上缓解 MoE 路由中的负载不均问题？
• RQ5 统一系统是否能够在 MoE 训练与推理中同时实现高吞吐量与低内存占用？

主要发现

• 在标准 MoE 工作负载下，SE-MoE 的训练吞吐量（每秒处理的 token 数）比 DeepSpeed 高出 33%。
• 在通用 MoE 任务中，SE-MoE 的推理吞吐量比 DeepSpeed 高出 13%。
• 在如 UFO 模型等负载不均衡的 MoE 工作负载下，SE-MoE 的吞吐量比 DeepSpeed 高出 64%，内存占用降低 18%。
• 该系统成功在 48 张 A100 GPU 上用 8 天时间训练完成一个 120 亿参数的统一特征优化（UFO）MoE 模型。
• 在 V100 GPU 上，数据并行中的嵌入层分区技术将 GPU 内存使用量最多降低 26.3%，并将吞吐量最多提升 15.6%。
• CPU-GPU 内存环推理机制通过轮询式任务调度，实现了对超出 GPU 内存容量的 MoE 模型的高效执行。