4.1 原理解析

GPipe与PipeDream原理解析¶

学习目标¶

理解GPipe架构.
理解GPipe的底层原理.
理解PipeDream架构.
理解PipeDream底层原理.

微批次流水线并行¶

微批次(MicroBatch)流水线并行技术: 与3.1小节中的朴素流水线技术几乎相同, 但是通过将传入的小批次(minibatch)分块为微批次(microbatch), 并人为创建流水线来解决GPU空闲问题, 从而允许不同的GPU同时参与计算过程, 可以显著提升流水线并行设备利用率, 减小设备空闲状态的时间.

目前工业界常见的流水线并行方法GPipe和PipeDream都采用微批次流水线并行方案! 如下图所示:

GPipe¶

GPipe, 即Google Pipe(全称为Easy Scaling with Micro-Batch Pipeline Parallelism), 是由谷歌提出的一种流水线并行方案. 最早, 谷歌在Lingvo框架下开源了GPipe, 基于Tensorflow实现的. 后来, Kakao Brain的工程师用Pytorch实现了GPipe, 并开源出来, 也就是torchgpipe. 再后来Facebook(Meta)将相关库集成到了Pytorch 1.8.0之后的版本中. 相关代码在torch/distributed/pipe/sync目录下.

一个栗子: 基于Pytorch使用包含两个FC层的模型跨GPU0和GPU1进行流水线并行:

import os
import torch
import torch.nn as nn

# 首先初始化RPC框架
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '29500'
torch.distributed.rpc.init_rpc('worker', rank=0, world_size=1)

# 构建模型
fc1 = nn.Linear(16, 8).cuda(0)
fc2 = nn.Linear(8, 4).cuda(1)
model = nn.Sequential(fc1, fc2)

# 导入流水线并行的模块
from torch.distributed.pipeline.sync import Pipe

# chunks表示micro-batches的大小, 默认值为1
model = Pipe(model, chunks=8)
input = torch.rand(16, 16).cuda(0)
output = model(input)

GPipe流水线并行主要用来解决两个问题:
- 第一: 提高模型训练的并行度.
- 第二: 通过重计算降低显存消耗.

第一: 提高模型训练的并行度. GPipe在朴素流水线并行的基础上, 利用数据并行的思想, 将minibatch细分为多个更小的microbatch, 送入GPU进行训练, 来提高并行程度. 下图展示了朴素流水线并行与GPipe微批次流水线并行对比, 通过GPipe可以有效降低流水线并行bubble空间的比例:

其中, F的第一个下标表示GPU编号, 第二个下标表示microbatch编号. 假设将 minibatch划分为M个, GPipe流水线并行下, GPipe流水线bubble时间为O((K-1)/(K+M-1)). K为设备数量, M为将minibatch切分成多少个microbatch, 当M>>K时, 这个时间可以忽略不计.

缺点分析: 把batch拆分小了之后, 对于那些需要统计量的层(比如Batch Normalization), 就会导致计算变得麻烦, 需要重新实现. 在GPipe中采用的方法是在训练时计算和利用的是microbatch里的均值和方差, 同时持续追踪全部minibatch的移动平均和方差, 以便在测试阶段使用. 这样Layer Normalization不受影响.

第二: 通过重计算(Re-materialization)降低显存消耗. 在模型训练过程中的前向传播时, 会记录每一个算子的计算结果, 用于反向传播时的梯度计算.

Re-materialization可以不保存中间层输出的激活值, 在计算梯度的时候会重新计算出来这些激活值从而可以计算梯度. 在GPipe中, 应用了这个技术后, 如果一个设备上有多层, 那么就可以只保存多层中的最后一层的输出值. 这样就降低了每个设备上内存占用峰值, 同样的模型尺寸需要的显存就少了.

核心: Re-materialization并非是不需要中间结果, 而是有办法在求导过程中实时的计算出之前被舍弃掉的中间结果.

简而言之, GPipe通过纵向对模型进行切分解决了单个设备无法训练大模型的问题. 同时又通过微批量流水线增加了多设备上的并行度. 还使用re-materialization降低了单设备上的显存峰值.

流水线并行策略¶

流水线并行根据执行的策略, 可以分为两种模式:
- F-then-B模式: GPipe, 朴素流水线并行
- 1F1B模式: PipeDream

F-then-B模式¶

F-then-B模式, 先进行前向计算, 再进行反向传播. F-then-B模式由于缓存了多个micro-batch的中间变量和梯度, 显存的实际利用率并不高, 具体如下图所示:

1F1B模式¶

1F1B(One Forward pass followed by One Backward pass)模式, 一种前向计算和反向传播交叉进行的方式. 在1F1B模式下, 前向计算和反向传播交叉进行, 可以及时释放不必要的中间变量.

1F1B流程如下图所示, 以stage4的F42(stage4的第2个micro-batch的前向计算)为例, F42在计算前, F41的反向B41(stage4的第1个micro-batch的反向计算)已经计算结束, 即可释放F41的中间变量, 从而F42可以复用F41中间变量的显存:

研究表明: 1F1B模式相比于F-then-B方式, 峰值显存可以节省37.5%, 对比朴素流水线并行峰值显存明显下降, 设备资源利用率显著提升.

PipeDream¶

GPipe流水线有以下几个问题:
- 第一: 讲mini-batch切分成m份micro-batch后, 将带来更频繁的流水线刷新(Pipeline flush), 这降低了硬件效率, 导致空闲时间的增加. 如下图所示:

将mini-batch切分成m份micro-batch后, 需要缓存m份activation, 这将导致内存增加. 原因是每个micro-batch前向计算的中间结果activation都要被其后向计算所使用, 所以需要在内存中缓存. 即使使用了重计算技术, 前向计算的activation也需要等到对应的后向计算完成之后才能释放.

微软在DeepSpeed中提出了PipeDream, 针对上述问题的改进方法就是1F1B策略.
- 这种个改进策略可以解决缓存activation的份数问题, 使得activation的缓存数量只跟stage数量相关, 从而进一步节省显存, 训练更大的模型.
- 解决思路就是努力减少每个activation的保存时间, 这就需要每个micro-batch数据尽可能早的完成后向计算, 从而让每个activation尽可能早的释放掉.

注意: 微批次在GPipe中叫micro-batch, 在PipeDream中叫mini-batch. 为了避免干扰, 本讲义统一使用micro-batch.

PipeDream具体方案¶

PipeDream具体方案如下:
- 1: 一个阶段stage在做完一次micro-batch的前向传播之后, 立即进行micro-batch的反向传播, 然后释放资源, 这样就可以让其他stage尽可能早的开始计算, 即1F1B策略. 类似于把整体同步变成了众多小数据块上的异步, 而且众多小数据块都是独立更新.
- 2: 在1F1B的稳定状态下(steady state), 会在每台机器上严格交替的进行前向计算/反向传播, 这样使得每个GPU上都会有一个micro-batch数据正在处理, 从而保证资源的高利用率(整个流水线比较均衡, 没有流水线刷新Pipeline flush), 这样能确保以固定周期执行每个阶段的参数更新.
- 3: 面对流水线带来的异步性, 1F1B 使用不同版本的权重来确保训练的有效性.
- 4: PipeDream还扩展了1F1B, 对于使用数据并行的stage, 采用轮询(round-robin)的调度模式将任务分配在同一个stage的各个设备上, 保证了一个小批次的数据的前向计算和反向传播发生在同一台机器上, 即1F1B-RR(one forward one backward-round robin).

结论: 相比GPipe, 表面上看PipeDream在Bubble率上并没有优化, PipeDream流水线Bubble时间仍然是O((K-1)/(K+M-1)). 但节省了显存后, 在设备显存一定的情况下, 就可以通过增大M的值(增大micro-batch的个数)来降低Bubble率了.

PipeDream-2BW¶

在之前的流水线方案GPipe和PipeDream存在如下问题:
- 1: GPipe维护模型权重的单一版本, 输入的小批次被分成更小的微批次. 权重梯度是累积的, 不会立即应用, 流水线会定期刷新, 以确保不需要维护多个权重版本. GPipe提供类似于数据并行的权重更新语义, 但是定期的流水线刷新可能会很昂贵, 从而限制了吞吐量. 减轻这种开销的一种方法是在流水线内执行额外的累积, 但这并不总是实用的.
- 2: PipeDream使用权重存储方案来确保相同输入的前向和后向传播中使用相同的权重版本. 在最坏的情况下, 隐藏的权重版本总数为d, 其中d是流水线深度, 这对于大模型来说成本太高了. 而且使用PipeDream默认的权重更新语义, 每个阶段(state)的权重更新都有不同的延迟项, 同时, 流水线内不会执行累积.

基于上述问题的分析, 有了PipeDream-2BW(double-buffered weights), 其在流水线之中只维护两个版本的模型权重 .

PipeDream-2BW会为每m个微批次生成一个新的权重版本(m>=d), d为流水线深度, 但是因为有些剩余的反向传播计算仍然依赖于旧版本模型, 所以新的模型版本无法立即取代旧版本. 因此, 新生成的权重版本需要缓冲以供将来使用. 然而, 需要维护的权重版本总数最多为2, 因为用于生成新权重版本的权重版本可以立即被丢弃(通过该阶段的后续的输入不再使用旧的权重版本), 由于只保存了两个版本, 这极大的降低了内存的占用.

PipeDream-Flush¶

在PipeDream 2BW论文中(Memory-Efficient Pipeline-Parallel DNN Training), 还提到了一种变体PipeDream-Flush, 使用Flush更新权重. 它的内存占用量低于PipeDream 2BW, 但代价是吞吐量较低. 该调度重用了微软的PipeDream 中的1F1B调度策略. 但是, 同GPipe一样, 只维护单个权重版本并引入定期流水线刷新(pipeline flush), 以确保权重更新期间的权重版本保持一致, 通过这种方式以执行性能为代价降低了峰值内存. 两个流水线阶段的PipeDream-Flush和GPipe的时间线, 如下图所示:

下图展示了GPipe, PipeDream-Flush, PipeDream 2BW流水线并行方法的吞吐量对比:

1F1B调度模式¶

在PipeDream中使用1F1B策略时, 存在两种调度模式:
- 非交错调度
- 交错式调度

下图上部显示了非交错调度(non-interleaved schedule), 下部显示了交错式调度(interleaved schedule):

非交错调度(non-interleaved schedule)分为三个阶段:
- 1: 第一阶段是热身阶段, 处理器进行不同数量的前向计算.
- 2: 第二阶段, 处理器进行一次前向计算, 然后进行一次反向计算.
- 3: 第三阶段, 处理器完成反向计算.

参考上面示例图进行举例说明: 若网络共16层(编号 0-15), 4个Device, 谷歌的GPipe和微软的PipeDream是分成4个stage, 按编号0-3层放置于Device1, 4-7层放置于Device2, 以此类推.

注意: 微软的PipeDream使用非交错式1F1B 调度. 虽然这种调度模式比GPipe更节省内存, 然而它需要和GPipe一样的时间来完成一轮计算!

交错式调度(interleaved schedule), 每个设备可以对多个层的子集(称为模型块)进行计算, 而不是一个连续层的集合. Megatron-LM基于PipeDream-Flush提出了一个优化: 交错式1F1B调度, 是Megatron-LM论文中最主要的一个创新点.

传统的流水线并行通常会在一个设备(Device)上放置几个连续的模型层(比如Transformer层). 但Megatron这篇论文采用虚拟流水线(virtual pipeline)进行交错式1F1B并行. 在设备数量不变的情况下, 分出更多的流水线阶段(pipeline stage), 以更多的通信量, 换取流水线Bubble比率降低.

那虚拟流水线(virtual pipeline)是怎么做到的呢? 该方案要求一个小批次中的微批次数量是管道并行大小(即流水线中的设备数量)的整数倍. 例如, 对于4个设备, 一个小批次中的微批次数量必须是4的倍数.

英伟达的virtual pipeline是按照论文中提出的virtual_pipeline_stage概念减小切分粒度, 以virtual_pipeline_stage=2为例:
- 将0-1层放置于Device1;
- 将2-3层放置于Device2;
- 将4-5层放置于Device3;
- 将6-7层放置于Device4;
- 然后8-9层继续放置于Device1;
- 10-11层继续放置于Device2;
- 12-13层继续放置于Device3;
- 14-15层继续放置于Device4, 如下图所示:

分布式训练框架流水线并行方案小结:
- 同步流水线(Sync-PP): GPipe, PipeDream-Flush
- 异步流水线(Async-PP): PipeDream, PipeDream 2BW

同步流水线: 与数据并行具有相同的权重更新语意, 但是需要引入流水线bubble(空闲等待时间), 会降低训练的吞吐量.

异步流水线: 大幅度降低了训练timeline中的bubble, 但是需要引入不同的权重版本来解决权重过期的问题.

几个著名框架中采用的流水线并行方案:
- Pytorch: 采用的是GPipe方案, 使用F-then-B调度策略.
- DeepSpeed: 采用PipeDream-Flush方案, 使用非交错式1F1B调度策略.
  - 使用这个调度方案, 是为了促进最大规模的模型进行训练, 在模型训练过程中, 存储多个权重缓冲可能会令人望而却步, 我们的首要目标希望是一个"精确"的方法, 而不需要收敛权衡. 当然, DeepSpeed引擎组件抽象出了流水线调度, 你也可以自行实现其他的流水线调度方案.
- Megatron-LM: 采用改进后的PipeDream-Flush方案, 使用交错式1F1B调度策略.
- Colossal-AI: 基于Megatron-LM的交错式1F1B方案, 提供了非交错和交错两种调度策略.