nanopipe流水线并行
问题分析
流水线并行是AI大模型大规模分布式训练的关键组成部分之一,但其效率收到流水线中bubble的影响,为了提高吞吐量,需要降低其bubble比例。
解决方案
在大模型流水线调度中,反向的input梯度和weight梯度通常是一起调度计算的,然而,通过分析他们计算的依赖关系,可以发现其实只有input梯度的计算存在相互层间的依赖关系。因此,通过独立调度反向的input梯度和weight梯度的计算,我们可以减少流水线调度的bubble。
反向input梯度和weight梯度一起调度的Interleaved 1F1B如下图所示:
独立调度input梯度和weight梯度的nano-pipe如下图所示:

独立调度weight计算展示图如下图所示:

解决思路:
- 分离weight梯度计算流程,通过修改RowParallelLinear和ColumnParallelLinear的backward实现,将对weight的梯度计算进行剥离,先存储在调度器的dw计算队列中。
- 在需要对dw计算时,从调度器的dw计算队列中pop出一个计算,然后计算对应的梯度。
使用场景
在训练模型时,降低bubble的比例,从而提升计算效率,达到更好的流水线并行。此特性暂未适配--use-mcore-models。
使用方法
nanopipe依赖于vpp,设置--num-layers-per-virtual-pipeline-stage N。要求--pipeline-model-parallel-size > 2
设置--use-nanopipe,默认为False,根据用户需求配置。
使用效果
提升计算效率,减少bubble占比。如下表所示:
| device | TP | SP | PP | SEQ | hidden-size | Nano vs vpp收益 |
|---|---|---|---|---|---|---|
| 单机 | 1 | 关 | 4 | 4096 | 4096 | 3.24% |
| 双机 | 4 | 开 | 4 | 8192 | 8192 | 1.02% |
nanoswap
问题分析
使用nano时grad从前向到反向需要持续存储在npu上,生命周期过长,多次累加会增大npu内存的峰值。
解决方案
将过多的张量做offload动作存储到cpu上,在内存峰值过后再将其张量reload回npu上。
解决思路
在前向时将上一轮过多的张量offload到cpu,再在连续的反向运算中途reload回npu上,通过swap流控制不会让reload和offload出现顺序错误。
完整nanopipe-swap原理图如下图所示:

使用方法
基于nanopipe的基础上再开启--use-nanopipe-swap。
使用效果
优化设备内存峰值,如下表所示:
| device | TP | SP | PP | SEQ | hidden-size | mc2 | Nano内存峰值 | swap内存峰值 | Nano vs swap内存峰值下降 |
|---|---|---|---|---|---|---|---|---|---|
| 单机 | 2 | 开 | 4 | 1024 | 4096 | 开 | 5520.62 | 5177.72 | 6.21% |