Microsoft 开源 Zero Redundancy Optimizer (ZeRO-2)

Microsoft 开源了 Zero Redundancy Optimizer 第二版（ZeRO-2），这是一种分布式深度学习优化算法，能够随集群规模实现超线性扩展。使用 ZeRO-2，Microsoft 成功训练了一个包含 1000 亿参数的自然语言处理（NLP）模型，速度比之前的分布式学习技术快 10 倍。

ZeRO-2 的主要特点

ZeRO-2 是 Microsoft 开源深度学习训练优化库 DeepSpeed 的一部分。该算法通过优化训练过程中的内存消耗，支持训练高达 1700 亿参数的模型。此外，ZeRO-2 减少了分布式集群中工作节点之间的通信，实现了超线性并行加速，将训练时间缩短了最多 10 倍。在一个包含 1024 个 GPU 的集群上，DeepSpeed 团队使用 ZeRO-2 在 44 分钟内训练了一个 BERT 自然语言模型，比 NVIDIA 的结果提高了 30% 以上。

NLP 研究的趋势

近年来，NLP 研究的趋势表明，通过在更大数据集上训练更大的模型可以提高准确性。OpenAI 提出了一套“缩放定律”，表明模型准确性与模型大小呈幂律关系，并通过创建包含 1750 亿参数的 GPT-3 模型验证了这一观点。由于这些模型太大，无法适应单个 GPU 的内存，训练它们需要一组机器和跨集群分配参数的模型并行训练技术。

ZeRO-2 的内存优化策略

ZeRO-2 通过三种策略减少训练所需的内存：减少模型状态内存需求、将层激活卸载到 CPU 以及减少内存碎片。ZeRO-2 通过跨并行进程分区梯度和参数，可以将模型状态内存需求减少多达 8 倍。层激活值从正向训练传递中保存，以便在反向传递中使用，但 ZeRO-2 暂时将它们从 GPU 的内存移动到主机 CPU 的内存。最后，即使内存可用，如果可用内存不连续，内存分配也可能失败。ZeRO-2 通过预分配连续内存块来减少碎片，用于临时用途，如激活和梯度。

实验结果

这些内存优化可以减少所需的模型并行性，从而减少节点间通信开销，当与数据并行训练一起使用时，提供了超线性加速。DeepSpeed 团队发现，增加训练中使用的 GPU 数量可以提高每个 GPU 的总体吞吐量。在大型 NLP 模型的实验中，团队观察到每个 V100 GPU 的平均吞吐量为 37 teraflops，“适用于 20 亿到 130 亿参数的模型大小”。对于给定的模型大小和集群，团队指出，ZeRO-2 训练模型的速度比基线 Megatron 方法快 10 倍。使用 1024 个 V100 GPU，团队在 44 分钟内训练了一个 BERT 模型，打破了 NVIDIA 之前使用 1472 个 V100 GPU 的 47 分钟记录。

社区讨论

DeepSpeed 团队成员 Jeff Rasley 在 Hacker News 上参与了讨论，回答了社区的问题。Rasley 指出，DeepSpeed 在 Microsoft 内部有“数百名用户”，他们使用它来训练生产中的模型。当被问及对 TPU 的支持时，Rasley 回答说：

ZeRO 技术与 TPU 或集群设置中的任何加速器兼容，但我们尚未在 TPU 上测试过。可能需要对 DeepSpeed 进行一些小的重构才能使其与 TPU 一起工作。我们目前没有内部计划支持它们，但当然完全欢迎社区的贡献。

获取 DeepSpeed

DeepSpeed 库可在 GitHub 上获取。