DeepSpeed ZeRO-3 卸载
今天我们宣布发布ZeRO-3 Offload，这是一个高效且易于使用的ZeRO Stage 3和ZeRO Offload结合的实现，旨在通过使高效的大规模深度学习训练对所有人可用，继续实现我们民主化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 PetaFlops的计算性能，达到49 TFlops/GPU。
- 与单GPU上的ZeRO-2 Offload相比，吞吐量提高了最多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是一种异构深度学习训练技术，它与ZeRO-2协同工作，将分区的优化器状态和梯度卸载到CPU内存中。即使在单个GPU上，ZeRO-2 Offload也能提供ZeRO-2的全部内存优势，同时在多GPU设置上提供ZeRO-2的出色可扩展性。自2020年9月以来，DeepSpeed库一直提供ZeRO-2 Offload。详情请见下文：
ZeRO-3 卸载
随着今天发布的ZeRO-3 Offload，我们增加了对参数分区和卸载的支持，除了DeepSpeed中ZeRO-2 Offload已经支持的优化器状态和梯度分区。通过参数分区，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（例如，32GB V100 GPU + 1.5TB CPU内存）上高效训练超过400亿参数的模型。这比当前最先进的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 是唯一能够独立高效扩展到超过一万亿参数的深度学习并行技术，无需混合并行策略，极大地简化了深度学习训练的系统堆栈。
ii) ZeRO-3 Offload 几乎不需要模型科学家进行模型重构，使数据科学家能够将复杂模型扩展到数千亿到数万亿参数。
卓越的训练效率
多节点上每个GPU的高性能吞吐量：ZeRO-3 Offload 在多个节点上为数十亿和万亿参数模型提供了出色的训练效率。它在由512个NVIDIA V100 GPU组成的32个DGX2节点上运行时，每个GPU的持续吞吐量可达50 Tflops（见图2）。相比之下，使用PyTorch进行标准数据并行训练时，对于1.2B参数模型（这是仅使用数据并行可以训练的最大模型），每个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上可训练的模型大小从13B增加到40B。其次，对于两种解决方案都可训练的模型大小，ZeRO-3 Offload相比ZeRO-2 Offload提供了加速（例如，13B模型的2.3倍加速）。这些结果总结在图3中。
图3. 在单个V100 GPU上进行数十亿参数模型的训练
跨GPU的超线性扩展性：此外，ZeRO-3 Offload还保留了我们在之前所有ZeRO技术（ZeRO Stage 1、ZeRO Stage 2和ZeRO Offload）中展示的超线性扩展特性。ZeRO-3 Offload可以在多GPU训练配置中利用GPU和CPU之间的聚合PCI-E带宽，同时也可以利用所有节点的聚合CPU计算能力。因此，CPU-GPU-CPU通信时间以及优化器更新时间分别随着GPU和节点数量的增加而线性减少，使得ZeRO-3 Offload能够展现出超线性扩展（参见图4）。
图4. 200B参数模型的ZeRO-3卸载超线性可扩展性。
如何使用ZeRO-3卸载
与许多其他现有的DeepSpeed功能一样，一旦用户模型已转换为使用DeepSpeed，启用ZeRO-3 Offload就像在DeepSpeed配置文件中打开几个标志一样简单。支持高级功能，如权重共享，或启用需要跨GPU/节点分区以适应GPU/CPU内存的极大模型，只需使用ZeRO-3 Offload API进行几行额外的代码更改即可实现。
如果您已经是DeepSpeed用户，您可以在下面找到我们关于ZeRO-3 Offload的详细教程。如果您是DeepSpeed的新手，我们建议您在尝试我们的ZeRO-3 Offload教程之前，先从入门页面开始。
DeepSpeed团队非常高兴与深度学习社区分享ZeRO-3 Offload。