某机构研究人员优化分布式训练工具，使其在弹性结构适配器网络接口上高效运行。
大多数现代自然语言处理应用都构建于预训练语言模型之上，这些模型编码了整个语言的词序列概率。随着时间的推移，这些模型规模越来越大，参数数量达到数十亿甚至数万亿。要在合理的时间内训练这些模型，需要非常庞大的计算集群，其巨大的通信量可能会阻塞计算，导致GPU利用率低下或效率不高。因此，需要仔细管理GPU之间的通信，以避免其成为性能瓶颈。
微软的DeepSpeed分布式训练库引入了一种名为“零冗余优化器”（ZeRO）的管理技术。ZeRO通过将机器学习模型的状态分区到分布式工作节点上，并在训练过程中按需从其他工作节点获取必要的模型状态。ZeRO具有多个“阶段”，每个阶段都通过降低内存需求来支持更大模型的训练，但通常会增加通信量。
虽然微软的研究人员能够利用该技术实现理想的扩展性能，但他们的实验报告仅基于使用昂贵的高速InfiniBand网络（具体为Nvidia DGX系统）的专业超级集群。为降低需要高性能计算的客户的成本，某机构使用弹性结构适配器网络而非InfiniBand。在某机构p4d.24xlarge计算基础设施实例上可用的EFA，其通信带宽低于Nvidia DGX超级集群上的InfiniBand，因此对于带宽密集型任务，我们预期性能会有所下降。然而，在尝试复现微软的结果时，我们发现ZeRO第三阶段的性能相对下降是第二阶段的两倍。
我们通过分析训练过程来寻找瓶颈，并观察到在ZeRO Stage 3中，通信占据了训练的主导时间。为了缩小与配备InfiniBand的DGX集群上获得的结果之间的性能差距，我们对ZeRO Stage 3进行了一系列优化。下表展示了优化带来的整体性能提升，这是在某机构p4d.24xlarge实例上训练RoBERTa语言模型时测量的结果。
| 模型 | GPU数量 | 每GPU TFLOPS |
|---|---|---|
| RoBERTa-10B | 64 | 优化后：123 万亿次浮点运算/秒 优化前：73 万亿次浮点运算/秒 |
| RoBERTa-50B | 64 | 优化后：154 万亿次浮点运算/秒 优化前：89 万亿次浮点运算/秒 |
今年一月，我们将这些优化合并到DeepSpeed代码库中，供公众使用。
优化措施
我们的优化大致可分为三类：（1）改进通信与计算的重叠；（2）提高带宽利用率；（3）提升内存效率。
同步/并行性
通信流与计算流之间更细粒度的同步 在低带宽或大型集群中，通信时间占主导地位，通过计算与通信的重叠来掩盖通信成本至关重要。通过分析，我们发现ZeRO过于粗粒度的同步限制了这种重叠。这导致两种分布式计算操作（allgather和reduce-scatter）的重叠水平不佳。allgather用于聚合网络中所有工作节点的数据（本例中为模型参数），而reduce-scatter则用于跨工作节点归约数据（本例中为梯度求和）。由于通信不断阻塞计算操作，这两个操作导致了较差的GPU利用率。为此，我们对参数收集和梯度归约-分散路径进行了重大修改，以减少或消除同步，同时保持正确性。完成这些更改后，我们实现了更好的重叠，从而大大减少了计算气泡的数量和大小。
Python获取和分区决策的预计算/缓存 在训练过程中，需要做出许多复杂的决策，例如应获取哪些参数、接下来将使用哪些参数、哪些参数可能很快被重用而应保留、哪些可以释放。这些操作速度较慢，常常使Python进程无法持续为GPU提供工作，从而产生大的计算气泡。我们通过预先计算或尽可能多地缓存决策来优化这一点，加速了这些决策的计算，使其不再成为训练吞吐量的影响因素。
通信/带宽利用
批处理allgather/reduce-scatter调用 我们发现，将集合操作（allgather和reduce-scatter）进行批处理可以更有效地利用带宽，并分摊执行这些操作的计算内核的固定成本。为实现集合操作的批处理，我们将张量数据展平到一个连续的缓冲区中，以便在单个事务中发送。每个集合操作都需要一个特殊的交织方案，以确保每个工作节点接收到正确的数据。
内存
我们的ZeRO实现与微软的实现一样，都使用了统一计算设备架构（CUDA）。CUDA内存分配既是同步的也是缓慢的，因此PyTorch使用缓存分配器来避免不断重新分配内存的高昂成本。如果分配请求没有可用的缓存或空闲块，分配器将刷新其缓存。这带来几个问题：
- 在刷新开始之前，需要多次调用
cudaEventSynchronize以允许在持有内存上的计算完成。此操作及随后的cudaFree调用可能需要数秒。 - 无法保证不同工作节点同时刷新其缓存。这意味着对于任何集合操作，即使只有一个工作节点正在刷新其缓存，其他
N-1个工作节点也会被阻塞，等待该工作节点加入。随着集群规模增大，任何给定集合操作中至少有一个工作节点正在刷新其缓存的可能性也随之增加。 - 缓存刷新后，后续分配需要
cudaMalloc调用，如前所述，这些调用既同步又缓慢。
因此，内存效率对于性能至关重要。
内存高效的批处理PyTorch集合操作
虽然我们使用批处理集合操作显著减少了内核启动开销并提高了带宽利用率，但它也因将批处理的张量展平到一个额外的缓冲区而增加了内存消耗。为避免在PyTorch集合操作中进行冗余的展平操作，我们使用了集合操作的*_base变体，这些变体接受预先展平的张量，从而避免了内部额外分配展平缓冲区的需要。在未来的工作中，我们计划使用Nvidia集合通信库的基于组的批处理操作来消除所有展平操作。
初始化时更激进的参数分区碎片整理 即使在有超过10GB的GPU空闲内存时，我们仍然观察到分配器缓存刷新的迹象，这表明存在内存碎片。为了减少这种情况，我们进行了初始化时的碎片整理更改，将所有持久化张量移动到一个连续的缓冲区中。
其他
除上述优化外，我们还：
- 通过减少主机与设备间的数据传输和同步，并将数学运算从for循环中取出，放入一个具有并行化计算的单一内核启动中，从而优化了梯度归一化。
- 移除了通过字符串格式化添加到调试消息中的张量操作（
.norm()）。这些操作会导致从主机到设备的复制，意味着数据传输和主机-设备同步。
通过使DeepSpeed ZeRO Stage 3在广泛可用的公共云服务上实现高性能，我们希望进一步推动大型语言模型训练的普及。FINISHED
