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

某机构的DeepSpeed分布式训练库引入了一种称为零冗余优化器（ZeRO）的管理技术。ZeRO通过将机器学习模型状态分区存储在分布式工作节点上，并在训练期间按需从其他节点获取所需模型状态。ZeRO具有多个"阶段"，每个阶段都通过降低内存需求来支持训练更大模型，但通常会增加通信量。

虽然某机构研究人员使用专用超算集群（配备昂贵的高速InfiniBand网络）实现了理想的扩展性能，但某中心Web服务（AWS）使用Elastic Fabric Adapter（EFA）网络来降低高性能计算成本。我们发现，在EFA网络上ZeRO第三阶段的性能下降是第二阶段的两倍。

通过分析训练过程，我们发现ZeRO Stage 3中通信时间占主导地位。为此我们进行了一系列优化：

1. 同步/并行优化：

   • 改进通信与计算流之间的细粒度同步
   • 预计算和缓存Python获取与分区决策

2. 通信/带宽利用：

   • 批量处理allgather和reduce-scatter调用
   • 实现特殊的数据交错方案确保正确传输

3. 内存效率提升：

   • 采用内存高效的批处理PyTorch集合操作
   • 初始化时更积极的参数分区碎片整理

优化后，在AWS p4d.24xlarge实例上训练RoBERTa语言模型的性能显著提升：

这些优化已合并到DeepSpeed代码库中供公开使用。通过使DeepSpeed ZeRO Stage 3在公共云服务上高效运行，我们希望能进一步推动大型语言模型训练的普及化。

其他改进包括：

• 优化梯度归一化，减少主机-设备数据传输
• 移除调试消息中的张量操作以避免不必要的数据拷贝

这些优化显著提升了在低成本硬件上的分布式训练效率，为更广泛的研究社区提供了训练大型模型的可能性。