【大模型02---Megatron-LM】

---

Megatron-LM

Megatron是当前大模型训练时经常使用的一种分布式并行框架，它通过采用DP,TP,PP等来加速模型的训练，反正就是一个字，好。
大模型在训练的时候，显存占用是非常大的，比如一个175B的模型，假设模型参数用FP32表示，即4Byte,那逛模型参数就要700G，梯度700G，Adam两个动量1400G，很明显，装都装不下，所以采用这些并行技术使得它可以在有限的资源下进行训练。

补充：Adam优化器：
主要是用来动态调整学习率和梯度更新的方向，每一次更新的时候需要保存其一阶矩和二阶矩，其中一阶矩是梯度的平均值，二阶矩是梯度的平方，主要用来动态调整学习率。一阶矩和二阶矩分别表示为m和v,其相对SGD训练更加稳定。每一次更新的时候，都需要依赖上一次的计算出的一阶矩和二阶矩，也就是每一个参数都对对应一个一阶矩和二阶矩。

数据并行

假设有N张卡，每一张卡都保存一个模型的副本，现在可以将一个batch的数据分割为多个mini-batch,然后分发给每一个模型副本，进行前向传播，并计算损失和梯度，然后通过All-Reduce操作进行通信和广播，对每一个GPU计算的梯度进行规约（同步加平均），然后将梯度分发给每一个GPU，每张卡独立更新，单独更新模型参数，此时由于更新的梯度相同，模型的初始参数相同，经过更新后，每一个GPU上模型的参数也相同。
但是这种数据并行有哪些限制呢？
第一个是可以使用的GPU数量受限于batch的大小，假设batch是64，那你最多也只能用64张卡了。
另一个就是GPU利用率可能拉不满，如果batch的数量固定式512，你GPU太多，分发给每一个GPU的数据量太小，那GPU 更新块，但是通信频率也就增加了，可能会限制训练速度。

所有 GPU 必须一起等到梯度计算完成，才能开始 all-reduce，否则会造成阻塞卡顿。 数据并行最适合模型较小、计算量大、batch size 较大时使用。模型较大或 GPU 太多时要考虑混合并行或 ZeRO 分布式技术。 All-Reduce 是一种通信操作，由 NCCL / MPI / Gloo 等通信库负责实现。它不是库，而是库提供的功能

模型并行

梯度累积：
主要是用来模仿大batch进行更新的操作，因为大batch更新往往更见稳定，但是受限于显存，所以可以用梯度累积的方式，当累积到固定数量的batch之后再进行优化器更新，它通过将多个小 batch 的梯度累加，然后在累积到设定的步数后进行一次更新，从而 模拟了大 batch 的梯度平均效果。

激活检查点：
主要是用来缓解激活值对显存的占用压力，因为按照反向传播公式，每一个参数更新时，都需要前一层的激活值，这样的话，每一次更新的时候就需要存储每一个节点的激活值，对显存的占用太高，所以就采用梯度检查点的方式，每隔一定的步数保存激活值，两个激活值之间的没有保存的激活值，通过前向传播再算出来，这也就是用时间换空间了，总的来说，现在所占用的空间复杂度降为$Q(\sqrt{N})$,相当于在更新的时候再做了一次前向传播。
这也是为什么模型在模型在训练的时候，不考虑激活值占用显存的问题，因为它可以通过技术原因绕过去。

张量并行

如果一个GPU装不下一整个模型，那么就可以对模型进行拆分，相当于横着或者竖着来一刀，一般来说，对于大矩阵的乘积计算，我们可以将其分成多个小矩阵的乘积和加和，根据拆分方式的不同可以分为行并行和列并行，一般来说，列并行更好一些，因为在计算激活值的时候不需要先进行通信。通信的原因还是GELU是非线性函数，需要根据全局的信息进行计算。

行并行：

列并行：

流水线并行

流水线并行通过将模型按网络层划分为多组，每一组在一个GPU上。
目前主流的流水线并行方法包括Gpipe和PipeDream,降低空泡率。Megatron用的时Visual pipeline.1F1B,一前向一反向。
实际上流水线并行和张量并行是正交的，可以同时存在。

3D并行

3D并行就是混合数据并行DP，张量并行TP和流水线并行PP。四路张量，四路流水线，2路数据