DeepSpeed ZeRO-3 Offload
今天我们宣布发布 ZeRO-3 Offload，这是一个高效且易于使用的 ZeRO Stage 3 和 ZeRO Offload 的组合实现，旨在通过让每个人都能使用高效的大规模 DL 训练来实现我们的目标：让 AI 民主化。ZeRO-3 Offload 的主要优势是
- 前所未有的内存效率，可以在有限的 GPU 资源上运行超大型模型——例如，在一个 GPU 上微调超过 400 亿个参数的模型，在 512 个 GPU 上微调超过 2 万亿个参数的模型！
- 极易使用
- 扩展到超过万亿个参数，无需以复杂的方式组合多种并行技术。
- 对于现有的 DeepSpeed 用户，只需在 DeepSpeed 配置文件中设置几个标志即可开启 ZeRO-3 Offload。
- 分布式训练中每个 GPU 的高性能吞吐量以及跨 GPU 的超线性可扩展性。
- 在 1 万亿个参数的情况下，ZeRO-3 Offload 在 512 个 NVIDIA V100 GPU 上的计算性能达到 25 拍每秒，每个 GPU 达到 49 TFlops。
- 与 ZeRO-2 Offload 相比，单个 GPU 的吞吐量提高了 2 倍
ZeRO 技术系列概述
ZeRO 冗余优化器（简称 ZeRO）是一系列针对大型分布式深度学习的内存优化技术。与数据并行（高效但只能支持有限的模型规模）或模型并行（可以支持更大的模型规模，但需要大量的代码重构，同时增加通信开销，限制效率）不同，ZeRO 允许在不进行代码重构的情况下将更大的模型拟合到内存中，同时保持高效性。ZeRO 通过消除数据并行中固有的内存冗余，同时将通信开销降至最低，从而实现这一点。ZeRO 通过在数据并行进程之间划分三个模型状态（优化器状态、梯度和参数）而不是复制它们，来消除数据并行进程之间的内存冗余。通过这样做，它提高了与经典数据并行相比的内存效率，同时保留了其计算粒度和通信效率。ZeRO 有三个阶段，对应于三个模型状态，如图 1 所示：第一阶段（ZeRO-1）只划分优化器状态，第二阶段（ZeRO-2）划分优化器状态和梯度，最后阶段（ZeRO-3）划分所有三个模型状态（有关更多详细信息，请参阅 ZeRO 论文）。
图 1. ZeRO 内存节省概述
除了这三个阶段之外，ZeRO 技术系列还包括 ZeRO-2 Offload。ZeRO-2 Offload 是一种异构 DL 训练技术，它与 ZeRO-2 协同工作，将已划分的优化器状态和梯度卸载到 CPU 内存中。ZeRO-2 Offload 即使在单个 GPU 上也能提供 ZeRO-2 的全部内存优势，同时在多 GPU 设置中提供 ZeRO-2 的出色可扩展性。DeepSpeed 库自 2020 年 9 月起提供 ZeRO-2 Offload。有关详细信息，请参阅以下内容
ZeRO-3 Offload
随着今天发布 ZeRO-3 Offload，我们添加了对划分和卸载参数的支持，除了 ZeRO-2 Offload 在 DeepSpeed 中已经支持的优化器状态和梯度划分之外。通过参数划分，ZeRO-3 Offload 实现了 ZeRO 三个阶段中的全部功能，允许模型规模随着 GPU 数量线性增长。此外，ZeRO-3 Offload 还可以选择将所有这些模型状态卸载到 CPU 以进一步减少 GPU 内存消耗，利用 CPU 和 GPU 来最大限度地提高整个系统的内存和计算效率。
我们认为 ZeRO-3 Offload 在三个方面为大型模型训练带来了巨大的飞跃
i) 前所未有的模型规模，
ii) 轻松支持超大型模型，以及
iii) 达到出色的训练效率。
前所未有的模型规模
与 ZeRO-2 和 ZeRO-Offload 必须将参数拟合到单个 GPU 的内存中不同，ZeRO-3 Offload 可以跨 GPU 划分参数，并将它们卸载到 CPU，支持比单个 GPU 上的内存更大的模型规模。此外，ZeRO-3 Offload 超越了最先进的混合 3D 并行（数据并行、模型并行和流水线并行的组合）。虽然 3D 并行受到总 GPU 内存的限制，但 ZeRO-3 Offload 可以利用 GPU 和 CPU 内存，后者比 GPU 内存更大、更便宜。这使得 ZeRO-3 Offload 能够在给定的 GPU 和 CPU 资源下训练比任何其他现有技术更大的模型规模。
单个 GPU 上的模型规模：ZeRO-3 Offload 可以高效地在单个 GPU 上训练超过 400 亿个参数的模型（例如，32GB V100 GPU + 1.5TB CPU 内存）。这是当前最先进的 ZeRO-2 Offload 的 3 倍。
多 GPU 上的模型规模：使用 ZeRO-3 Offload，您可以在 NVIDIA 32GB V100 DGX-2 集群上分别使用 256 个 GPU 和 512 个 GPU 训练万亿参数模型和两万亿参数模型。相比之下，最先进的 3D 并行分别需要 800 个 GPU 和 1600 个 GPU 来拟合相同规模的模型。这表示拟合超过万亿个参数的模型所需的 GPU 数量减少了 3 倍。
轻松支持超大型模型
从系统角度来看，训练具有数百亿和数万亿个参数的模型极具挑战性。数据并行无法将模型规模扩展到超过数十亿个参数，模型并行（使用张量切片）由于巨大的通信开销，无法有效地将模型规模扩展到单个节点边界之外，而流水线并行无法扩展到模型中可用的层数，这限制了模型规模和它可以扩展到的 GPU 数量。
唯一现有的并行技术，它可以在大规模并行 GPU 集群上扩展到超过万亿个参数，是 3D 并行，它以复杂的方式组合了数据并行、模型并行和流水线并行。虽然这样的系统可以非常有效，但它需要数据科学家进行主要的模型代码重构，以将模型拆分为负载均衡的流水线阶段。这也使得 3D 并行在它可以支持的模型类型方面变得不灵活，因为具有复杂依赖图的模型无法轻松地转换为负载均衡的流水线。
ZeRO-3 Offload 通过两种方式解决了这些挑战
i) 凭借突破性的内存效率，ZeRO-3 和 ZeRO-3 Offload 是唯一一种可以有效地扩展到超过万亿个参数的 DL 并行技术，它不需要混合并行策略，从而大大简化了 DL 训练的系统堆栈。
ii) ZeRO-3 Offload 几乎不需要模型科学家进行模型重构，从而让数据科学家能够将复杂的模型扩展到数百亿到数万亿个参数。
出色的训练效率
多个节点上的每个 GPU 的高性能吞吐量：ZeRO-3 Offload 为多节点上的数十亿和数万亿参数模型提供了出色的训练效率。它在 32 个 DGX2 节点（包含 512 个 NVIDIA V100 GPU）上运行时，每个 GPU 的持续吞吐量高达 50 Tflops（见图 2）。相比之下，使用 PyTorch 的标准数据并行训练对于 12 亿参数模型（使用数据并行单独可以训练的最大模型）每个 GPU 只能达到 30 TFlops。
图 2. ZeRO-3 Offload：512 个 V100 GPU 上的数十亿和数万亿参数模型吞吐量
ZeRO-3 Offload 尽管对于固定批次大小，ZeRO Stage 3 的通信开销比标准数据并行训练高 50%，但它仍然获得了高效率。这得益于以通信重叠为中心的設計和实现，这使得 ZeRO-3 Offload 能够隐藏几乎所有通信量，同时利用更大的批次大小来提高由于更好的 GPU 内存效率而导致的效率。
单个 GPU 上的多十亿参数模型训练效率：ZeRO-3 Offload 通过使单个 GPU 上的多十亿参数模型的训练效率更高，进一步推动了 AI 民主化。对于单个 GPU 训练，ZeRO-3 Offload 在两个维度上都比 ZeRO-2 Offload 具有优势。首先，ZeRO-3 Offload 将单个 V100 上可训练的模型规模从 130 亿增加到 400 亿。其次，对于 ZeRO-3 Offload，它提供了比 ZeRO-2 Offload 更高的加速（例如，对于 130 亿，加速 2.3 倍），这两个解决方案都可以训练这些规模的模型。这些结果总结在图 3 中。
图 3. 一个 V100 GPU 上的多十亿参数模型训练
跨 GPU 的超线性可扩展性：此外，ZeRO-3 Offload 还保留了我们在所有以前的 ZeRO 技术（ZeRO Stage 1、ZeRO Stage 2 和 ZeRO Offload）中展示的超线性可扩展性特征。ZeRO-3 Offload 可以利用多 GPU 训练配置中所有 GPU 之间总的 PCI-E 带宽，同时它还可以利用所有节点的总 CPU 计算。因此，CPU-GPU-CPU 通信时间以及优化器更新时间分别随着 GPU 和节点数量的增加而线性减少，使得 ZeRO-3 Offload 可以表现出超线性扩展（见图 4）。
图 4. 2000 亿参数模型的 ZeRO-3 Offload 超线性可扩展性。
如何使用 ZeRO-3 Offload
与许多其他现有的 DeepSpeed 功能一样，一旦用户模型已转换为使用 DeepSpeed，启用 ZeRO-3 Offload 就如同在 DeepSpeed 配置文件中打开几个标志一样简单。支持像权重共享这样的高级功能，或启用需要跨 GPU/节点进行划分才能拟合到 GPU/CPU 内存的超大型模型，只需使用 ZeRO-3 Offload API 更改几行代码即可。
如果您已经是 DeepSpeed 用户，您可以在下面找到有关 ZeRO-3 Offload 的详细教程。如果您是 DeepSpeed 的新手，我们建议您在尝试使用 ZeRO-3 Offload 教程之前先从入门页面开始。
DeepSpeed 团队很高兴与 DL 社区分享 ZeRO-3 Offload。