现代自然语言处理应用大多基于预训练语言模型，这些模型参数量已增长至数十亿甚至数万亿级别。在合理时间内训练这些模型需要大型计算集群，但高通信量会导致GPU利用率低下。微软DeepSpeed库提出的零冗余优化器(ZeRO)技术通过分区模型状态解决了这一问题，但其理想性能仅在配备高速InfiniBand的专用超算集群上实现。
某中心研究团队发现，在使用弹性结构适配器(EFA)网络的p4d.24xlarge实例上，ZeRO第三阶段性能下降幅度是第二阶段的两倍。通过分析发现，通信时间在第三阶段占据了训练过程的主导地位。团队实施了三类核心优化：
- 同步/并行化优化
- 改进通信与计算流之间的细粒度同步
- 重构参数收集和梯度reduce-scatter路径
- 预计算Python获取与分区决策
- 通信/带宽优化
- 批量处理allgather/reduce-scatter调用
- 采用特殊数据交错方案确保传输准确性
- 内存效率提升
- 使用预扁平化张量的集体操作变体
- 初始化时参数分区碎片整理
- 构建持久化张量的连续缓冲区
优化效果在RoBERTa模型训练中得到验证：在64 GPU上，10B参数模型单GPU算力从73提升至123万亿次，50B模型从89提升至154万亿次。这些改进已并入DeepSpeed代码库，显著降低了大规模语言模型的训练硬件门槛。
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）
公众号二维码
