Nano-Pipe流水线并行

背景与挑战

流水线并行是AI大模型分布式训练的关键组成部分之一，但其效率受到流水线中空泡的影响，为了提高吞吐量，需要降低其空泡比例。Nano-Pipe展示了另一个角度下的优化思路。

解决方案

在大模型流水线调度中，反向的输入梯度和权重梯度通常是一起调度计算的，然而，通过分析他们计算的依赖关系，可以发现其实只有输入梯度的计算存在相互层间的依赖关系。因此，通过独立调度反向的输入梯度和权重梯度的计算，可以减少流水线调度的空泡。

分离权重梯度计算流程，可通过修改RowParallelLinear和ColumnParallelLinear的backward实现，将对权重的梯度计算进行剥离，先存储在调度器上权重的梯度计算队列中。在需要对权重的梯度计算时，从调度器上权重的梯度计算队列中出队列中，依次出队，然后计算对应的梯度。

使用场景

在训练模型时，降低空泡的比例，从而提升计算效率，达到更好的流水线并行。此特性暂只适配Legacy分支，即脚本中配置--use-legacy-models。

使用方法

1. Nano-Pipe依赖于虚拟流水并行，使用前需检查是否开启流水线并行和虚拟流水并行。

   --num-layers-per-virtual-pipeline-stage N   

   N表示每个虚拟流水线阶段的层数。

2. 在脚本中配置以下参数，使能Nano-Pipe流水线并行。

   --use-nanopipe

使用效果

通过流水线并行，Nano-Pipe可以更快地处理大量的数据，提高数据处理的效率和准确性，提升计算的速度，减少空泡占比。