前文提要：

上一篇：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中被引入。