Azure 使用 DeepSpeed 实现易于使用、高性能和超大规模模型训练

简介

近年来，在大量数据上训练的大型基于 Transformer 的深度学习模型在多个认知任务中取得了显著成果，并推动了增强人类能力的新产品和功能的出现。这些模型在过去五年中在规模上增长了几个数量级。从最初 Transformer 模型的几百万个参数，到最近的 5300 亿参数 Megatron-Turing 模型，如图 1所示。客户对以空前规模训练和微调大型模型的需求不断增长。

图 1：大型模型和硬件能力的景观

为了训练这些模型，用户需要设置和维护一个复杂的分布式训练基础设施，这通常需要执行多个手动步骤，且容易出错。这些导致在可用性和性能方面都体验不佳。我们最近宣布我们在简化这一过程方面取得了重大进展，并在 Azure 上实现了易于使用和高性能的 1K+ GPU 规模训练。

在这篇扩展文章中，我们分享了 DeepSpeed 用户如何在 Azure 上使用新的易于使用、简化、可扩展且高性能的分布式训练体验来训练万亿参数模型的详细信息。我们还分享了实验设置、模型配置、其他性能趋势的详细信息，并指导用户如何在自己的环境中运行这些实验。

使用 DeepSpeed 在 Azure 上加快和简化分布式训练

我们在图 2中将现有的手动且容易出错的工作流程与我们为 DeepSpeed 在 Azure 上提出的易于使用的工作流程进行了比较。客户现在可以使用易于使用的训练管道来启动大规模训练作业。如果用户依赖推荐的AzureML配方，新工作流程可将步骤数量从 11 个减少到 1 个。

图 2：使用 DeepSpeed 在 Azure 上实现易于使用且简化的分布式训练体验

对于使用 Azure VM 或Azure VMSS构建自定义环境的用户，只需要两个步骤

• 1) 运行集群设置脚本（将在未来几周内发布）
• 2) 使用 Azure VMSS配方来启动训练。

主要性能优势

我们已经在 Azure公告中分享了我们主要性能结果的摘要。与其他云提供商相比，我们能够训练 2 倍更大的模型规模（2 万亿参数 vs. 1 万亿参数），扩展到 2 倍的 GPU 数量（1024 个 vs. 512 个），并提供高达 1.8 倍的计算吞吐量/GPU（150 TFLOPs vs. 81 TFLOPs）。

DeepSpeed 在 Azure 上提供了近乎线性的可扩展性，无论是模型规模的增加还是GPU 数量的增加。如图 3a所示，结合 DeepSpeed 的ZeRO-3，其新颖的 CPU 卸载功能，以及由 InfiniBand 互连和 A100 GPU 提供支持的高性能 Azure 堆栈，我们能够以近乎线性的方式保持高效的吞吐量/GPU（>157 TFLOPs），因为模型规模从 1750 亿个参数增加到 2 万亿个参数。另一方面，对于给定的模型规模，例如 175B，当我们从 128 个 GPU 增加到 1024 个 GPU 时，我们实现了近乎线性的扩展，如图 3b所示。关键的结论是，Azure 和 DeepSpeed 共同打破了 GPU 内存壁垒，使我们的客户能够轻松高效地训练大规模的万亿参数模型。

图 3：(a) 当我们从 1750 亿个参数增加到 2 万亿个参数时，吞吐量/GPU 几乎保持完美（BS/GPU=8）。(b) 175B 模型中 GPU 设备数量增加时的近乎完美性能扩展（BS/GPU=16）。这两种情况下的序列长度均为 1024。

实验设置

我们分享了实验设置的详细信息以及我们遵循的一些最佳实践。用户可以直接使用它们来重现我们的结果，或者修改它们以适应自己的模型规模和配置的 Azure 硬件规模。

硬件 (Azure 实例)

我们在实验中使用了NDm A100 v4 系列 实例。每个实例包含两个插座 AMD EPYC 7V12 64 核 CPU、1.7TB 主内存和八个 A100 80GB GPU。该系统具有平衡的 PCIe 拓扑结构，将 4 个 GPU 设备连接到每个 CPU 插座。VM 中的每个 GPU 都配备了专用的、与拓扑无关的 200 Gb/s NVIDIA Mellanox HDR InfiniBand 连接，提供加速的 200 Gbps 高速结构。DeepSpeed 库利用卸载功能，其中激活和优化器状态分配在主内存中。因此，每个节点 1.7TB 的内存容量有助于我们扩展到大型模型规模。

使用 AzureML 进行训练设置

用户可以直接使用 AzureML 工作室并使用我们发布的配方来运行实验，无需任何其他设置。这是在 Azure 上运行实验最简单且推荐的方式。

使用 Azure VMSS 进行训练设置

现有的 VMSS 客户和其他拥有基于 Azure VM 的自定义环境的用户可以按照以下设置进行。简化这些步骤的脚本将在未来几周内发布。使用 Azure 虚拟机规模集 (VMSS) 创建一个集群，以配置所需数量的计算节点，这些节点运行新的 Azure HPAI VM 映像，该映像专门用于使用表 1中列出的软件堆栈进行极大规模的深度学习应用程序。

名称	描述（版本）
PyTorch	1.10.2（从源代码安装）
DeepSpeed	0.6.2（从源代码安装）
Megatron-LM	https://github.com/microsoft/Megatron-DeepSpeed
Apex	0.1
NCCL	2.12.10
CUDNN	8.2.4.15
CUDA	11.4
CUDA 驱动程序	R470.82
VM 映像	Ubuntu-HPC 20.04 映像

表 1：Azure HPC VM 映像中软件包的详细版本信息

用户可以创建最多包含 600 个 VM 实例的 VMSS，从而支持最多 4800 个 A100 GPU。除了计算节点的 VMSS 之外，我们还使用廉价的 D4s v4（或类似的）实例配置了一个独立的登录节点，该实例具有 4 核 Intel VCPU，运行相同的映像，用于编译、启动和监控作业。登录节点、计算节点和共享存储文件系统分组在 Azure 虚拟网络 (vnet) 中，允许 VM 通过 SSH 和共享 NFS 卷相互连接，如图 4所示。

图 4：基于 VMSS 的实验设置的组织方式

对各种模型配置进行性能评估

我们使用四种不同的模型规模（175B、530B、1T 和 2T）运行了实验，使用了表 2中所示的配置。

模型规模	175B	530B	1T	2T
层数	96	105	128	160
隐藏维度	12,288	20,480	25,600	32,768
注意力头数	96	128	160	128

表 2：模型配置

对于每种配置，我们使用 TFLOPs/GPU 作为主要性能指标，报告系统的峰值吞吐量。为了计算 TFLOPs，我们使用 Megatron 论文中使用的公式，如下所示。

FLOPs/GPU = 96 * B * s * l * h2 * (1 + s/6h + V/(16*l*h))

B 是批次大小，s 是序列长度，l 是层数，h 是隐藏大小，V 是词汇量。

扩展 175B 和 530B 模型

图 5a和5b分别显示了序列长度为 512 和 1024 的 175B 模型的结果。我们只将 512 个 GPU 扩展到序列长度为 512，因为添加更多 GPU 显示出相似的性能。另一方面，对于序列长度为 1024，我们观察到线性性能提高到 1024 个 GPU。总的来说，204.49 TFLOPs/GPU的峰值吞吐量是在 256 个 GPU 上使用 32 的微批次大小和 512 的序列长度实现的。

图 5：175B 模型在 512 个和 1K 个 GPU 上的性能特征。彩色列表示不同的微批次大小。

接下来，我们报告 530B 模型的扩展情况。使用 DeepSpeed 和 Megatron-LM 在 Selene 超级计算机上的 280 个 DGX A100 服务器上进行的 530B MT-NLG 模型的先前结果显示峰值吞吐量为 126 TFLOPS/GPU。但是，我们能够超过该吞吐量，并在 128 个 NDm A100 v4 系列 A100 系统（即 1024 个 GPU）上实现了高达171.37 TFLOPs/GPU的性能，如图 6所示。

这个 530B 模型的优势在于它的并行化配置更简单，因为没有张量/流水线并行。使用 ZeRO 驱动的并行，所需的启发式方法更少，以便最佳地配置分布式模型。此外，对于微批次大小 >1，超过 140 TFLOPs/GPU 的一致稳态性能证明了强大的软件和硬件平台。

图 6：使用 530B 参数模型在 512 个和 1024 个 GPU 上实现的吞吐量，每个 GPU 的微批次大小为 1、2、4 和 8，序列长度为 1024。

扩展 1T 和 2T 模型

1T 参数模型包含 128 层，具有 160 个注意力头。训练如此极大规模的模型并非易事。图 7显示了我们在 512 个和 1024 个 GPU 上探索的每个模型配置所实现的吞吐量。在 1024 个 GPU 上使用 8 的微批次大小实现了165.36 TFLOPs/GPU的峰值吞吐量，并且模型在最初的 3-4 个迭代中达到了稳态性能。

图 7：1T 参数模型在 512 个和 1024 个 GPU 上的性能特征，微批次大小为 1、2、4 和 8，序列长度为 1024。

2T 参数模型包含 160 层、32k 隐藏维度和 128 个注意力头。鉴于模型规模很大，并且在 1024 个 GPU 上需要相当长的时间，我们将 2T 模型的基准测试运行限制在每个 GPU 8 个批次大小，序列长度为 1024。我们在 1024 个 GPU 上实现了 157 TFLOPs/GPU 的性能。

如何在 Azure 上运行训练实验？

我们认识到 DeepSpeed 用户是多元的，并且拥有不同的环境。在本教程中，我们专注于简化计划在 Azure 上运行大型模型训练实验的用户的工作流程。

在 Azure 上进行模型训练最简单的方法是通过 Azure ML 配方。作业提交和数据准备脚本已在此处提供提供。用户只需按照指南设置 Azure ML 工作区，并使用 aml_submit.py 文件提交实验。

一些用户拥有在 Azure VM 和基于 VMSS 的集群之上构建的自定义环境。为了简化在这些设置上的训练，我们正在努力开发一个易于使用的集群设置脚本，该脚本将在未来几周内发布。如果您已经拥有正在运行的集群设置，可以使用Azure 配方来训练 175B 模型和 1T 模型。可以轻松修改这些配方来训练其他模型配置。

致谢

这篇博客文章由 DeepSpeed 团队与 AzureML 和 AzureHPC 团队合作撰写。我们感谢以下使这项工作成为可能的个人

• AzureHPC 团队：Russell J. Hewett、Kushal Datta、Prabhat Ram、Jithin Jose 和 Nidhi Chappell
• AzureML 团队：Vijay Aski、Razvan Tanase、Miseon Park、Savita Mittal、Ravi Shankar Kolli、Prasanth Pulavarthi 和 Daniel Moth
• DeepSpeed 团队：Ammar Ahmad Awan、Jeff Rasley、Samyam Rajbhandari、Martin Cai 和 Yuxiong He
• 首席技术官办公室：Gopi Kumar 和 Luis Vargas