优化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语言模型时优化所带来的整体性能提升。模型GPU数量每GPU TFLOPSRoBERTa-10B64优化后123 teraflops未优化73 teraflopsRoBERTa-50B64优化后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更多精彩内容 请关注我的个人公众号 公众号办公AI智能小助手或者 我的个人博客 https://blog.qife122.com/对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号网络安全技术点滴分享