权重更新通信隐藏

问题分析

大模型训练过程中，通常会使用数据并行。在进行梯度更新时，数据并行组中的通信要等反向计算完成后再进行。这样的串行执行顺序会造成计算和通信流存在一定的空闲等待时间，导致执行效率较低。

解决方案

通过计算和通信任务并行的方式来实现相互的流水掩盖。

a. 仅打开 --use-distributed-optimizer

仅打开分布式优化器时（--use-distributed-optimizer），运行流程如下图所示，前向和反向计算完成后，会有独立的通信时间，进行梯度的reduce-scatter、计算权重、进行权重的all-gather，获得权重之后再进入下一轮的前向计算。

b. 打开 --use-distributed-optimizer 和 --overlap-grad-reduce

在打开--use-distributed-optimizer的同时打开--overlap-grad-reduce，运行流程如下图所示，对梯度的reduce-scatter过程与反向计算过程并行，从而节省了单独的reduce-scatter过程，提高了计算-通信并行效率。

c. 打开 --use-distributed-optimizer 和 --overlap-grad-reduce 和 --overlap-param-gather

在打开--use-distributed-optimizer和--overlap-grad-reduce的基础上进一步打开--overlap-param-gather，运行流程如下图所示，对权重的all-gather过程与下一轮的前向计算并行，从而节省了单独的all-gather过程。

以上流程对比发现，打开--overlap-param-gather后，通信与计算完全并行，极大提高了计算-通信并行效率，进而提升了模型训练效率。

使用场景

在数据并行场景可以开启该特性。

使用方法

设置--overlap-param-gather即可调用该算法。 确保同时开启了--use-distributed-optimizer和--overlap-grad-reduce。

使用效果

使用该特性可以提升性能。

注意事项

开启该特性后，attention层init的顺序会更正为先创建linear_qkv再创建linear_proj，这是为了修复Megatron的错误init顺序，该bug会导致当linear_qkv和linear_proj被分配在不同bucket时，overlap-param-gather可能会在权重未完成更新时进行下一轮前向计算。 legacy下，--overlap-param-gather暂不支持和reuse_fp32_param一起使用。