混合专家(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% 以上的加速。
MoE 技术前景广阔,它可以基于许多领域的技术,实现整体训练,例如系统路由、海量节点的网络平衡,甚至可以从基于 GPU 的加速中受益。微软亚洲研究院展示了一个高效的 MoE 实现方式 ——Tutel,它比 fairseq 框架有着显著的进步。
现在,Tutel 也已集成到了 DeepSpeed 框架中。相信 Tutel 和相关的集成将使更多 Azure 服务受益,特别是对于那些希望高效扩展自有大型模型的客户。今天的 MoE 仍处于早期阶段,需要付出更多努力才能充分发挥 MoE 的潜力,因此研究员们也将继续优化 Tutel,希望未来可以给大家带来更多令人兴奋的研究和应用成果。