异步DDP

问题分析

在大规模模型训练场景下,传统数据并行策略往往采用串行执行模式,即在各数据并行组内完成反向传播计算后,再统一执行梯度通信(未启用分布式优化器时为AllReduce,启用时为ReduceScatter)。这种顺序执行方式不可避免地导致计算与通信之间的等待时间,降低了资源利用率与整体训练效率。

解决方案

为突破上述瓶颈,引入了异步DDP(Distributed Data Parallel)优化机制,通过将计算与通信任务分解为更细粒度的子任务,实现两者的流水线式重叠执行,有效提升资源利用率。其核心原理如下图所示:
图1 异步DPP并行计算示意图

解决思路

Bucket机制:设立一个临时存储区(Bucket),用于暂存反向传播产生的梯度。一旦Bucket达到预设容量,即刻触发内部梯度的通信任务,无需等待所有反向计算完成。此机制允许后续反向计算与当前通信任务并行运行,显著提升了计算与通信资源的利用率。

使用场景

使用该特性的前提是模型开启数据并行和虚拟流水并行,即脚本中进行如下设置:

--num-layers-per-virtual-pipeline-stage N     # N表示每个虚拟流水线阶段的层数,要求模型总层数L % N == 0,且--pipeline-model-parallel-size > 2

使用方法

只需在训练脚本中加入以下参数,即可激活异步DDP优化算法:

--overlap-grad-reduce

使用效果

异步DDP优化通过巧妙结合计算与通信,有效消除了传统数据并行策略中的等待时间,显著提升了大规模模型训练的执行效率与资源利用率。