最近比较火的chatGPT内部使用的模型是GPT系列。GPT属于一种Large Language Model，专注于自然语言生成任务，运用了transformer的decoder模块堆叠而成。与它对应的是BERT，专注于自然语言理解任务，运用了transformer的encoder模块堆叠而成。

　　在我看来，其实不论chatGPT还是BERT还是AIGC，背后的核心技术，或者说门槛比较高的原因，一个是海量训练数据，一个是大算力。这些大模型的训练，一般需要用到上千个GPU，服务器集群。本文主要介绍一下在大模型训练过程中用到的一些算法层面的技术。

　　大家都知道，数据多了模型大了就需要分布式训练。核心是我们常说的数据并行，模型并行等技术。数据并行表示一张卡能放下，然后多张卡保存多个模型，组成一个大的batch进行训练。而模型并行是一张卡放不下，需要多张卡共同跑一个模型。模型并行又包含pipeline并行和tensor并行。

　　目前绝大多数Large Language Model都是在Nvidia的软硬件体系下完成的。所以本文主要介绍一下Nvidia发布的专注于LLM训练的megatron-lm框架，看一看它到底用了哪些技术。

　　介绍具体方法之前，先有两个基本知识作为铺垫，一个是NN backward，特别是gemm的backward的简单说明，一个是矩阵分块技术。

　　根据链式法则，每层的反向过程输入是上层传过来的梯度。输出就是对x和对w的梯度。计算w的梯度，需要用到当前层的输入x。计算x的梯度，需要用到w。

　　1． 反向过程的执行时间可以粗略估计为前向过程的2倍（因为前向算一个输出，反向算两个）

　　2． 反向过程中权重的梯度计算需要用到输入tensor（也就是上一层的输出），所以需缓存住前向过程的输出。

　　模型并行包含pipeline并行和tensor并行。pipeline并行表示对模型按照层做切分但不对一层算子做切分，可以理解为一张卡上跑的层都是完整的。而tensor并行表示把一层进行切分，这样尽量让多个层在一个卡中，这样的话，如果模型能从头切到尾，那一张卡就保存了整个模型层但只做部分tensor数据，如果中间层不进行卡间通信的情况下，性能还是非常高的。megatron-lm引入的是tensor并行技术。

　　当然，model parallel是可以和data parallel共同使用的。比如一个模型分布在8张卡上，用了模型并行，但是有8台8卡机器，这8台机器上就可以做data parallel。

　　矩阵乘可表示为Y=X*A，其中，X为输入hidden_state，shape是b,s,h，A是权重，矩阵乘的主要作用是扩大hiEBET易博真人dden size 4倍，所以权重一般的shape都是h,4h。由于A比较大，对于tensor并行模式来说，一般是需要考虑A的切分。

　　有两种切分方法，一种是按照row来切分，那么X也需要切分，以切成两部分为例，表示为：

　　如果按照A column切分，考虑第二个GEMM的输入，相当于X被按照列进行切分了，所以它的权重也需要按照行来切分，具体就如下所示：

　　这样的话，做完整个MLP之后多卡是需要进行all-reduce的。总结下来MLP的tensor parallel流程就如下所示：

　　对于反向传播，如果是2卡划分，B1的shape就是row/2, col, Y1的梯度计算等于传过来的Z1的梯度乘以B1的转置，输出的Y1梯度shap就是col，row/2，可以理解为Y按列划分。传到计算X的梯度时，由于A1是按列划分的，乘以A1的转置之后，shape又恢复了X输入的shape，那么卡间要做all-reduce完成梯度累加。和前向传播正好时共轭的。

　　可以看到，这里softmax的计算被划分到了不同的卡上，而进行softmax的其实是q*k转置得到的shape为seq_len, seq_len的矩阵，按照切分流程，这里的q和k其实只是左半边或者右半边，那么分别进行softmax肯定不是数学上等价的。但是需要注意的是，这是一个多头自注意力模块，qkv weight被分成了多头，而每个头的softmax是独立处理的，这里的切分不会破坏一个头对应的weight，而是把多个头放到同一个卡中去处理。所以softmax不会出现不等价的计算结果。

　　其中，f代表前向过程不执行操作，反向过程执行all-reduce，代表前向过程执行all-reduce，反向过程不执行操作。

　　在流程图中可以看到，LayerNorm等模块是没被并行执行的。虽然这些模块计算量不大，但是为了减少通信，每张卡上都会计算一份完整的layernorm等操作，这样造成的activation占用mem还是比较大的，公式中的10sbh就是由这些引起。

　　具体来说，由于Layernorm的作用范围是每个token自己统计均值方差然后norm，所以sequence尺度上来看是相互独立的。所以在mlp之后，插入一个reduce-scatter的操作，该操作可描述为，每个卡将自己的数据分为同等大小的数据块，然后每个卡将根据index得到的数据做一个规约操作，即先做Scatter再做Reduce。

　　体现在具体shape上，就是将大小为seq x h的输出数据（每张卡上只有自己那个part，并不是最终的完整数据），先在seq方向划分为t份，然后只取自己对应的seq/t长度的数据，然后卡间all-reduce，每个卡就有了seq/t的完整数据。所以，在流程图上看，插入的通信算子由f变成了g。

　　为了进一步减小mem占用，大部分算法都使用了activation重计算的技术，具体就是说存储一份该group的初始输入activation，而不存储group中每层中间的activation，在backward执行过程中，重新走forward，把中间的activation计算出来给梯度计算使用。传统recompute方法是每个模块都全部重新计算，这样的线%之间。

　　通过act占用mem公式可以看出，5as/ht是由attention部分引起的（包括除去生成qkv的线as/ht=80，对于MT-NLG大模型，a=128，s=2048，h=20480，计算出5as/ht=64，这些值和公式中的34相比都要大一倍多。如果该方法把这部分采用重计算（如下图红色虚线框所示），那么mem将节省65%~70%。

　　考虑到反向传播要计算每层intput data和weight的梯度，计算量是前向的2倍。这样，整体的计算量如下所示：

　　基本的pipeline并行表示将模型中按层进行切分，但不对单层进行切分。具体如下图所示，有4个device，就把一个模型分成四份，并将需要处理的一个大的batch分为8个小的micro-batches（由于需要在训练尺度来看，还应该是一个大batch在做处理，所以引入了Gradient Accumulation的技术，就是每个micro-batch前后向计算后的梯度累加，在最后一个micro-batch累加结束后，统一更新模型）：

　　之所以需要划分batch，就是为了并行起来，设备1处理完一个microbatch的数据之后，设备2马上就可以进行后面几层的处理。但是这样做有两个问题，一是需要缓存多个microbatch的activation，比较费显存。二是pipeline bubble问题。

　　想让这个比值变小，最直观的就是增大m，但是增大m会带来显存占用增加的问题。

　　此方法的核心思想是一个microbatch前向完成之后马上执行反向。从而减少了需要缓存的显存占用。但是这个方法并未减少pipeline bubble，megatron基于此问题又提出了interleaved schedule，具体流程如下图所示:

　　该方法的核心思想是重排每个stage占用的layer index。比如原来1f1b策略中layer是这样分布的：

　　以每一个microbatch来说，每个设备需要执行四层才能启动下一个设备的执行。Interleaved schedule修改了layer的排布，变成：

　　不错此方法会带来更多的通信开销，所以v设多大也是一个trade off的问题。

　　结合计算量与访存量的trade off。最终选择三种方法相结合的混合并行方法。一个block的整体流程如下所示：

　　上述过程未考虑pipeline parallelism的融合，对于超大模型来说，一般是需要多种并行技术的融合使用，以参数规模1T的GPT而言，有128层的Transformer Layer，这个超大超深的网络被分割成了64个stage，每个stage跑在6台DGX-A100上，其中6台机器之间进行数据并行，每台机器内部的8张卡之间做模型并行，整个集群的3072张A100按照机器拓扑被划分成了[6 x 8 x 64]的矩阵，同时使用数据并行&模型并行&流水并行进行训练。

　　其中Hardware FLOP表示引入重计算方法，计算量肯定是增加一些的。可以看到megatron中使用的方法比传统重计算的throughput要提高30%左右。

　　除了megatron-lm，还有一些性能非常好的框架比如ColossalAI，OneFlow等。它们也使用了数据并行，模型并行等技术的融合。具体后面再进行介绍。