优化 DeepSpeed ZeRO 在低成本硬件上的运行效率

现代自然语言处理应用大多基于预训练语言模型，这些模型参数量已增长至数十亿甚至数万亿级别。在合理时间内训练这些模型需要大型计算集群，但高通信量会导致 GPU 利用率低下。微软 DeepSpeed 库提出的零冗余优化器(ZeRO)技术通过分区模型状态解决了这一问题，但其理想性能仅在配备高速 InfiniBand 的专用超算集群上实现。

某中心研究团队发现，在使用弹性结构适配器(EFA)网络的 p4d.24xlarge 实例上，ZeRO 第三阶段性能下降幅度是第二阶段的两倍。通过分析发现，通信时间在第三阶段占据了训练过程的主导地位。团队实施了三类核心优化：

1. 同步/并行化优化

• 改进通信与计算流之间的细粒度同步

• 重构参数收集和梯度 reduce-scatter 路径

• 预计算 Python 获取与分区决策

2. 通信/带宽优化

• 批量处理 allgather/reduce-scatter 调用

• 采用特殊数据交错方案确保传输准确性

3. 内存效率提升

• 使用预扁平化张量的集体操作变体

• 初始化时参数分区碎片整理

• 构建持久化张量的连续缓冲区

优化效果在 RoBERTa 模型训练中得到验证：在 64 GPU 上，10B 参数模型单 GPU 算力从 73 提升至 123 万亿次，50B 模型从 89 提升至 154 万亿次。这些改进已并入 DeepSpeed 代码库，显著降低了大规模语言模型的训练硬件门槛。

更多精彩内容 请关注我的个人公众号 公众号（办公 AI 智能小助手）公众号二维码

办公AI智能小助手