前文提要:
上一篇:LLM分布式训练第三课-模型并行之流水线并行 (qq.com)
上上一篇:LLM分布式训练第二课(数据并行) (qq.com)
上上上一篇:LLM分布式训练第二课(数据并行) (qq.com)
张量并行不是张亮并行,不是麻辣烫
如果上一节介绍的流水线并行是把模型基于层给进行了划分,来让多张GPU的显存可以承载规模较大的模型,那么这一节介绍的张量并行就正好用另外一个角度来解决单张GPU显存不足的问题。
张量并行其实也有两个细分的子流派,行并行和列并行。
我们用 GEMM 来拆解模型如何并行,以Y =XA 举例,对于模型来说,X 是输入,A是权重,Y是输出。
行并行(Row Parallelism):
行并行简单说就是把权重A给按照行来分割为2部分,为了输入X要去匹配A被按行切分的状态来进行计算,所以把X也给切成2部分,因为要矩阵乘,所以X得竖着切,如下图所示。
而Y=XA就被拆解成:
行并行
X1和A1就被放在一块GPU上进行计算,而X2和A2就放另一块GPU上进行计算,通过这个方式完成了模型的并行。
行并行计算-1
行并行计算-2
如上两图所示,两块GPU分别算出来了各自的矩阵Y1和Y2,然后用矩阵加法将两个矩阵的值进行相加,得到和原来计算结果等值的矩阵Y。
列并行(Column Parallelism):
列并行计算-1
列并行计算-2
列并行的计算方式和行并行有很大的区别,期中最重要的就是三点:
1- 列并行的A是按着矩阵的列去进行拆分的
2- 列并行的输入X是不需要拆分的,因为矩阵乘,行乘以列,A进行列切分,列维度没变,维度是相等的。
3- 最后的Y1和Y2不是相加的关系,是contact的关系,将两个矩阵合为一个矩阵Y
目前看起来似乎行并行和列并行没有什么太大区别,得到的值也是一样的,而且列并行需要把X复制两份分别和A1和A2进行矩阵乘会消耗更多的显存。
但是如果考虑了激活函数呢?
比如要连续过两个激活函数层,例如2层以上的Transformer,每一层都会有一个MLP,就要过一遍Relu或者Gelu函数,我们以Gelu为例:
上面的式子在列并行的情况下:
因为列并行并没有进行任何的输入拆分,所以只要把A激活函数层和B激活函数层划分好,就可以独立计算,在计算出GeLu(GeLu(XAi)Bi)后(i=和X计算的被拆解的子矩阵号,如1,2,3…),最后进行contact就可以,换个说法只要在得到最终结果之前通信一次就行。
如果是行并行呢?
由于GeLu是非线性的函数,所以:
也就是说,在整个计算流程中,每经过一个全连接层,都必须要通过通信来聚合成最终的结果,然后才能进入到下一个层来进行计算,过大的通信量会极大的降低模型的训练速度,增加延迟。
2D/2.5D和3D并行:
其实对于2D并行有两种解释的说法,比如TP+PP,或者TP+DDP都算2D并行的范畴,因为是从两个维度来支持更好的分布式,降低单卡显存和计算压力。
另外一种关于2D并行的解释是专门针对ColossalAI来讲的,一般称为2D张量并行。
一般会把基于Megatron的Tensor方式称为1D并行,1D并行的一个弊端是,对于刚才的函数Y=XA,在计算的过程中,并没有对激活Activation进行划分,导致激活这部分会消耗大量的显存,也就是每块GPU虽然参数被分开了,但是激活还是每块都有,还有一个重要的点是,如果采用1D并行,那么所有的GPU都要和其他的GPU进行通信,all-reduce或者其他的源语通信,通信成本越高,整体训练的性能越差。
基于以上的原因,Colossal-AI引入的2D张量并行的概念。
还是一个简单的函数Y=XA,如果我们拥有P个GPU,P必须满足q的平方,比如拥有4个GPU,那么q就是2,q*q=4,在这个前置条件下,我们把输入X和权重A都拆成q*q的子矩阵,即2个拥有用4个子矩阵的矩阵。
这个计算一共包含q个步骤,如上式而言实际上是2个步,我们首先让X矩阵的第一列和A矩阵的第一行在所有的4个GPU中进行广播即:
然后让上式的每2个子矩阵在相应的GPU上进行矩阵乘,在单位时间里,这个计算是并行的,并且4个GPU的任何一个都没有保存其他GPU的Activation的必要。
同样的在第2步,我们可以得到:
最终我们把Y=XA分解为:
虽然两个大矩阵中间要进行串行操作,但是在大矩阵内部的4个子矩阵都是进行并行的操作。
加入有1万张GPU卡,如果是1D并行的话,其中任意一张卡都要和其他9999个机器通信,而2D并行划分了子单元,每个机器理论上只需要和96个机器进行通信,极大的节省了通信的代偿和开销。
2.5D并行其实就是在2D并行的基础上加了一个维度,如图所示
2.5D并行
还是以Y=XA这个函数为例,这次P个GPU被分解成d*q*q,为了计算流程看起来清楚,假设d=q也为2,所以这个tensor为[2,2,2]。
现在把输入的X划分为d*q*q,来满足P个GPU均匀分布,得到下面式子:
这个式子其实可以被表达为下面两个子矩阵的contact,我们把大矩阵拆解成下面两个子矩阵。
然后权重A被分解成q*q个单位:
对于X的每一层,我们都使用2D算法和A做矩阵乘,就得到以下两个式子:
将这两个式子进行垂直contact就能得到最终的结果。
2.5D主要能进一步优化2D的通信代偿,但是实际生产中使用的不多,仅作为算法让大家理解它的原理。
3D并行在2.5D的基础上,把A矩阵也拆解成d*q*q,或者可以理解为X和X矩阵都被拆解成q*q*q,如下式:
每升一个维度,通信代偿都会得到进一步的下降,看明白原理就可以,在这里就不赘述了。
关于3D并行业界比较通用的解释是立体的并行手段,如图所示:
3D并行
3D并行,目前业界共识的叫法主要是针对用多种并行训练方式来进行训练,如上图所示,在GPU4/12/20/28这个维度,通过流水线将模型切割成不同的stage,然后在每个stage内部,又通过模型并行来进行横向的划分,然后在GPU0和GPU4这之间,因为他们的模型参数都是相同的,所以又可以采用数据并行来增大训练的dataset吞吐量,提升训练速度,这是一个典型的3D并行训练的案例。
Pytorch TORCH.DISTRIBUTED.TENSOR.PARALLEL
不管是 Megatron 还是Colossal 中的张量并行,都是基于 Transformer模型来实现的张量并行,不具备通用性。Pytorch 作为一个通用框架,提出了自己的并行方式,Dtensor,即TORCH.DISTRIBUTED.TENSOR.PARALLEL,可以更简单的在 SPMD(单程序多设备)中进行分布式计算。
Dtensor在Pytorch2.0中被引入。