Press "Enter" to skip to content

让Pipeline在Transformer LM上沿着Token level并行起来——TeraPipe

这次我写的浅显易懂一些,论文传送门:

 

前言

 

如今,大模型不断刷榜,围绕着大模型并行训练的话题也越来越热,除了最基本的数据并行策略,还有随着Megatron-LM,Mesh TensorFlow项目火起来的Operator partition并行策略,以及GPipe,PipeDream,DAPPLE所提出的Pipeline并行策略。我们也都知道,这些策略在理论上相互正交,可以组合起来形成1+1大于2的效果,而且还具有泛化性和通用性——适合于不同种类的模型,这正是系统性能研究的价值所在。

 

系统追求通用性是职责所在,但过分强调这个原则可能会失掉一些机会。一个基本事实是,大模型基本离不开Transformer结构,不变的结构与重复的层层堆叠已是主流。以NLP的语言模型为例,Transformer-LM已经是一个大家族。保持其结构不变,仅仅在参数量上做调整就可以认为是不同的模型。何况,这样一种模型的downstream task又十分丰富。显而易见,如果在通用性方面做一点牺牲,仅是挖掘Transformer-LM类模型的特殊并行模式,是具有一定产品价值的。

 

这篇文章注意到了这一点,对Transformer-based LM提出了新的流水线并行方式——基于token level的Pipeline并行模式。在这种Pipeline并行模式下,不再将batch size这一维度做切分形成多个micro-batch,而是在sequence length这一维度上做切分,实现这一维度上的流水并行。

 

Token level的Pipeline并行

 

Batch切分Pipeline并行的缺陷

 

回忆一下,同步Pipeline并行策略会引入Bubble,而一个iteration内Bubble的占比决定了计算效率的理论上限。为了尽可能降低Bubble的占比,通常需要推入更多的micro-batch,但这也会带来一些问题:

 

 

    1. 更多micro-batch可能意味着更大的batch size,这对收敛可能不利

 

    1. 当显存成为制约瓶颈,比如sequence length必须设置很长,导致micro-batch size被迫很小,GPU利用率会很低

 

 

短句子大batch与长句子小batch的Bubble占比

事实上,模型越大,需要处理的sequence length就越大。 所以,本文提出了一种新的解决办法,即切分sequence length维度,即使 长时序的单个sample 也可以做流水并行。使用这种切分方法,可以使Pipeline的Bubble占比更小,计算效率更高,如下图所示。

 

TeraPipe的Token level切分

但这个切分并不是简单的将batch dimension换到sequence length dimension就行的,还有其他与切batch dimension迥然不同的地方。

 

为什幺Token Level切分可行?

 

切分sequence length之所以能够形成Pipeline,是因为Transformer-based LM的特殊结构。Transformer-based模型一般是Self-Attention和Feed Forward Network。其中,观察某一层中Self-Attention在某个时刻 的token计算,它只依赖于 的结果。而Feed Forward Network只和 时刻有依赖。可以回顾下这两个Layer的公式:

 

公式说明时刻t的状态只取决于t之前的状态

在这样的特性下,如果我们让每个Transformer Block(即Self-Attention和Feed Forward Network)放在不同的Device上,就可以实现这样一种并行状态:时刻 在layer 与时刻 在layer 之间的overlap计算。下图(d)就是这一并行的图示,TeraPipe采用的就是这样的并行模式。读者可以对比下(a)(b)(c)三种并行模式,加深理解。

 

所以,我们可以切 形成token level的pipeline并行,让不同时刻的token在不同的layer上同时计算,并行起来。再次强调,Token level的并行是针对Transformer-based LM特殊结构而设计的,并不是一个适合所有模型的技术。

 

Token Level切分的挑战

 

均匀切分还是均匀计算?

 

至此,你可能觉得思路如此简单,直接切sequence length维度不就行了?当然不是。请先思考下面问题。

 

是否应该在sequence length维度上做均匀切分?

 

如果沿用对Batch size切分成多个micro-batch的惯性思维回答这个问题,你会不假思索认为:应该对sequence length维度做均匀切分。然而这是错的。

 

如果在Batch维度上切分成均匀的若干个micro-batch,那幺每个micro-batch的计算量是完全相同的。 但均匀切分sequence length维度会带来不均匀的计算量。 这是因为时刻 的计算和 有关,越到后面的时刻,计算量就越大。显然,我们的目的是让每个切片的计算量均匀,而不是sequence length切分形式上均匀。

 

均匀切分sequence length会引起不均匀的计算量,从而让Pipeline的Bubble占比增多

因此,这里有一个不得不通过 搜索 才能解决的问题,那就是寻找最佳切点。这篇文章作者Zhuohan Li也是Alpa论文的作者,所以在Pipeline方面的建模思路上是非常相似的,有兴趣的读者可以看我另一篇关于Alpa文章的解析。

 

寻找最佳切点——DP Formulation

 

作者使用 动态规划 的方法进行Pipeline latency的建模,为此先做一些Annotation。

已经按layer切分好的Transformer-based model(称为cell),记为:

input squence的各段切分长度,记为:

每个切分的sub squence slice,记为: ,并且

每个slice 在cell 上的耗时,记为: ,可见这个变量和cell没有关系,因为它可以建模成和previous sequence的关系

显然, 。现在列出Pipeline latency的表达式:

 

形式与Alpa相同,推导可以看Alpa那篇文章

这里的 需要建模。某个sequence 无论在哪个cell,它的耗时应该都是相同的。因为它只和当前 的长度 和所有前驱的长度只和 相关。如果记 为forward time,那幺可以建模:

 

现在你可能发现了,同Alpa的建模一样,此处只考虑了计算,并没有把通信建模进来。事实证明不引入通信也可以搜索到很好的切分解。

 

问题来了,我们需要换一种formulation来求解问题。同Alpa文章的思路:要求解上面的问题,也需要枚举不同的 值,然后用DP求解公式(5),其实就是等价于求解第一项的最小值。改写如下:

 

的意义可以表达为:在固定 情况下, 可以表达的总长度限定为 的各种切分方案中,能使total latency最小的方案的time cost值。因为sequence length的切割是线性的,所以很容易套用最优子结构。

 

自此,DP建模完成。可以看出,无论是formulation还是DP的迭代求解思路,与Alpa同模板。分析下复杂度, 对于固定 ,求解的复杂度为 。那幺 有多少种枚举的可能性呢?答案也是 ,因为 的 和 都有 种选择。所以总体时间复杂度是 。

 

重要的工程问题

 

 

    1. 剪枝 :同Alpa,从小打到枚举 ,当 大于best 时则停止枚举

 

    1. Cost估计 :

 

 

对 的估计分成两个部分:

 

第一项比较好估,把所有的情况都跑一遍就行了。对第二项比较麻烦,因为还要有extra context input,所以建立一个简单的模型去拟合,这样就可以仅用一些subset即可拟合出来:

 

与其他策略的融合

 

显然,对于Operator partitioning,data parallel这些策略的结合是显而易见的,因为他们在策略上是正交的,互相可以叠加。而对和micro-batch based pipeline parallelism的结合可能比较迷,但其实他们也是正交的。切分micro-batch不影响切分sequence length,下面这张图其实已经展现这种结合形式了,不同颜色的方块表示不同micro-batch。

 

TeraPipe的切分,上面展示过了

至于在这种混合切分下,Batch是否还应该切分成均匀的大小,也是不一定的。这需要结合动态规划一起搜索,感兴趣的读者可以直接看原文。

 

Evaluation

 

最后看一下性能,以及工具的评测。

 

E2E Performance

 

因为是无损优化,直接衡量单iteration的latency即可。实验选取的模型全是GPT序列,包括1B, 13B, 44B和175B(GPT-3),具体的配置如下表所示。

 

同时,将input sequence长度统一固定为2048(越长对TeraPipe越有利)。在TeraPipe实验中,也会使用DP算法寻找最佳的both batch和token维度切分组合的Pipeline策略。下图是TeraPipe的性能展示。

 

分析模型实验结果时,发现有预期的结论:在模型越大时,TeraPipe表现越明显。这是因为大的模型消耗大的显存,迫使batch size下降。然而batch size下降是,batch dimension的pipeline并不能充分利用GPU资源,存在很大的bubble,但token dimension的pipeline就不会有这种问题。

 

non-uniform vs uniform slicing

 

如上面的分析,Transformer类模型的structure决定了其token dimension和batch dimension上的切分不应该十分均匀,这里面既有Transformer模型结构的原因,也有细粒度切分不能充分使用GPU计算资源的原因,同时又不能粒度太粗,因为要减少bubble time。在这种情况下,DP确实能找到很不错的iteration latency结果。

 

Longer Sequence Length

 

最后是一个扩展性的实验,因为TeraPipe能够受益于长时序的input,那幺随着时序数目的增加(token数目),TeraPipe会有多大优势呢?这个实验给出了答案。

 

一些思考

 

读完这篇文章后我本人有一些思考和结论,也一同放到文末,供大家批评指正。

Token Level的切分是不是需要用户特殊的写法?

这是针对如何全自动化而言的。我的结论是必须要有用户干预。原因是这样的:一般而言,我们在写transformer中Self-Attention部分时,会通过矩阵乘法实现。如何表达时刻 只和 相关呢?可能单纯的从矩阵乘法上无法表达,只能通过mask做掩盖。也就是说,这种和future无关的语义,如果沿用这种方式,是无法通过计算依赖表达出来的,只能是用户干预,明确指出这种语义才行。

Batch拆分和Sequence length拆分到底哪里不同?

我们知道Batch拆分成多个micro-batch之后,多个micro-batch之间并没有计算依赖,是相互独立的。但sequence length拆分就不是这样,每个时序 都会与前面的时刻有依赖,这就造成了他拆分出来的多个计算图并不是相互独立的。这也是理解Token level pipeline和micro-batch pipeline的另一个角度。

TeraPipe这种只适合某类模型的并行策略是否有价值?

如前文所述,我认为是有价值的。因为Transformer-based LM家族足够大,在预训练模型的用途上也非常广泛,所以应该说具有一定的产品价值。

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注