在大模型训练场景下,传统Megatron数据并行策略往往采用串行执行模式,即在各数据并行组内完成反向传播计算后,再统一执行梯度通信(未启用分布式优化器时为AllReduce,启用时为ReduceScatter)。这种顺序执行方式导致计算与通信之间的等待时间,降低了资源利用率与整体训练效率。
Bucket机制:设立一个临时存储区(Bucket),用于暂存反向传播产生的梯度。一旦Bucket达到预设容量,即刻触发内部梯度的通信任务,无需等待所有反向计算完成。此机制允许后续反向计算与当前通信任务并行运行,显著提升了计算与通信资源的利用率。
需在训练脚本中加入以下参数,即可激活异步DDP优化算法。
--use-distributed-optimizer --overlap-grad-reduce
异步DDP通过实现计算与通信任务的重叠执行,显著改善了大模型训练中的资源利用率和训练效率,对于LLaMA2-70B模型,端到端性能提升约2-3%。