优化DeepSpeed ZeRO，让大模型训练更经济

现代自然语言处理应用大多建立在预训练语言模型之上，这些模型编码了整门语言的词序列概率。随着时间的推移，这些模型变得越来越大，达到数十亿甚至数万亿参数的规模。在合理的时间内训练这些模型需要非常庞大的计算集群，而集群间巨大的通信量可能会阻塞计算，导致GPU利用率低下。因此，需要精心管理GPU之间的通信，以避免其成为性能瓶颈。

微软的DeepSpeed分布式训练库引入了一种名为“零冗余优化器”（ZeRO）的管理技术。ZeRO的工作原理是将机器学习模型的状态分区到分布式工作节点上，并在训练过程中根据需要从其他工作节点获取必要的模型状态。ZeRO有多个“阶段”，每个阶段都通过降低内存需求来训练更大的模型，但这通常是以增加通信量为代价的。

虽然微软研究人员使用这项技术实现了理想的扩展性能，但他们仅在专用的超大规模集群上进行了实验报告，该集群使用了昂贵的高速InfiniBand网络（具体来说，是某图形处理器厂商的DGX系统）。为了降低需要高性能计算的客户成本，某机构使用弹性结构适配器（EFA）网络替代InfiniBand。其p4d.24xlarge计算基础设施实例上可用的EFA，其通信带宽低于某DGX超大规模集群上的InfiniBand，因此预计带宽密集型任务的性能会有所下降。然而，当尝试重现微软的结果时，发现ZeRO第三阶段的性能相对下降幅度是第二阶段的两倍。

对训练过程进行了分析以寻找瓶颈，并观察到在ZeRO第三阶段中，通信占据了训练的主导地位。为了缩小与配备InfiniBand的DGX集群上所获结果的性能差距，对ZeRO第三阶段进行了一系列优化。下表展示了在某机构p4d.24xlarge实例上训练RoBERTa语言模型时，优化所带来的整体性能提升。

未优化：73 teraflops |

| RoBERTa-50B | 64 | 优化后：154 teraflops

未优化：89 teraflops |

今年一月，将这些优化合并到了DeepSpeed代码仓库中，供公众使用。

优化

优化大致可分为三类：（1）改善通信与计算的重叠；（2）优化带宽利用率；（3）提升内存效率。

同步/并行度

通信与计算流的细粒度同步

在带宽较低或规模较大的集群中，通信时间占主导地位，因此通过计算与计算重叠来掩盖通信成本至关重要。通过分析，发现ZeRO过于粗粒度的同步限制了这种重叠。

这导致了两种分布式计算操作的重叠水平不理想：allgather（从网络中的所有工作节点聚合数据，此处为模型参数）和reduce-scatter（跨工作节点归约数据，此处为梯度求和）。由于通信持续阻塞计算操作，这两个操作导致了GPU利用率低下。为此，对参数收集和梯度reduce-scatter路径进行了重大更改，以减少或消除同步，同时保证正确性。

经过这些更改后，能够实现更好的重叠，从而显著减少计算空闲时间。

Python获取和分区决策的预计算/缓存

训练期间需要做出许多复杂决策，涉及应获取哪些参数、接下来将使用哪些参数、哪些参数可能很快被重用而应保留，以及哪些可以释放。这些操作速度慢，经常导致Python进程无法持续为GPU提供任务，从而产生大量计算空闲时间。

通过预计算或缓存尽可能多的决策来优化这一点，使其计算速度大幅提升，以至于不再影响训练吞吐量。

通信/带宽使用

批量化allgather/reduce-scatter调用

发现批量化集合通信操作——allgather和reduce-scatter——可以更有效地利用带宽，并分摊运行执行这些操作的计算内核的固定成本。为了实现批量化集合通信，将张量数据扁平化为单个连续的缓冲区，以便在一次事务中发送。每个集合通信都需要一个特殊的交错方案，以确保每个工作节点接收到正确的数据。

内存

实现ZeRO的方式，与微软的实现一样，都使用了某图形处理器厂商的并行计算平台。其内存分配既是同步的又是缓慢的，因此PyTorch使用缓存分配器来避免不断重新分配内存的巨大成本。如果某个分配请求没有缓存或空闲块，分配器将刷新其缓存。这会造成灾难性后果，原因如下：

• 在刷新开始之前，需要进行多次同步调用，以允许完成对持有内存的计算。这及其后的空闲调用可能需要数秒时间。
• 不同的工作节点不能保证同时刷新其缓存。这意味着对于任何集合通信，只要有一个工作节点正在刷新其缓存，其他N-1个工作节点就会阻塞等待该工作节点加入。随着集群规模增大，对于任何给定的集合通信，至少有一个工作节点正在刷新其缓存的概率也会增加。
• 缓存刷新后，后续分配需要调用malloc函数，如前所述，该函数既同步又缓慢。

出于这些原因，内存效率对性能至关重要。

内存高效的批量化PyTorch集合通信

尽管使用批量化集合通信显著减少了内核启动开销并提高了带宽利用率，但由于将批量化张量扁平化为一个额外的缓冲区，它也增加了内存消耗。

为了避免在PyTorch集合通信中进行冗余的扁平化操作，使用了集合通信操作的*_base变体，这些变体接受预先扁平化的张量，从而避免了在内部分配额外的扁平化缓冲区。在未来工作中，计划使用某厂商通信库中基于组的批量化操作，以消除所有扁平化操作。

初始化时更积极地对参数分区进行碎片整理

即使有超过10GB的可用GPU内存，仍然观察到分配器缓存刷新的迹象，这暗示了内存碎片问题。为了减少这种情况，对初始化时的碎片整理进行了更改，将所有持久化张量移动到一个连续的缓冲区中。

其他优化

除了上述优化之外，还：

• 通过减少主机与设备间的数据移动和同步，并将数学运算从for循环中移出，整合为一次内核启动与并行计算，从而优化了梯度归一化。
• 移除了通过字符串格式化添加到调试消息中的张量操作，因为这些操作会导致从主机到设备的复制，意味着数据移动和主机-设备同步。

通过使DeepSpeed ZeRO第三阶段在广泛可用的公共云产品上高效运行，希望进一步推动大语言模型训练的普及。

致谢：Zhen Zhang, Stephen Rawls, Yida WangFINISHED