大模型分布式并行技术–数据并行优化

作者：hanscalZheng 2023-11-01 20:10:53

云计算

分布式通信和计算的重叠通常是将通信和计算算子调度到不同的流 (stream) 上实现的。通信算子调度到通信流，计算算子调度到计算流，同一个流上的算子间是顺序执行的，不同流上的算子可以并行执行，从而实现反向中梯度通信和计算的并行重叠。

通信融合

从上文知道数据并行中需要同步每一个模型梯度，这是通过进程间的 Allreduce 通信实现的。如果一个模型有非常多的参数，则数据并行训练的每一个 step 中会有非常多次的 Allreduce 通信，下图为融合梯度同步示例。

融合梯度同步示例

通信的耗时可以从通信延迟(lantency) 和数据传输时间消耗两方面考虑。单次通信延迟时间相对固定，而传输时间由通信的数据量和带宽决定。减少总的通信消耗，可以通过减少通信频率来实现，通信融合是一个可行的手段，通过将 N 个梯度的 Allreduce 通信合并成一次 Allreduce 通信，可以减少 N- 1 次通信延迟时间。

常用的 Allreduce 融合实现方式是在通信前将多个梯度 tensors 拼接成一个内存地址连续的大 tensor，梯度同步时仅对拼接后的大 tensor 做一次 Allreduce 操作。参数更新时将大 tensor 切分还原回之前的多个小 tensors，完成每个梯度对应参数的更新。

通信计算重叠

除了降低绝对的通信耗时，还可以从降低整体训练耗时角度来优化，可以考虑通信和计算的异步流水实现。数据并行中的梯度同步 Allreduce 通信是在训练的反向过程中进行的，而 Allreduce 后得到的同步梯度是在训练的更新过程中才被使用，在反向中并没有被使用。也就是说上一个梯度的通信和下一个梯度的计算间并没有依赖，通信和计算可以并行，让两者的耗时相互重叠掩盖，减少反向的耗时，下图为通信计算并行相互重叠示例。

通信计算并行相互重叠示例。

通信和计算的重叠通常是将通信和计算算子调度到不同的流 (stream) 上实现的。通信算子调度到通信流，计算算子调度到计算流，同一个流上的算子间是顺序执行的，不同流上的算子可以并行执行，从而实现反向中梯度通信和计算的并行重叠。需要注意的是，当通信和计算被调度在不同的流上执行时，需要考虑两个流之间依赖和同步关系。