详解MegatronLM Tensor模型并行训练(Tensor Parallel)_megatron-lm

1. 背景介绍

MegatronLM的第一篇论文【Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism】是2020年出的，针对billion级别的模型进行训练，例如具有38亿参数的类GPT-2的transformer模型和具有39亿参数的BERT模型。

分布式训练的模型并行有两种方式，一种是层间并行(inter-layer)，也就是Pipeline流水线并行，相当于下图对整个模型竖切后每个device各保存3个layer(0,1,2和3,4,5)；一种是层内并行(intra-layer)的方式进行，也就是Tensor模型并行，相当于下图横切后每个device各保留6个layer的一半。

在实际中由于Pipeline并行和Tensor并行是正交的，所以可以同时使用，如下图pipeline并行是竖切，tensor并行是横切，每个色块代表一个device，一个模型运行在4个device上。

接下来重点来看Megatron Tensor并行在transformer模型上的实现。

2. 详细介绍

2.1 Tensor并行计算方法介绍

Tensor计算要进行并行计算，主要方法是通过合理的方式对输入矩阵和参数矩阵进行分块，然后对不同分块分别进行计算。

2.1.1 对参数weight矩阵进行横切(Row Parallel Linear Layer)

以下图为例，$X$ 看成是输入矩阵，$A$ 看成是参数weight矩阵，这里是对 $A$ 横向切分成两个小的矩阵 $A_1$ 和 $A_2$ ，然后为了相乘对应 $X$ 也切分为 $X_1$ 和 $X_2$。

X × A = [ X 1   X 2 ] × [ A 2 A 1 ] = X 1 ⋅ A 1 + X 2 ⋅ A 2

$$\begin{gather*} \begin{aligned} X \times A &= \left[ X_1 \ X_2\right] \times \left[ {}^{A_1}_{A_2} \right] \\ &= X_1 \cdot A_1 + X_2 \cdot A_2 \end{aligned} \end{gather*}$$ X×A​=[X1​ X2​]×[A2​A1​​]=X1​⋅A1​+X2​⋅A2​​​

假设 $X$ 的shape大小是 $(100，300)$, $X1、X2$ 的shape大小都是 $(100,150)$；$A$ 的shape大小是 $(300,200)$； $A1、A2$ 的shape大小是 $(150,200)$；$Y_1、Y_2$ 的shape大小是 $(100,200)$ , 其中 $Y_1=X_1\cdot A_1$ 以及 $Y_2=X_2\cdot A_2$；$Y$ 的shape大小是$(100,200)$。

从前后向的角度来看按行切分 $A$ 参数矩阵的过程，$f$ 函数前向会对 $X$ 输入进行切分两份 $X_1,X_2$，在反向会对回传的梯度通过all-gather方法进行拼接；分开计算完有两部分结果需要合并相加成最终的结果，所以在后面有一个 $g$ 函数前向过程会通过all-reduce方法对结果进行累加，在反向的时候会分别求梯度(也就是identity)。

2.1.2 对参数矩阵进行纵切(Column Parallel Linear Layer)

以下图为例，$X$ 看成是输入矩阵，$A$ 看成是参数weight矩阵，这里是对 $A$ 纵向切分成两个小的矩阵 $A_1$ 和 $A_2$ ，$X$ 是整个参与计算。

X × A = X × [ A 1 ,   A 2 ] = [ X ⋅ A 1 ,   X ⋅ A 2 ]

$$\begin{gather*} \begin{aligned} X \times A &= X \times \left[ {A_1}, \ {A_2} \right] \\ &= \left[ X \cdot A_1,\ X \cdot A_2 \right] \end{aligned} \end{gather*}$$ X×A​=X×[A1​, A2​]=[X⋅A1​, X⋅A2​]​​

假设 $X$ 的shape大小是 $(100，300)$；$A$ 的shape大小是 $(300,200)$； $A1、A2$ 的shape大小是 $(300,100)$；$Y_1、Y_2$ 的shape大小是 $(100,100)$ , 其中 $Y_1=X\cdot A_1$ 以及 $Y_2=X\cdot A_2$; $Y$ 的shape大小是 $(100,200)$。

从前后向的角度来看按列切分 $A$ 参数矩阵的过程，$f$ 函数前向会重复使用输入 $X$ 进行两部分的计算(identity)；在反向会对回传的梯度通过all-reduce方法进行累加；分开计算完有两部分结果需要拼接成最终的结果，所以在后面有一个 $g$ 函数前向过程会通过all-gather方法对结果进行拼接，在反向的时候会对梯度矩阵进行split成两部分，再往后回传。

2.2 Tensor并行在GPT Transformer中的应用

2.2.1 GPT Transformer结构

在GPT Transformer结构中是由一个Attention模块和MLP模块组成。在Attention模块中先是有self-attention层加上dropout组成；在MLP模块有两个MLP层，第一个MLP把维度从H变为4H，第二个MLP把维度从4H变回H，中间是采用了非线性的激活GeLU；每层的连接上也使用了像Resnet的残差连接。。

2.2.2 对MLP模块进行Tensor并行

在进行Tensor并行过程中，要选择哪种方式来对MLP模块中的矩阵进行切分？

先看MLP模块中的第一个MLP层，对应的计算操作可以表达成 $Y=GeLU(XA)$。如果对 $A$ 按行(Row)进行横向切分的话，$X=\left[X_1,X_2\right],A=\left[{}_{A_2}^{A_1}\right],Y=GeLU(X_1{A}_1+X_2{A}_2)$, 由于GeLU不是线性的不好进行后续的并行【$GeLU(X_1{A}_1+X_2{A}_2)\ne GeLU(X_1{A}_1)+GeLU(X_2{A}_2)$】, 所以不采用按行的Tensor切分方式。所以要采用对每一个MLP层的 $A$ 按列(Column)进行纵向切分的方式，对应 $A=\left[A_1,A_2\right],Y=\left[Y_1,Y_2\right]=\left[GeLU(X{A}_1),GeLU(X{A}_2)\right]$。

再来看MLP模块中的第二个MLP层，在第一个MLP层在做完GeLU操作后是有两部分结果，还需要$g$函数进行合并操作，但是可以和第二个MLP层一起计算，这样第一个MLP层的 $g$ 函数和第二个MLP层的 $f$ 函数都可以去掉；对于第二个MLP层没有GeLU操作，要采用对 $A$ 按行(Row)进行横向切分才能接上第一个MLP层。如下图阴影部分是去掉的部分，对应公式如下：

A ( 1 ) = [ A 1 ( 1 ) , A 2 ( 1 ) ] Y 1 ( 1 ) = G e L U ( X A 1 ( 1 ) ) Y 2 ( 1 ) = G e L U ( X A 2 ( 1 ) ) Y ( 1 ) = [ Y 1 ( 1 ) , Y 2 ( 1 ) ] = [ G e L U ( X A 1 ( 1 ) ) , G e L U ( X A 2 ( 1 ) ) ] A ( 2 ) = [ A 2 ( 2 ) A 1 ( 2 ) ] Y ( 1 ) × A ( 2 ) = [ Y 1 ( 1 )   Y 2 ( 1 ) ] × [ A 2 ( 2 ) A 1 ( 2 ) ] = [ Y 1 ( 1 ) ⋅ A 1 ( 2 ) , Y 2 ( 1 ) ⋅ A 2 ( 2 ) ] Y ( 2 ) = Y 1 ( 1 ) ⋅ A 1 ( 2 ) + Y 2 ( 1 ) ⋅ A 2 ( 2 )

$$\begin{gather*} \begin{aligned} \\ A^{(1)} &=\left[ A^{(1)}_1, A^{(1)}_2 \right] \\ Y^{(1)}_1 &= GeLU(X A^{(1)}_1) \\ Y^{(1)}_2 &= GeLU(X A^{(1)}_2) \\ Y^{(1)} &=\left[ Y^{(1)}_1, Y^{(1)}_2 \right] \\ &=\left[ GeLU(X A^{(1)}_1), GeLU(X A^{(1)}_2) \right] \\ \\ A^{(2)} &= \left[ {}^{A^{(2)}_1}_{A^{(2)}_2} \right] \\ Y^{(1)} \times A^{(2)} &= \left[ Y^{(1)}_1 \ Y^{(1)}_2\right] \times \left[ {}^{A^{(2)}_1}_{A^{(2)}_2} \right] \\ &= \left[ Y^{(1)}_1 \cdot A^{(2)}_1, Y^{(1)}_2 \cdot A^{(2)}_2 \right] \\ Y^{(2)} &= Y^{(1)}_1 \cdot A^{(2)}_1 + Y^{(1)}_2 \cdot A^{(2)}_2 \\ \\ \end{aligned} \end{gather*}$$ A(1)Y1(1)​Y2(1)​Y(1)A(2)Y(1)×A(2)Y(2)​=[A1(1)​,A2(1)​]=GeLU(XA1(1)​)=GeLU(XA2(1)​)=[Y1(1)​,Y2(1)​]=[GeLU(XA1(1)​),GeLU(XA2(1)​)]=[A2(2)​A1(2)​​]=[Y1(1)​ Y2(1)​]×[A2(2)​A1(2)​​]=[Y1(1)​⋅A1(2)​,Y2(1)​⋅A2(2)​]=Y1(1)​⋅A1(2)​+Y2(1)​⋅A2(2)​​​

最终对应整体的图如下：

上图MLP模块中的 $f$ 和 $g$ 函数在PyTorch中的伪码如下：

class f(torch.autograd.Function):
    def forward(ctx, x):
        return x
    def backward(ctx, gradient):
        all_reduce(gradient)
        return gradient
        
class g(torch.autograd.Function):
    def forward(ctx, x):
        all_reduce(x)
        return x
    def backward(ctx, gradient):
        return gradient
1
2
3
4
5
6
7
8
9
10
11
12
13

2.2.3 对Attention模块进行Tensor并行

回顾下Attention的计算过程有哪些矩阵计算, 首先是在前面$QK/V$计算中每个Head中$Q/K/V$对应有三个weight矩阵，然后是在最后对$Z_n$ 进行汇总时要用到weight矩阵 $W^0$。

Attention模块中对Tensor并行切分方式跟MLP类似，先分别对 $Q/K/V$按列进行切分，计算的结果是多个独立的部分，然后对最终的weight矩阵进行按行切分，得到的结果进行累加操作(all_reduce)得到最终的结果。因为在前向输入是多次复用的，所以在反向时需要对梯度进行累加操作(all_reduce)。这样attention计算中也有两次all_reduce操作。

最终对应整体的图如下：

在整个Transformer结构中通信上共需要4次All-Reduce操作。

2.2.4 对输出Embedding层进行Tensor并行

以下是GPT计算的整体公式：$W_e$是输入的embedding层, $W_p$是position embedding层，$W_e$同时也在最终输出时复用了，embedding层的shape大小是 $hidden-size(H)\times vocabulary-size(v)$ 。

在GPT-2中词表大小是50257; 为了加速并行对embedding的权重$E_{H\times v}$ 按vocabulary的维度(也就是按列切分)进行拆分，结果 $E=\left[E_1,E_2\right]，GEMM[Y1,Y2]=[X{E}_1,X{E}_2]$，并行时通过all-gather通信得到最终结果：$Y=all-gather([Y_1,Y_2])$, 然后再计算交叉熵的loss陨失，通信量是【batch-size x sequence-length x vocabulary-size】，这里vocabulary-size往往过大造成通信代价大。为了降低通信量，将 GEMM[Y1, Y2] 和 cross entropy loss进行fuse融合，可以降低通信量到 【batch-size x sequence-length】。

3. 参考

• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• MEGATRON-LM: TRAINING BILLION PARAMETER LANGUAGE MODELS WITH GPU MODEL PARALLELISM
• The Illustrated Transformer