BMTrain类Megatron+DeepSpeed原理学习

这一章节虽然是BMTrain，不是目前常用的Megatron+DeepSpeed，但是对于了解原理，也是很有帮助。

BMTrain

数据并行

一般数据并行

上图，把数据切为3份，每张显卡处理一部分数据，每张显卡利用得到的数据进行前向传播和反向传播，得到各自的梯度，为了让模型学到这份数据的所有知识，就需要把这些梯度的信息进行一个聚合，也就是取平均的操作。那么我们得到聚合好的参数去更新模型，就能学到切分3部分数据合起来的完整的数据的知识。
具体来说：在数据并行的过程中，我们有一个参数服务器，里面保存了模型的参数，还有完整的一批数据，前向传播的过程中，参数服务器上的参数，会被复制到所有的显卡上，每张显卡上都得到了很参数服务器上一样的模型参数，然后把数据切分为3份，每张显卡上各拿到一部分数据，然后每张显卡用完整的模型参数和一部分的数据去进行前向传播和反向传播，我们就能够得到每张显卡上各自的梯度，最后将这个梯度进行聚合，将聚合后的梯度传回我们的参数服务器。那么参数服务器上面有了原始的模型参数和这个聚合好的模型的完整的梯度，我们就可以用优化器去对模型的参数进行更新，那么更新后的参数又会进入下一轮的模型的训练迭代

集合通信

1. broadcast
   把数据从一张显卡传到其它所有的显卡上
2. reduce
   规约可以是求和、平均、MAX等，把各张显卡上的数据进行一个规约，然后把规约得到的结果，放到我们其中一张指定的显卡里面。
3. all reduce
   和reduce几乎相同，不同点就是把结果发送到所有显卡上面（reduce是发送到指定的一张显卡上）
4. reduce scatter
   跟all reduce相比，相同之处是它们都把规约得到的结果发给所有显卡，不同之处在于，reduce scatter，最后每张显卡上只得到一部分的规约结果。
   如上图，0号显卡，会得到in0的前1/4的参数，加上in1的前1/4的参数，加上in2的前1/4的参数，加上in3的前1/4的参数，得到out0。其他同理。
5. all gather
   可以跟all reduce类比下，把各张显卡上的数据进行收集，然后进行一个拼接，然后广播到所有显卡上，所有显卡得到了一个搜集后的结果。

分布式数据并行

对数据并行进行了优化，没有参数服务器。
每张显卡各自完成参数更新。然后保证参数更新后的结果一致。

数据并行显存优化

计算过程的中间结果是跟batch乘以句子长度和模型维度相关的一个显存占用。使用数据并行的时候，把一批数据分成了N份，让每张显卡只处理其中的一部分数据，等效于每张显卡上所处理的batch的大小将变成了原来显卡数分之一，那么通过把这个输入的维度进行了一个降低，那我们模型整体的中间结果量（每张卡），也进行了一个降低。
但是这个方法有一个缺点，就是为了支持模型的训练，它至少需要训练1条数据，最极端情况下，每张显卡只得到一个数据的时候，由于我们的参数哈需要完整的保存在显卡上，梯度也需要完整保存在显卡上，优化器也需要完整保存在显卡上，那么即使中间结果一点都不在显卡上，我们模型仍然有可能无法在一张显卡上进行计算。

模型并行

一张显卡无法放下模型的所有参数、所有梯度、所有优化器，就想办法把一个模型分成很多个小部分
思路：针对线性层的矩阵乘法的例子，上图左上角，3行2列的矩阵乘以2行1列的矩阵，本质上可以把它的结果分成三个 部分。
可以把3X2的矩阵看成线性层的参数W，把2X1向量看成线性层的输入。

通过上面的方法，我们的线性层的参数，就可以划分到多张显卡上。而且我们需要保证多张显卡上模型的输入是一样的。那么我们就不能使用数据并行的那一套方式对数据进行划分了，我们需要保证每一张显卡上得到的输入是一样的，所以他们是同一批数据。
得到多个子结果，然后使用all gather收集算子。
模型并行对显存的优化：首先，每张显卡只需要保存原来N分之一的模型参数，N是我们GPU的个数。由于每张卡只保留一部分参数，梯度、优化器也只用保存对应的一部分（变少了）。
由于这个矩阵乘法，需要保证数据的输入是一致的，不能使用数据并行的那一套，去把数据切分成很多个小份，因此我们模型的计算的中间结果实际上是没有减少的，这也是这个方法所带来的弊端。既然我们的中间结果无法减小，当batch size开的比较大的时候，仍然会出现显存溢出的问题。

ZeRO

ZeRO-Stage1

基于数据并行建立起的一套框架。

数据并行之中，需要对模型的梯度进行规约，为了保证每一轮迭代之后，每张显卡上得到参数仍然是一致的，我们就让每一张显卡上都得到了规约后的参数。然后每张显卡各自去进行参数的更新。可以发现1号显卡用的是同样一批参数和梯度去进行参数的更新，而2号显卡也是用的同样的一批参数和梯度。其实所有卡用的都是同样一批参数和梯度，他们各自去进行参数优化是不是就带来了计算上的重复和冗余。

为了消除这个样的冗余，提出的方法就是（上图下半部分）每张显卡上只获得一部分的梯度，然后只去更新一部分的参数，这样我们各张显卡可以合作，每张显卡只更新一部分的参数，最后他们进行一个交流，就能把这个模型的完整梯度给恢复出来。

完整过程如下：

由于是基于数据并行的架构，那么它仍然是每张显卡上保存完整的模型参数，然后有一部分的数据，通过前向传播和反向传播的到得到各自的梯度，从这之后的步骤开始就不一样了，它不是使用all reduce让每张显卡得到相同规约后的梯度，而是使用reduce scatter，是让每张显卡得到一部分的结果，比如1号显卡得到了规约后的梯度的前1/3，2号显卡得到了规约后的梯度的中间1/3，3号显卡得到了规约后的梯度的后1/3。
我们可以让1号显卡用这前1/3的梯度和前1/3的梯度对应的optimizer，去更新前1/3的模型的参数。然后2、3号显卡同理。每张显卡各自更新得到了一部分的模型参数之后，我们需要通过一个收集的操作，把各显卡分工合作的结果进行一个收集后，再进行一个拼接。拼接之后的结果通过All,把这它告诉所有的显卡，经过这样的操作之后，每张显卡上又得到了完全一样的参数和一致的结果。

ZeRO-Stage2

stage2进行了一个优化，在stage1中，需要再反向传播得到所有的梯度之后，对这个梯度进行reduce scatter，那我们reduce scatter之后，每张显卡上各自得到了一部分规约后的梯度。得到这个gradient*之后，原来的gradient实际上是不需要保存在显卡上了，因为后面参数更新的过程中，优化器只需要用到gradient*去进行一个参数的优化，那么这部分gradient其实已经可以从显存中移除。

什么时候移除呢？

stage1阶段，在反向传播结束之后，才去把这个gradient移除。那可不可以在这个反向传播的过程中，就把gradient*提前算出来，然后把这个gradient删掉呢？具体来说，假如有一个24层的transformer，在反向传播结束第24层的时候，就用第24层各卡上的梯度进行一个reduce scatter，之后，第24层的梯度就可以不要了，到第23层的时候，计算第23层各卡上的模型梯度，进行一个reduce scatter，得到那一层的gradient*，这个时候，就又可以只保存第23层的gradient*到显卡上，然后把第23层的gradient从显卡上移除。通过这样的操作，把gradient*从显卡上移除的操作提前了，就可以进一步节省显存空间。

ZeRO-Stage3

这个阶段对模型的参数进行了进一步的划分，在stage1和stage2中，都没有避免前面模型并行所解决的那个问题(模型的参数仍然要完整的保存在我们的显卡里面)，前面也计算过11B的模型，它保存在显卡里面就已占用了约40GB的显存。

为了解决这个问题，就需要把模型的参数也进行一个划分，因为每张显卡上只保留了一部分的梯度去进行参数更新，参数更新也只更新一部分的模型的参数，那么我们实际上每张显卡就可以只保存它自己参数更新所负责的那一部分参数，那么每张显卡上就会只有一部分的参数和一部分的数据、一部分的梯度和一部分的优化器。
那么每张卡上只有一部分模型参数，怎么进行模型的前向传播、反向传播呢？这两个过程中，需要把模型的参数进行一个all gather的操作。在用到模型参数的时候，我们去临时把每张卡上的那一部分参数进行一个收集拼接，当我们用完的时候，就把这个参数从显卡中释放。

具体而言，在transformer中一个线性层，这个线性层的参数可能被分到3张显卡上，需要计算这个线性层的时候，临时把跟这个线性层有关的一部分参数从各张显卡上收集起来，进行一个线性层的计算，进行完这一层线性层计算之后，模型参数就不需要放在显卡里了，又可以释放掉了，又恢复层每张显卡只保留一部分线性层参数的状态。

需要注意的是，使用这样的方式，我们在方向传播计算模型梯度的时候，实际上也需要用到模型的完整的参数，也是需要gather操作。相比stage2来说（它只在参数更新那一步进行了all gather操作），而在stage3中，它虽然在阐述更新的时候没有使用all gather，但是它在前向传播和反向传播的时候，各需要使用一次all gather，实际上是增加了一次通信，相比stage2，stage3是一个用时间换空间的一个算法。

显存分析

分析下各阶段显存占比。

stage1

在stage1中， 每张显卡上只需要处理一部分的模型梯度，优化器降低到原来的卡数分之一，由于ZeRO是一个基于数据并行的架构，我们把中间结果的量，也降低到了原来的卡数分之一。

stage2

进一步把模型的梯度的划分提前，把reduce scatter提前到了反向传播的过程中，这样，就不需要保留完整的梯度，没删一层，我们就可以把梯度划分到卡数分之一。

stage3

由于每一张显卡只负责一部分的参数更新，我们进一部分的把参数划分。这样参数、梯度、优化器和中间结果都得到了划分，每张显卡只需要保存这些参数的卡数分之一的部分，已经是一个非常完整的优化方法。

流水线并行

跟模型并行有类似之处，模型并行的方法通过把线性层分成很多个小的矩阵，然后把这些小的矩阵分到各张显卡上。
流水线并行的方法，模型的不同的层分给不同的显卡，具体的说，假如我们有一个三层的transformer，可以把transformer的第一层，分到我们第一张显卡上，第二层分到第二张显卡上，第三层分到第三张显卡上，那我们进行模型前向传播的过程中，我们需要再第一张显卡上，完成第一层的模型计算，然后把计算的结果传到第二张显卡，第二张显卡拿到第一张显卡的输出，去进行这一层的计算。

我们可以看到这样的方法，它的显存占比，仍然是参数、梯度和优化器、中间结果都得到了划分，因为每张显卡上只保留了某些层的参数，也只用保留对应的梯度和对应的优化器。而且由于模型的层数变少了，虽然没有使用数据并行的方法，模型层数变少之后，我们的中间结果仍然也是得到了减少。但是这种方法存在一个弊端在于，我们在1号显卡计算的时候，其他显卡实际上是处于一个空闲的状态，这会是一个计算资源上的浪费。后续对流水线并行的一些优化方法，对资源浪费的问题进行了一定的补充和优化（bubble）

优化技术

混合精度

在实际模型训练过程中，模型的参数一般不会超过千这个数量级的，那么完全可以在FP16的数值表示范围之内。
FP32转成FP16，进行速度提升，其实是面临一些问题的：
就是我们参数更新的时候，参数更新的量，其实是约等于我们的梯度乘上学习率，学习率一般是1e-3,1e-4，甚至是1e-5的数量级，FP16的数值的最小值，它是一个1e-5的数量级的数，假如我们梯度乘上一个非常小的学习率，就会低于这个F16的表示范围，那么我们参数更新量是否就会产生丢失，也就是我们所说的下溢。
假如相乘后结果小于6e-5，那么它在FP16表示下就会变成零。那它对参数的更新就会被忽略。
为了解决这个问题，我们需要扬长补短，既然FP32能表示储更大范围和更高的精度，我们就把FP16的梯度乘上学习率，得到的这个参数更新量表示为FP32，我们用更高的精度去表示这个参数的更新。我们肯定还是不能把这个更高精度的参数更新，去加到我们的更低精度的模型上，因为实际上这样还是会产生浮点的下溢，所以我们需要在优化器上额外保存一个，单精度的参数。
具体来说， 我们在一般的模型训练中，一个模型会有一个FP32的参数，FP32的梯度，然后优化器会使用FP32的梯度去进行一个参数优化，更新到FP32的参数里面（如下图左边）。
而在混合训练中（如下图右边），为了加速模型的前向传播和后向传播，在模型中会使用半精度的参数和半精度的梯度，通过半精度的梯度传到我们的优化器里面进行一个优化器的更新，这个优化器的更新量需要把它保存为FP32类型，
我们把这个FP32类型通过优化器里临时创建的FP32的参数，进行一个有效累积，有效累计之后，我们再把它重新转回我们的FP16参数，来供模型进行计算。

offloading

前面提到优化器的参数，比如以adam为例，它会是一个相比于参数量是一个至少两倍的关系， 我们把它放在显卡上，它就会使一个显卡上的显存占用的大头，能不能想办法让它从显卡中移除？

可以把它从显卡上放到CPU中，需要先把模型参数的梯度从显卡中传给我们的CPU，然后在CPU上进行优化器的优化，然后将优化的结果重新传回给显卡上。

可能有人会说，既然GPU的计算速度比CPU快那么多，通过这样的一个offloading操作，会不会导致CPU变为整个模型训练速度的瓶颈。实际上使用ZeRO3之后，参数梯度以及优化器都已经划分到卡数分之一的大小。通过把一张显卡绑定到多张CPU上面，我们就可以把这个每个CPU上面的计算量降到足够的低，能够让CPU不会成为模型训练的一个瓶颈。

overlapping(通信的计算重叠)

在GPU中，它的memory操作一般是异步的，就是可以通过提前给memory发送一个请求，然后去进行其它的计算，进行其它计算完之后，来对那个memory的请求进行一个接收，这个两步的一个操作。
在模型前向传播的过程中，可能需要把layer1的参数通过gather操作进行一个获得，之后进行layer1，然后再对layer2的参数进行一个优化，在获得完第一层的参数之后，我们在第一层forward计算过程当中，提前把layer2这一层的参数的获得进行一个异步的提前，那么在layer1的前向传播计算完成后，layer2的参数也已经获得并放置在显卡中了，那么就可以立刻的去进行layer2的前向传播的计算，在layer2前向传播计算过程中，又可以提前把layer3的参数进行一个gather。

checkpointing

存档点，模型的中间结果，为了支持模型的反向传播，我们需要把模型的前向传播中的所有中间结果都保存在显存当中。那我们是否可以通过存档的方式去进行优化，我们不把所有的结果都保存在显存当中，而是只保存一定的存档点呢？

如上图左边，是一个transformer的结构简化示意图，每个transformer block会有一个模型的输入，输入之后会经过attention，feed forward等复杂的计算，得到这一层的一个输出，那我们只保留transformer每一个大层的模型的输入，作为我们的检查点。
那么在方向传播的过程中，如何对这个大层里面的这些线性层的梯度进行一个计算呢？我们可以进行一个重计算（上图右边）。重计算就是说通过transformer每个大层的输入，然后在方向传播的过程中，我们重新对它进行一个前向传播，临时的得到每一个大层里面，所有所有线性层的输入，那么我们得到的中间结果，我们就可以进行反向传播，我们反向传播过了这一层之后，就可以把这个检查点以及临时重计算的这些模型层里面的线性层的中间结果，从显存中删掉。假设我们有一个24层的transformer，每一层里面会有5个线性层，那我们通过检查点的方式，我们能把原来需要保存24X5等于120的中间结果，降低到了24的中间结果。

BMCook

模型压缩相关技术

知识蒸馏、模型剪枝、模型量化、参数共享、矩阵低秩分解、结构搜索

ALBERT: 参数共享；词表向量做一个矩阵分解，拆成两个更小的矩阵。

BMInf

目标是把百亿参数大模型跑到只有6GB的显卡上。

难点：

• 很高的显存占用
• 模型推理算力要求比较高

transformer分析

包含多个transformer block块，其实里面包含很多线性层做线性运算，对CPM-2模型进行统计，它虽然有百亿参数，但其实里面90%以上的参数，都是在线性层，也就是说有90亿的参数都是线性层的参数。另外反向模型推理的过程中，模型大概75%以上的时间都是在进行一个线性的运算。所以首当其冲的是优化线性层。
精度损失的前提下来优化整个线性层的运算效率。

量化

简单方法

举个例子，如下图

因为int8范围只有-128到127。原本矩阵中（上图左）中哪些浮点数，得把这些数缩放到一个-128到127之间，简单做法就是找到矩阵中最大的那个数（是绝对值最大，本例中是3.1），把它缩放到-127，相当于用3.1除以127得到一个缩放系数0.024，然后把左边的那些量除以这个缩放系数，得到一个比较大的值，然后通过四舍五入找到离它最近的整数，这样就能把左边浮点数转换成右边缩放后的INT8的矩阵。还原回去就是用INT8矩阵乘以缩放系数。

上图，两个IN8的矩阵（量化后得到）进行乘法，这里结果（上图右下角）INT8肯定是存不下，一般是用INT32来存，会得到一个整数结果，同时针对前两个相乘的INT8矩阵缩放的系统进行一个标量的惩罚，然后得到一个新的缩放系数，然后把相乘得到的结果（上图右下角）乘上得到的缩放系数，把它存下来，然后还原回去， 就得到一个INT8模拟的一个矩阵。
这个方法有误差，但是不大，通常在一些尺寸比较小的时候会表现的比较好，比如像CV里面卷积核，这个方法不会造成太大的损失。
但是直接把这个方法搬到transformer上，结果会很差。
transformer很大，如果还只使用一个缩放因子，相当于对整个矩阵中有好几百万种不同的值，现在变成最多只有256种值（-128到127）。这样的话，整个矩阵的表达能力会有一个直线的下降。所以transformer不能通过这种方法进行量化。

适配transformer的量化

为了适配transformer这种情况，需要做一个更精细化的一种缩放，这里使用了一种行级别的一种矩阵量化，相比之前整个矩阵量化，我们将量化的一个力度，从原来整个矩阵变成一行或者一列。

以前计算整个矩阵的缩放系数，现在计算单行(列)的缩放系数，将一行（列）给缩放到-128到127之间。发现在transformer上达到一个不错的效果。
如果取值更大的时候，可能还可以做到一些更细粒度的一些不同的方案，比如矩阵和矩阵的乘法，可以拆成子矩阵。
使用上述方法，成功将CPM-2从原本的半精度22GB减少到11GB，同时速度还提升了一倍。但是对于6GB显存的GPU还不够。

memory scheduling

类比操作系统的虚拟内存，这样的机制也可以用到GPU上，在进行百亿模型推理的时候，不会同时用到这11G参数中的所有东西，每次只用一部分，比如transformer中，我们每次指挥计算一层，其实只用到这一层的参数，那些不计算的层的话，其实没必要一直放在GPU上，可以把那些暂时不会用到的参数，给传回到CPU上或者放到CPU上，而把那些短时间内会用到的参数给重新传回到GPU上重新加载上来，这样的话，我们可以在计算每一层之前，将这一层的参数从CPU传到GPU上，算完之后再把这个参数给扔掉。我们就可以在一个更小的显卡限制下，把我们的模型跑起来。
目前主流的GPU中，都是支持同时去进行数据传输和运算的，也就是说可以一边算一层，一遍拷贝另一层参数（CUDA6里面已经实现，被称为unified memory），BMInf重写了这部分调度算法。
深入思考，如果我们能够在算一层的同时，去加载另外一层，那么理论上只需要两层，在最优情况下，整个模型就能完美地跑起来。算第零层的时候，把第一层加载进来，第零层算完了，然后第一层刚好下载好了，就把第零层扔掉，这是我们可以把扔掉的这部分空间给重新利用起来，去加载第二层。这样就能够在一个有限空间中把模型跑起来。
但是，实际操作过程中，发现传输一层参数的时间，往往超过了计算这一层所花费的时间。如果我们在GPU上放两层参数的话，那么它虽然占用空间会非常小，但是花费的时间会特别的长，可以说是长到不可以接受。
放两层不行，那是不是可以多放几层？来减少这一部分开销。想到的方案：假如GPU，它一共能放N层参数在上面，之前提到2层可以实现无限的调度，那么对于剩下的N-2层参数，可以从transformer中选取一些层，就固定放在GPU上（不把他们挪到CPU上）。这样的话，从的需要传输的量，在单次推理中，总的需要传输的量就少了，这些层就少了很多。换句话说，在通信上面就会快很多。
而则又引入另外一个问题，就睡说我们有两层用来调度，剩下N-2层固定，那么哪些层需要固定到GPU上，它会更好？选择哪些层被固定到GPU上的方案，可以是多种多样的。示例如下：

上图左边的方案是，6和10允许调度，789给固定到PGU上。
上图右边的方案是，6、8、10固定到GPU上。
两个方案不同点，就是两个不固定的层，它们之间的间隔，一个是间隔3层，一个是间隔1层。整体上说左边的方案不会比右边的方案差。
因为加载万layer6之后，中间有6 7 8三层计算，它们计算的时间用来加载第10层，也就是给加载第10层的时间更长了。
而右边的方案，我们中间只隔了一层，也就是加载第7层后，只有计算第八层的时间可以留给我加载第九层