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

作者：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) 上实现的。通信算子调度到通信流， 计 算算子调度到计算流， 同一个流上的算子间是顺序执行的， 不同流上的算子可以并行执行， 从而实现反向中梯 度通信和计算的并行重叠。需要注意的是， 当通信和计算被调度在不同的流上执行时， 需要考虑两个流之间依 赖和同步关系。

• 某个梯度 Allreduce 通信进行前，该梯度的反向计算已经完成。
• 某个梯度对应参数的更新计算开始前，该梯度的 Allreduce 通信已经完成。

在梯度同步的数据并行场景中，开发者需要需要通过 stream 间的同步功能保证：

以上两个方法是数据并行中常用的减少通信时间消耗， 提高并行加速比的优化策略。如果能做到通信和计 算的重叠程度越高，那么数据并行的加速比越接近 100% ，多卡并行对训练吞吐提升的效率也就越高。