微软亚洲研究院发布高性能MoE库Tutel，为大规模DNN模型开发提速！
编者按：作为目前唯一被证明能够将深度学习模型扩展到万亿以上参数的方法，MoE 能让模型学习更多信息，并为计算机视觉、语音识别、自然语言处理和机器翻译系统等提供支持。近期，微软亚洲研究院发布了一个高性能 MoE 库——Tutel，并在 8x 80GB NVIDIA A100 和 8x 200Gbps 的 InfiniBand 网络的 Azure NDm A100 v4 节点上进行了实验。让我们来看一看，这个用于促进大规模 DNN 模型开发的高性能 MoE 库有哪些优势？其表现又如何？
混合专家（Mixture-of-Experts，简称 MoE）是一种深度学习模型架构，其计算成本与参数的数量呈次线性关系，因此更容易扩展。MoE 是目前唯一被证明能够将深度学习模型扩展到万亿以上参数的方法，它能让模型学习更多信息，并为计算机视觉、语音识别、自然语言处理和机器翻译系统等提供支持，从而以全新的方式为人类社会提供帮助。
近日，微软亚洲研究院发布了一个用于促进大规模 DNN 模型开发的高性能 MoE 库——Tutel，并针对已普遍使用的新 Azure NDm A100 v4 系列进行了高度优化。借助 Tutel 多样化和灵活的 MoE 算法支持，AI 领域的开发人员可以更轻松、高效地执行 MoE。与最先进的 MoE 实现方式，如 fairseq（Meta的Facebook AI 研究院基于 PyTorch 的 Sequence to Sequence 工具包）相比，对于单个 MoE 层，Tutel 在具有8个 GPU 的单个 NDm A100 v4 节点上实现了8.49倍的加速，在具有512个 A100 GPU 的64个 NDm A100 v4 节点上实现了2.75倍的加速。在端到端性能方面，得益于 all-to-all 通信优化，Tutel 在 Meta（原 Facebook 公司）的1.1万亿参数的 MoE 语言模型中使用64个 NDm A100 v4 节点实现了40%以上的加速。
Tutel 具有良好的兼容性和丰富的功能，确保了其在 Azure NDm A100 v4 群集上运行时可以发挥出色的性能。目前，Tutel 已开源，并已集成到 fairseq 中。
Tutel GitHub 链接：
https://github.com/microsoft/tutel
Tutel MoE的三大优势
作为 fairseq、FastMoE 等其他高水平 MoE 解决方案的补充，Tutel 主要专注于优化面向 MoE 的计算和 all-to-all 通信，以及其他多样化和灵活的 MoE 算法支持。Tutel 具有简洁的接口，可以轻松集成到其他 MoE 解决方案中。当然，开发人员也可以从头开始，利用 Tutel 的接口将独立的 MoE 层合并到他们自己的 DNN 模型中，直接从高度优化的、最先进的 MoE 功能中受益。
与现有的 MoE 解决方案相比，Tutel 具有以下三个主要优势：
优化面向 MoE 的计算。由于缺乏高效的实现方法，目前基于 MoE 的 DNN 模型依赖于深度学习框架（如 PyTorch、TensorFlow 等）提供的多个现成 DNN 运算符的拼接来组成 MoE 计算。由于需要冗余计算，这种做法会导致显著的性能开销。Tutel 设计并实现了多个高度优化的 GPU 内核，为面向 MoE 的计算提供了运算符。例如，Tute l将调度“输出选通（gating output）”的时间复杂度从 O(N^3)降低到 O(N^2)，显著提高了数据调度的效率。Tutel 还实现了快速 cumsum-minus-one 运算符（fast cumsum-minus-one operator），与 fairseq 实现方式相比，达到了24倍的加速。此外，Tutel 还利用 NVRTC（CUDA C++ 的运行时编译库）进一步实时优化了定制的 MoE 内核。
图1对比了 Tutel 与 faireseq 在 Azure NDm A100 v4 平台上的运行结果，如前所述，使用 Tutel 的单个 MoE 层在8个 A100 GPU 上实现了8.49倍的加速，在512个 A100 GPU 上实现了2.75倍的加速。
图1：与 fairseq 相比，对于单个 MoE 层，Tutel 在具有8个 GPU 的单个 NDm A100 v4 节点上实现了8.49倍的加速，在具有512个 A100 GPU 的64个 NDm A100 v4 节点上实现了2.75倍的加速。详细设置为：batch_size = 32, sequence_length = 1,024, Top_K = 2, model_dim = 2,048, ands hidden_size = 2,048
在 Azure NDm A100 v4 群集上实现底层 all-to-all 通信优化。Tutel 针对 Azure NDm A100 v4 群集上的大规模 MoE 训练，优化了 all-to-all 聚合通信（collective communication），其中包括 CPU-GPU 绑定和自适应路由（AR）调整。在非一致存储访问结构（NUMA）系统上，尤其是在 NDv4 VM上，正确的 CPU-GPU 绑定对于 all-to-all 性能非常关键。但是，现有的机器学习框架（TensorFlow、PyTorch 等）并未提供高效的 all-to-all 通信库，导致大规模分布式训练的性能下降。
Tutel 可以自动优化绑定，并为用户微调提供简洁的接口。此外，Tutel 在 NDm A100 v4 集群上使用了多路径技术，即 AR。对于 MoE 中的 all-to-all 通信，每个 GPU 通信的总数据流量规模并不会发生变化，但每个 GPU 之间的数据规模会随着 GPU 数量的增加而变小。而更小的数据规模会在 all-to-all 通信中产生更大的开销，导致 MoE 训练性能下降。借助 Azure NDm A100 v4 集群提供的 AR 技术，Tutel 提高了小消息组的通信效率，并在 NDv4 系统上提供了高性能的 all-to-all 通信。得益于 CPU-GPU 绑定和AR调整，Tutel 使用512个 A100 GPU，对通常用于 MoE 训练的每次交换的数百兆字节的消息，实现了2.56倍到5.93倍的 all-to-all 加速，如图2所示。
图2：在应用 Tutel 前后，具有64个 NDm A100 v4 节点（512个 A100 GPU）的不同消息大小的 all-to-all 带宽。Tutel 使用512个 A100 GPU，对大小为数百兆字节的消息实现了2.56倍到5.93倍的 all-to-all 加速。
多样、灵活的 MoE 算法支持。Tutel 为最先进的 MoE 算法提供了多样化且灵活的支持，包括：
- 为 Top-K gating 算法设置任意K值（大多数实现方法仅支持 Top-1 和 Top-2 ）。
- 不同的探索策略，包括批量优先路由、输入信息丢失、输入抖动。
- 不同的精度级别，包括半精度（FP16）、全精度（FP32）、混合精度等（下一个版本中将支持 BF16）。
- 不同的设备类型，包括 NVIDIA CUDA 和 AMD ROCm 设备等。
Tutel 将积极整合来自开源社区的各种新兴 MoE 算法。
Tutel与Meta的MoE语言模型集成
此前，Meta 就开源了自己的 MoE 语言模型，并利用 fairseq 实现了 MoE。微软亚洲研究院与 Meta 合作将 Tutel 集成到了 fairseq 工具包中。Meta 也一直用 Tutel 在 Azure NDm A100 v4 上训练其大型语言模型，该模型中基于注意力的神经架构类似于 GPT-3。
微软亚洲研究院的研究员们使用 Meta 的语言模型评估了 Tutel 的端到端性能。该模型有32个注意力层，每个层都有32个128维注意力头（32x128-dimension heads）。每2层中包含1个 MoE 层，而每个 GPU 都配有一名专家（expert）。表1总结了模型的详细设置参数，图3则显示了Tutel的加速效果。由于 all-to-all 通信成为瓶颈，所以随着 GPU 数量的增加，Tutel 带来的改进从8个 A100 GPU 的131%提升至512个 A100 GPU 的40%。研究员们将在下一个版本中对此做进一步优化。
表1：512个 A100 (80G) GPU 的 MoE 语言模型配置
图3：在端到端性能方面，Tutel 在 Meta 的1.1万亿参数的 MoE 语言模型中使用64个 NDm A100 v4 节点实现了40%以上的加速。
Tutel将继续进化，充分发挥MoE潜力
MoE 技术前景广阔，它可以基于许多领域的技术，实现整体训练，例如系统路由、海量节点的网络平衡，甚至可以从基于 GPU 的加速中受益。微软亚洲研究院展示了一个高效的 MoE 实现方式——Tutel，它比 fairseq 框架有着显著的进步。
现在，Tutel 也已集成到了 DeepSpeed 框架中。相信 Tutel 和相关的集成将使更多 Azure 服务受益，特别是对于那些希望高效扩展自有大型模型的客户。今天的 MoE 仍处于早期阶段，需要付出更多努力才能充分发挥 MoE 的潜力，因此研究员们也将继续优化 Tutel，希望未来可以给大家带来更多令人兴奋的研究和应用成果。
作者：系统与网络组
来源：微软研究院AI头条
