文章目录

• 简介
• 3D并行
• • 前置知识点：通信算子
  • 1. DP 数据并行
  • • 显存效率优化（ZeRO ）
    • 计算效率优化（梯度累计减少通信）
  • 2. TP tensor并行（算子内）
  • • 前置知识点：矩阵分块并行计算
    • MLP
    • self-attention
    • embedding层
    • 小节
  • 3. PP pipeline并行（算子间）
  • • 显存效率：
    • 计算效率：
• 3D并行的合并
• • 合并数据并行和流水线并行
  • 3D混合并行
• 混合精度
• • 精度选择
  • 训练
• megatron-3，节约激活值的显存
• • 序列并行sequence parallelism
  • 选择性激活重算 selective activation recomputation
  • Checkpointing Skipping
  • 各种优化措施 与 激活值显存占用

简介

Megatron-LM是一个基于PyTorch的框架，用于训练基于Transformer架构的巨型语言模型。它实现了高效的大规模语言模型训练，主要通过以下几种方式：

• 模型并行：将模型参数分散在多个GPU上，减少单个GPU的内存占用23。
• 数据并行：将数据批次分散在多个GPU上，增加训练吞吐量4。
• 混合精度：使用16位浮点数代替32位浮点数，减少内存和带宽需求，提高计算速度
• 梯度累积：在多个数据批次上累积梯度，然后再更新参数，降低通信开销
• Megatron-LM还可以与其他框架如DeepSpeed结合，实现更高级的并行技术，如ZeRO分片和管道并行5。这样可以进一步提升训练效率和规模。

超大规模语言模型训练 = GPU + PyTorch + Megatron-LM + DeepSpeed

3D并行

前置知识点：通信算子

上一步纵切并行计算后，得到 Y = [ Y 1 , Y 2 ] Y = [Y1, Y2] Y=[Y1,Y2],

• 对B纵切的话，需要通信得到完整Y。
• 对B横切的话， [ Y 1 , Y 2 ] [ B 1 , B 2 ] T = a d d ( Y 1 B 1 , Y 2 B 2 ) = Y 1 B 1 + Y 2 B 2 ， Y = D r o p o u t ( Y 1 B 1 + Y 2 B 2 ) [Y1, Y2] [B1, B2]^T = add(Y1B1, Y2B2) = Y1B1 + Y2B2， Y=Dropout(Y1B1 + Y2B2) [Y1,Y2][B1,B2]T=add(Y1B1,Y2B2)=Y1B1+Y2B2，Y=Dropout(Y1B1+Y2B2) 所以仅需Dropout前做all-reduce

所以选择对B纵切，这样整个MLP前向传播仅需一次all-reduce，在向后传递中执行一次all-reduce操作。

原mlp 的计算如下：
[ b , s , h ] × [ h , 4 h ] = [ b , s , 4 h ] [ b , s , 4 h ] × [ 4 h , h ] = [ b , s , h ] [b, s, h] \times [h, 4h] = [b, s, 4h] \\ [b, s, 4h] \times [4h, h] = [b, s, h] [b,s,h]×[h,4h]=[b,s,4h][b,s,4h]×[4h,h]=[b,s,h]

将 [h, 4h] 切分成 [h, 4h/p]，把 [4h, h] 切分成 [4h/p, h]，对[b, s, h]复制到p个卡上。计算步骤变为：
[ b , s , h ] × [ h , 4 h / p ] = [ b , s , 4 h / p ] [ b , s , 4 h / p ] × [ 4 h / p , h ] = [ b , s , h ] [b, s, h] \times [h, 4h/p] = [b, s, 4h/p]\\ [b, s, 4h/p] \times [4h/p, h] = [b, s, h] [b,s,h]×[h,4h/p]=[b,s,4h/p][b,s,4h/p]×[4h/p,h]=[b,s,h]
此有 p 个 [b, s, h]，需要做一次 allreduce 得到最终的 [b, s, h]

self-attention

（1）溢出错误
（2）因精度不足而带来舍入错误：当a与b均用FP16表示时，a=1与b=0.0001相加时，其结果是错误的，因为a/b=10000> 2 11 2^{11} 211，此时b会因尾数右移而变为0，导致结果出错。

当进行浮点数进行加减运算时，首先要使两个数的阶码相同，即小数点的位置对齐，这个过程称为对阶。在对阶时规定使小阶向大阶看齐，通过小阶的尾数算术右移来改变阶码。对阶过程中，由于FP16只有10位的尾数，当小阶的尾数右移超过11 位时，会导致该数变为0，即以FP16表示的数，如果当大数与小数的比率为大于 2 11 2^{11} 211 时，加减法运算结果会出错。而FP32有23位尾数，可表达的精度范围更广，可有效地避免该问题。

精度选择

操作f(x)满足f(x)>>x, 应在FP32下进行，例如Exp, Square, Log, Cross-entropy，因为这些操作可能会在前向或者反向计算时有结果溢出；

FP16对于某些计算是安全的，例如ReLU，Sigmoid，Tanh，可以直接在FP16的精度下进行计算；

Convolution, MatMul等操作在FP16下进行，这些操作在Volta GPU上，可以使用Tensor Core进行计算，进一步提高计算性能；

大型的规约应在FP32下计算，因为累加很多FP16数值容易造成数据溢出；

Loss计算应在FP32中进行，因为Loss计算如Softmax, cross-entropy等操作中涉及到Exp, log等对FP16不安全的操作，使用FP16时可能会导致溢出。

训练

梯度更新
问题：梯度过小产生下溢导致梯度为0；学习率过小，而导致梯度与学习率的乘积过小产生下溢；权重相较于其更新值过大（ param/update> 2 11 2^{11} 211），会因FP16精度不足，在对阶过程中使更新值变为0，导致不会更新权重。

所以需要保留FP32的权重，每次迭代时，制作这些权重的FP16副本并使用它们用于前向计算和反向计算，更新时将梯度再转换为FP32并用于更新FP32的权重。

混合精度训练过程

使用FP16的输入；

前向计算：该场景中使用统一的FP32格式的初始化权重，训练时恢复checkpoint中的变量进行训练；将恢复的FP32参数转化为FP16供给前向运算；

Loss计算：该场景中loss计算复杂，涉及到许多exp, log等可能会发生FP16溢出的不稳定操作，因此loss计算在FP32精度下进行；

loss scale：将计算出来的loss乘以一个比例因子 scaled_loss = loss * loss_scale；

反向计算：grads = tf.gradients(scaled_loss , params,aggregation_method=aggmeth)

梯度计算完成后再将各梯度缩小相应的倍数：unscaled_grads = [(grad * (1. / loss_scale)) for grad in grads]；

梯度规约：累加很多FP16数值容易造成数据溢出，选择在FP32下规约；

参数更新：根据梯度是否溢出选择是否进行参数更新并动态调整 loss_scale，若更新，直接在FP32精度的参数上进行参数更新。由于该场景中就是使用PF32的参数，更新时直接应用于该参数即可。

为保证梯度落入半精度可表示范围内一个简单有效的方法将训练损失乘以比例因子，根据梯度的链式法则使得所有梯度也等比例放大。当然在权重更新之前，需要以相同的比例因子缩小梯度，再更新到权重上。即loss scale

注意：目前使用 bf16，有更大的表示范围，不需要 loss scale 了

megatron-3，节约激活值的显存

通过序列并行sequence parallelism 和选择性激活重算 selective activation recomputation。结合张量并行，这些技术几乎消除了重新计算激活的必要性。我们在规模达到一万亿个参数的语言模型上进行了评估，并展示了我们的方法可以将激活内存降低5倍，同时将激活重算的执行时间开销降低超过90%。

序列并行sequence parallelism

激活值所占显存减少的原因，megatron-LM中的layer-norms 和 dropouts是要将输入复制到各个卡上，重复运算，也就是说激活值复制了p份，而使用序列并行化，可以仅仅保留一份激活值，而且是分布在各个gpu上。

选择性激活重算 selective activation recomputation

把Transformer族模型的所有activation消耗算了一遍，然后发现在Transformer核里有一些操作是产生的激活值又大，但是计算量又小的，这些激活值不保存，反向传播时重新计算。其他的激活值存下来，以节省重计算量。

比如下图的红框区域，具有很大的输入大小和很大的激活值，但每个输入元素的浮点运算（FLOPs）数量非常低。这部分需要重新计算而不是存下来

Checkpointing Skipping

在GPU的显存没占满的时候，可以不做checkpointing，这么一来重计算所带来的额外计算代价会进一步减小。

各种优化措施 与 激活值显存占用

tensor parallel + sequence parallel + selective activation recomputation非常节省显存

https://zhuanlan.zhihu.com/p/614166245
https://zhuanlan.zhihu.com/p/498422407
https://zhuanlan.zhihu.com/p/343570325
https://zhuanlan.zhihu.com/p/513571706
https://zhuanlan.zhihu.com/p/617087561
https://zhuanlan.zhihu.com/p/68692579
https://arxiv.org/pdf/1909.08053.pdf
https://arxiv.org/pdf/2205.05198.pdf
https://zhuanlan.zhihu.com/p/628820408