ZeRO-Offload: Democratizing Billion-Scale Model Training

文章标题: ZeRO-Offload：让十亿规模模型的训练大众化
作者/机构: Jie Ren∗, Samyam Rajbhandari†, Reza Yazdani Aminabadi†, Olatunji Ruwase†, Shuangyan Yang∗, Minjia Zhang†, Dong Li∗, Yuxiong He†
†Microsoft, ∗University of California, Merced

A1 主要贡献

本文旨在解决大规模深度学习模型训练因需要复杂模型重构和昂贵GPU集群而仅限于少数研究者的核心问题。研究目标是通过开发一种新颖的异构深度学习训练技术 ZeRO-Offload，让大规模模型训练变得普及化，使数据科学家即便只有单块GPU也能进行。

核心创新点如下：

A3 背景知识与相关工作

大模型训练中的内存消耗
DNN模型训练过程中的内存消耗可分为两部分：模型状态（model states）和残差状态（residual states）【索引21, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. Samyam Rajbhandari et al. 2020. SC】。模型状态包括参数、梯度和优化器状态（如Adam优化器【索引13, Adam: A method for stochastic optimization. Diederik P. Kingma and Jimmy Ba. 2014】中的动量和方差）；残差状态包括激活值、临时缓冲区和不可用的内存碎片。

模型状态是大模型训练的主要内存瓶颈
对于像Megatron-LM（80亿参数）【索引28, Megatron-lm: Training multi-billion parameter language models using model parallelism. Mohammad Shoeybi et al. 2019. CoRR】、T5（110亿参数）【索引20, Exploring the limits of transfer learning with a unified text-to-text transformer. Colin Raffel et al. 2020】和Turing-NLG（172亿参数）【索引25, Turing-nlg: A 17-billion-parameter language model by microsoft. Corby Rosset. 2020】这样的大型Transformer模型，模型状态是主要的内存瓶颈。这些模型通常使用float-16（fp16）混合精度训练【索引16, Automatic Mixed Precision for Deep Learning. Nvidia. 2019】和Adam优化器【索引13, Adam: A method for stochastic optimization. Diederik P. Kingma and Jimmy Ba. 2014】。在混合精度训练中，通常会保留fp16和fp32两份参数副本。梯度以fp16存储。Adam优化器还需存储fp32的动量和方差。因此，一个拥有M个参数的模型总共需要 $16 \times M$ 字节的内存。据此计算，Megatron-LM、T5和Turing-NLG的模型状态分别需要128GB、176GB和284GB内存，这远超当前旗舰级NVIDIA A100 GPU（80GB）的内存容量。

现有大模型训练方法分类
近年来，为解决单GPU内存无法容纳模型和残差状态的问题，涌现了大量工作，可大致分为两类：横向扩展（scale-out）训练和纵向扩展（scale-up）训练。

横向扩展（Scale out）大模型训练
横向扩展训练利用多块GPU的聚合内存来满足大模型训练的内存需求。典型的例子是模型并行【索引5, Large scale distributed deep networks. Jeffrey Dean et al. 2012. NIPS】、【索引28, Megatron-lm: Training multi-billion parameter language models using model parallelism. Mohammad Shoeybi et al. 2019. CoRR】和流水线并行【索引7, Pipedream: Fast and efficient pipeline parallel dnn training. Aaron Harlap et al. 2018】、【索引10, Gpipe: Efficient training of giant neural networks using pipeline parallelism. Yanping Huang et al. 2018】，它们都将模型状态和残差状态划分到多个GPU上。模型并行垂直切分模型，而流水线并行则按层水平切分模型。这两种方法都需要修改用户模型，限制了易用性。最近的ZeRO【索引21, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. Samyam Rajbhandari et al. 2020. SC】提供了一种替代方案，它像数据并行一样在多GPU间划分训练批次，但它会对模型状态进行分区而不是复制，从而无需修改用户模型，并提供了更好的计算效率和扩展性。尽管这些技术能够训练大模型，但它们都要求GPU总内存足以容纳所有模型状态。相比之下，ZeRO-Offload通过将模型状态卸载到CPU内存来适配更大的模型，在单GPU上即可训练大10倍的模型，并且可以与ZeRO或模型并行结合以实现更好的扩展性。

纵向扩展（Scale up）大模型训练
现有工作主要通过三种方式在单GPU上扩展模型尺寸。第一种是通过激活检查点（activation checkpointing）技术【索引4, Training deep nets with sublinear memory cost. Tianqi Chen et al. 2016. arXiv: Learning】，用计算换取激活值（残差内存）的节省。第二种是使用压缩技术，如低精度或混合精度训练【索引16, Automatic Mixed Precision for Deep Learning. Nvidia. 2019】，以节省模型状态和激活值的内存。第三种是将外部存储（如CPU内存）作为GPU内存的扩展【索引8, Autotm: Automatic tensor movement in heterogeneous memory systems using integer linear programming. Mark Hildebrand et al. 2020. ASPLOS】、【索引9, Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping. Chien-Chin Huang et al. 2020. ASPLOS】、【索引11, Layer-centric memory reuse and data migration for extreme-scale deep learning on many-core architectures. Hai Jin et al. 2018. ACM Trans. Archit. Code Optim.】、【索引17, Capuchin: Tensor-based gpu memory management for deep learning. Xuan Peng et al. 2020. ASPLOS】、【索引23, Sentinel: Efficient Tensor Migration and Allocation on Heterogeneous Memory Systems for Deep Learning. Jie Ren et al. 2020. HPCA】、【索引24, vdnn: Virtualized deep neural networks for scalable, memory-efficient neural network design. Minsoo Rhu et al. 2016. MICRO-49】、【索引33, Superneurons: Dynamic gpu memory management for training deep neural networks. Linnan Wang et al. 2018. PPoPP】。ZeRO-Offload属于第三种。与这些工作不同，ZeRO-Offload不仅卸载数据，还卸载计算。近期一个名为L2L【索引18, Training large neural networks with constant memory using a new execution algorithm. Bharadwaj Pudipeddi et al. 2020】的工作通过逐层管理GPU内存来支持数十亿参数的训练，但其效率因额外的通信开销而受限，且不支持多设备扩展，并需要模型重构。

ZeRO驱动的数据并行训练
ZeRO-Offload与ZeRO协同工作以扩展到多GPU。ZeRO有三个阶段，分别对应优化器状态（ZeRO-1）、梯度（ZeRO-2）和参数（ZeRO-3）的划分。ZeRO-Offload与ZeRO-2共生，因此我们在此详细讨论ZeRO-2。在ZeRO-2中，每块GPU存储所有参数的副本，但在每个训练步骤结束时只更新其中互斥的一部分。因此，每块GPU只需存储与其更新部分对应的优化器状态和梯度。更新后，每块GPU通过all-gather通信原语将其更新后的参数部分发送给所有其他GPU。ZeRO-2的计算和通信调度如下：在前向传播中，每块GPU使用不同的小批量数据计算损失。在后向传播中，每个梯度计算出来后，会通过reduce操作在拥有该梯度的GPU上进行平均。后向传播结束后，每块GPU使用平均后的梯度更新其负责的参数和优化器状态部分。之后，执行一次all-gather操作以接收其他GPU计算的其余参数更新。

A2 方法细节

3. 独特的、最优的卸载策略

ZeRO-Offload的设计目标
ZeRO-Offload旨在通过将部分模型状态从GPU卸载到CPU内存，来支持在单块或多块GPU上进行高效的大模型训练。模型状态（参数、梯度和优化器状态）是大模型训练的主要内存瓶颈。通过卸载这些状态，ZeRO-Offload可以训练更大的模型。然而，确定最优的卸载策略并非易事，因为不同的策略在CPU计算和GPU-CPU通信方面有不同的权衡，这两者都可能限制训练效率。

通过数据流图和第一性原理分析寻找最优策略
为了识别最优卸载策略，ZeRO-Offload将深度学习训练建模为一个数据流图，并使用第一性原理分析来有效地在CPU和GPU设备间划分此图。ZeRO-Offload的图划分方式在三个关键方面达到最优：1) 它要求CPU上的计算量比GPU上少几个数量级，从而防止CPU计算成为性能瓶颈；2) 它保证了CPU和GPU内存之间的通信量最小化；3) 它在实现最小通信量的同时，可证明地最大化了GPU上的内存节省。实际上，ZeRO-Offload能达到与非卸载训练相当的效率，并且是唯一最优的，即没有其他方案能在不增加通信量或CPU计算量的情况下提供更多的内存节省。本节将详细阐述这一独特最优卸载策略的推导过程，该策略专为使用Adam优化器的大模型混合精度训练而设计。

图2：具有M个参数的全连接神经网络的数据流。我们使用激活检查点来减少激活内存，以避免激活在CPU和GPU之间迁移。
图2：具有M个参数的全连接神经网络的数据流。我们使用激活检查点来减少激活内存，以避免激活在CPU和GPU之间迁移。

3.1 将深度学习训练视为数据流图

数据流图的构建
深度学习训练工作负载可以表示为一个带权重的有向数据与计算图，如图2所示。其中，圆形节点代表模型状态（parameter16, gradient16, parameter32, momentum32, variance32），矩形节点代表计算（forward, backward, param update）。图中的边表示节点之间的数据流，边的权重是在任何给定训练迭代中流经该边的数据总量（以字节为单位）。对于一个有M个参数的模型，图中边的权重要么是2M（当源节点产生fp16模型状态时），要么是4M（当源节点产生fp32模型状态时）。

图划分与卸载策略
GPU和CPU之间的卸载策略可以表示为对该图的两路划分。一个分区中的计算节点将在拥有该分区的设备上执行，而数据节点将存储在该设备的内存中。GPU和CPU之间必须通信的总数据量由跨越两个分区的边的权重之和给出。存在多种划分此图的方式。在接下来的小节中，我们将使用第一性原理，基于三个不同的效率指标来简化数据流图以减少可能的选择：1) CPU计算开销，2) 通信开销，和 3) 内存节省。

3.2 限制CPU计算

避免卸载计算密集型组件
CPU的计算吞吐量比GPU慢几个数量级。因此，将大的计算图卸载到CPU会严重限制训练效率。我们必须避免将计算密集型组件卸载到CPU。

基于计算复杂度的划分
深度学习训练每次迭代的计算复杂度通常为$O(MB)$，其中M是模型大小，B是有效批量大小。为避免CPU计算成为瓶颈，只有计算复杂度低于$O(MB)$的计算才应被卸载到CPU。这意味着前向传播和后向传播（两者计算复杂度均为$O(MB)$）必须在GPU上完成，而剩下的计算，如范数计算、权重更新等（复杂度为$O(M)$），可以被卸载到CPU。基于这个简单的观察，我们将数据流图中的前向和后向节点融合成一个单一的超级节点（FWD-BWD），并将其分配在GPU上。

3.3 最小化通信量

通信开销的考量
CPU内存带宽比CPU和GPU之间的PCI-E带宽快至少一个数量级，而GPU内存又比CPU内存快一个数量级。因此，我们必须最小化CPU和GPU内存之间的通信量，以防止PCI-E带宽成为训练性能的瓶颈。为此，我们首先需要确定模型状态卸载策略的理论最小通信量。

理论最小通信量
任何模型状态卸载策略的最小通信量为4M。请注意，在将前向和后向传播融合成一个超级节点后（如3.2节所述），我们数据流图中的每个节点都成为一个环路的一部分。因此，对该图的任何划分都至少需要切断两条边，每条边的权重至少为2M，导致总通信量至少为4M。

为达到最小通信量而简化的划分策略
如果我们选择将通信量限制在这个最低限度，我们可以极大地简化数据流图，并将划分策略的数量减少到少数几个。
* 创建fp32超级节点：任何不将fp32模型状态与其生产者和消费者节点共同放置的划分策略都无法实现4M的最小通信量。这样的划分必须切断至少一条权重为4M的边和另一条权重至少为2M的边，导致通信量至少为6M。因此，为实现最小通信量，所有卸载策略必须将fp32模型状态与其生产者和消费者操作符共同放置，即fp32模型状态（momentum32, variance32 和 p32）必须与参数更新（Param Update）和float2half计算共同放置。这个约束使我们能将数据流图中所有上述fp32数据和计算节点视为一个单一的超级节点，我们称之为“更新超级节点”（Update Super）。简化后的数据流图如图2所示，只包含四个节点：FWD-BWD超级节点、p16数据节点、g16数据节点和Update Super节点。
* p16的分配：为了达到最小通信量，p16必须与FWD-BWD超级节点共同放置，因为这两个节点之间的边权重是4M。分离这两个节点会使通信量增加到6M（4M + 2M）。由于我们已经将FWD-BWD超级节点分配给GPU以限制CPU计算，因此p16也必须分配给GPU。

3.4 最大化内存节省

最终的划分选择
在为最小化通信量而简化数据流图后，只剩下g16和Update Super节点需要分配。此时，所有划分方案都将产生最小的通信量，因此我们可以根据最大化GPU内存节省来进一步筛选选择。表1显示了所有最小化通信量的有效划分策略的内存节省情况。通过将g16和Update Super都卸载到CPU，可以实现8倍的最大内存节省。

表1：与基线相比，最小化通信量的卸载策略的内存节省情况。
表1：与基线相比，最小化通信量的卸载策略的内存节省情况。

3.5 一个独特且最优的卸载策略

ZeRO-Offload的最终策略
ZeRO-Offload将所有fp32模型状态以及fp16梯度分配在CPU内存中，并且它还在CPU上计算参数更新。fp16参数保留在GPU上，前向和后向计算也都在GPU上完成。

策略的唯一最优性
我们通过简化数据流图并排除所有其他划分策略得出了这个卸载策略，因为其他策略要么没有限制CPU计算，要么没有最小化通信量，要么没有最大化内存节省。因此，ZeRO-Offload不仅在上述指标方面是最优的，而且是唯一的；不可能有其他策略能在不增加CPU计算复杂度或产生额外GPU-CPU通信量的情况下，提供比ZeRO-Offload更多的内存节省。

图3：ZeRO-Offload在单块GPU上的训练过程。
图3：ZeRO-Offload在单块GPU上的训练过程。

4. ZeRO-Offload调度

本节我们讨论基于前述卸载策略在单GPU系统上实现ZeRO-Offload的具体计算和通信调度。然后，我们展示如何通过将我们的卸载策略与ZeRO数据并行和模型并行相结合，来将此调度有效地扩展到多GPU系统。

4.1 单GPU调度

数据划分与训练流程
如第3节所述，ZeRO-Offload的数据划分策略是将fp16参数存储在GPU上，而将fp16梯度以及所有优化器状态（如fp32动量、方差和参数）存储在CPU上。训练开始时，通过前向传播计算损失。由于fp16参数已在GPU上，这部分计算无需CPU通信。在对损失进行后向传播期间，不同参数的梯度在后向调度中的不同时间点计算出来。ZeRO-Offload可以在每个参数的梯度计算完成后，立即将其单独或以小组形式传输到CPU内存。因此，在梯度被传输到CPU内存之前，只需要少量内存临时保存在GPU上。此外，每次梯度传输都可以与后向图剩余部分的后向传播重叠，使得ZeRO-Offload能够隐藏大部分通信成本。

参数更新与同步
后向传播结束后，ZeRO-Offload直接在CPU上更新fp32参数和其他优化器状态（如动量和方差），然后将更新后的fp32参数从CPU内存复制回GPU内存中的fp16参数。图3以图表形式展示了ZeRO-Offload每一步的计算和通信，图5则以伪代码形式给出了具体的调度。

图5：代表ZeRO-Offload的代码，它将独特的最优CPU卸载策略与ZeRO驱动的数据并行相结合。
图5：代表ZeRO-Offload的代码，它将独特的最优CPU卸载策略与ZeRO驱动的数据并行相结合。

图4：ZeRO-Offload在多GPU下的数据放置。
图4：ZeRO-Offload在多GPU下的数据放置。

4.2 扩展到多GPU

与ZeRO数据并行的共生集成
ZeRO-Offload的完整形态是第3节描述的ZeRO-Offload策略与第2节讨论的ZeRO驱动的数据并行的共生集成，这使得ZeRO-Offload能够高效地扩展到数百个GPU。ZeRO-Offload保留了ZeRO Stage-2的模型状态划分策略（即优化器状态和梯度划分），同时将划分后的梯度、优化器状态以及相应的参数更新卸载到CPU。这种先划分再卸载的关键好处在于，对于超过1个GPU的系统，每个数据并行进程只负责更新参数的一个子集。所有数据并行GPU到CPU的总通信量保持不变，并且CPU资源被并行使用，共同计算一次权重更新。因此，总的CPU更新时间随着数据并行度的增加而减少，因为CPU计算资源随计算节点数量的增加而线性增长。这使得ZeRO-Offload能够实现非常好的可扩展性，因为跨GPU的通信开销被CPU优化器步骤时间的减少所抵消。

多GPU调度流程
ZeRO-Offload在不同GPU之间划分梯度和优化器状态，每个GPU将其拥有的分区卸载到CPU内存中，并在整个训练过程中保持在那里。在后向传播期间，梯度被计算并通过GPU上的reduce-scatter进行平均，每个GPU只将其分区所属的平均梯度卸载到CPU内存。一旦梯度在CPU上可用，每个数据并行进程直接在CPU上并行更新其优化器状态分区。更新后，参数分区被移回GPU，随后在GPU上执行类似于ZeRO-2的all-gather操作以收集所有参数。图4展示了ZeRO-Offload的模型参数、梯度和优化器状态的数据放置模型，而图5中详细介绍了ZeRO-Offload数据并行的调度细节。图中描述的all-gather操作被展示为一系列广播操作。

与模型并行的结合
ZeRO-Offload还可以与基于张量切片的模型并行（MP）框架如Megatron-LM【索引28, Megatron-lm: Training multi-billion parameter language models using model parallelism. Mohammad Shoeybi et al. 2019. CoRR】协同工作。它通过卸载每个MP进程对应的梯度、优化器状态和优化器计算，使得ZeRO-Offload能够训练比单独使用模型并行大得多的模型。第6节提供了更多细节。

5. 优化的CPU执行

我们通过两种优化来加速参数更新的CPU执行时间。首先，我们使用高性能计算技术实现了一个快速的CPU Adam优化器，相比先进的Pytorch实现提供了显著的加速。其次，我们开发了一种单步延迟参数更新调度，将CPU参数更新计算与GPU上的前向和后向计算重叠，从而在启用时隐藏CPU执行时间。

5.1 实现CPU优化器

并行化优化技术
我们使用三个层次的并行来提升CPU优化器的性能：1) SIMD向量指令【索引15, Use of simd vector operations to accelerate application code performance on lowpowered arm and intel platforms. Gaurav Mitra et al. 2013】，以充分利用CPU架构支持的硬件并行性。2) 循环展开【索引31, The performance impact analysis of loop unrolling. G. Velkoski et al. 2014】，这是一种增加指令级并行度的有效技术，对于更好地利用内存带宽至关重要。3) OMP多线程，用于有效并行利用CPU上的多个核心和线程。通过这些技术，我们实现了一个比先进的PyTorch实现快得多的Adam优化器。

混合精度训练中的Adam优化器
ADAM是一种用于深度学习训练的优化算法，它利用损失梯度及其一阶和二阶矩来更新参数。因此，除了模型参数外，ADAM在训练期间还需要保存另外两个同样大小（M）的矩阵。在混合精度训练模式下，内存中存储了两个版本的参数：一个FP16版本（p16）用于在GPU上计算前向传播中的激活值，另一个是FP32的主副本（p32），由优化器（在CPU上）更新。在每个训练步骤中，p16通过float2half转换由p32更新。此外，动量和梯度的方差保存在FP32中（在CPU上），以防止更新参数时的精度损失。有关ADAM算法的更多细节，请参考【索引13, Adam: A method for stochastic optimization. Diederik P. Kingma and Jimmy Ba. 2014】。

优化实现细节
算法2详细说明了使用SIMD操作的ADAM实现细节。如算法所示，Adam函数接收优化器参数（如β1、β2和α）以及梯度、动量、方差和参数的主副本（p32）作为输入。我们还使用了一些实现特定的参数，如simd_width和unroll_width。Adam优化器返回更新后的方差、动量和FP16（到GPU）及FP32（到CPU）的参数。我们首先将数据（包括参数、梯度、动量和方差）读入向量寄存器（第7行）。然后，我们使用多个融合乘加（FMA）向量操作来执行主执行流水线，该流水线由展开宽度重复。注意，其余操作（如乘法、除法和平方根）也以向量模式运行。为获得最佳性能，我们根据自动调优结果使用AVX512 SIMD指令集和8的unroll_width。除了CPU-Adam优化器，我们还以分块方式实现了CPU到GPU的16FP参数拷贝（第15行）。我们通过并行化Adam计算和将参数复制到GPU来重叠CPU和GPU的执行。当我们在CPU上处理当前数据块的Adam计算时，我们将先前处理过的数据块的参数写回GPU。这样，我们减少了GPU开始处理下一个训练步骤的空闲时间。

Algorithm 2 CPU-ADAM Optimizer

Input: p32, g32, m32, v32, β1, β2, α , step, eps
Output: p16, p32, m32, v32
Parameter: tile_width, simd_width, unroll_width
1: biascorrection1 ← -α / (1 - β1^step)
2: biascorrection2 ← 1 / sqrt(1 - β2^step)
3: simd_count ← sizeof(p32) / simd_width
4: unroll omp parallel
5: for i in 1 to (simd_count / unroll_width) do
6: ...
7: gv, pv, mv, vv = g32[i], p32[i], m32[i], v32[i]
8: mv = FMA(gv, (1 - β1), β1 * mv)
9: vv = FMA(gv * gv, (1 - β2), β2 * vv)
10: gv = FMA(sqrt(vv), biascorrection2, eps)
11: gv = mv / gv
12: pv = FMA(gv, biascorrection1, pv)
13: p32[i], m32[i], v32[i] = pv, mv, vv
14:
15: IF (i == tile_width) Copy_to_GPU(p16, p32)
16: end for
5.2 单步延迟参数更新

解决小批量训练瓶颈
尽管使用了高度优化的CPU优化器，但在非常小的批量大小下，当GPU计算时间不比CPU计算时间长太多时，CPU计算开销仍可能成为瓶颈。针对这些有限的情况，我们开发了“单步延迟参数更新”（DPU），通过将参数更新延迟一步来重叠CPU和GPU的计算，从而隐藏CPU计算开销。我们在评估中验证了DPU不影响最终的训练准确性。

DPU训练调度
图6展示了使用延迟参数更新的ZeRO-Offload训练流程。
1. ➊ 前N-1步，不使用DPU进行训练，以避免在梯度变化迅速的训练早期阶段破坏稳定性。
2. ➋ 在第N步，我们从GPU获取梯度，但跳过CPU优化器步骤，也不更新GPU上的fp16参数。
3. ➌ 在第N+1步，我们使用第N步的梯度在CPU上计算参数更新，同时在GPU上使用第N-1步更新的参数并行计算前向和后向传播。从这一步开始，第(i+1)步的模型将使用由第(i-1)步的梯度更新的参数进行训练，而不是由第i步更新的参数，从而实现了CPU计算与GPU计算的重叠。

准确性权衡
由于DPU改变了训练的语义，因此有理由询问模型准确性与训练效率之间是否存在权衡。为了回答这个问题，我们在多个训练工作负载上评估了DPU，发现如果在几十次迭代后而不是在开始时引入DPU，DPU不会损害收敛性。我们在第6节的评估结果表明，与仅使用ZeRO-Offload的训练相比，使用延迟参数更新的训练在达到相同模型训练准确性的同时，具有更高的训练吞吐量。

图6：训练过程中的延迟参数更新。
图6：训练过程中的延迟参数更新。

A4 实验环境

表2：实验系统的硬件概览。
表2：实验系统的硬件概览。

表3：评估中的模型配置。
表3：评估中的模型配置。

A4 实验结果

模型规模

训练吞吐量

吞吐量可扩展性

优化的CPU执行

A5 结论

本文介绍了ZeRO-Offload，一种强大的GPU-CPU混合深度学习训练技术。它具备高计算效率和近线性的吞吐量可扩展性，使得数据科学家即便在单块GPU上也能训练数十亿参数的模型，且无需任何模型重构。通过将ZeRO-Offload作为DeepSpeed库的一部分开源，作者希望能够普及大规模模型训练，让全球的数据科学家都能发掘超大规模深度学习模型的潜力。

引用文献汇总