ZeRO-Offload 的工作原理
ZeRO-Offload 是基于 Zero Redundancy Optimizer (ZeRO) 构建的。ZeRO 是微软在 2020 年 2 月提出的一种万亿级模型参数训练方法，用于数据并行和模型并行训练中的内存优化，其中梯度、参数和优化器状态分布在多 GPU 内存中，没有任何冗余。这使得 GPU 之间的通信开销保持在比较低的状态。
与标准数据并行基准相比，ZeRO 在三个阶段中节省的内存和通信用量。
让我们来回顾一下 ZeRO：
为了解决数据并行和模型并行存在的问题，ZeRO 提供了三阶段的优化方法，分别为优化器状态分割、梯度分割、参数分割，三个阶段按顺序实施。
在优化器分割状态：ZeRO 降低了 3/4 的内存，通信量和数据并行相同；
加入梯度分割：降低了 7/8 的内存，通信量和数据并行相同；
加入参数分割：内存减少与数据并行度呈线性关系。例如，在 64 个 GPU 上进行分割的时候，可以将内存降至 1/64。在通信量上有 50% 的提升。
在去年 9 月份的博客中，微软这么介绍 ZeRO-Offload：
ZeRO-Offload 继承了 ZeRO-2 的优化器状态和梯度分割。但与 ZeRO-2 不同的是，ZeRO-Offload 不在每块 GPU 上保持优化器状态和梯度的分割，而是将二者卸载至主机 CPU 内存。在整个训练阶段，优化器状态都保存在 CPU 内存中；而梯度则在反向传播过程中在 GPU 上利用 reduce-scatter 进行计算和求均值，然后每个数据并行线程将属于其分割的梯度平均值卸载到 CPU 内存中（参见下图 g offload），将其余的抛弃。一旦梯度到达 CPU，则每个数据并行线程直接在 CPU 上并行更新优化器状态分割（参见下图 p update）。
之后，将参数分割移回 GPU，再在 GPU 上执行 all-gather 操作，收集所有更新后的参数（参见下图 g swap）。ZeRO-Offload 还利用单独的 CUDA 流来穷尽通信与计算中的重叠，从而最大化训练效率。
ZeRO-Offload 概览。
值得注意的是，ZeRO-Offload 专为使用 Adam 的混合精度训练而设计。也就是说，当前版本的 ZeRO-Offload 使用 Adam 的优化版本 DeepCPUAdam。其主要原因是避免 CPU 计算成为整个过程中的瓶颈。DeepCPUAdam 的速度是 Adam PyTorch 实现的 6 倍。会昌网景
