算法优化与并行注意力机制

算法优化是人工智能领域中的重要议题。本文将介绍并行注意力机制、滑动窗口Attention、AdamW优化、3D并行优化、数据加载优化、网络通信优化和集群容错等方面的优化方法。

并行注意力机制

在并行注意力机制中，通过将mlp融合到attention里，可以实现在62B模型上性能没有下降。主要对应的是网络结构的并行化改造。

滑动窗口Attention

滑动窗口Attention的原理是通过堆叠不同大小的窗口来捕获句子中的信息，从而减小所需的计算量。随着层级越高，理论上每个位置注意到的区域就越大，所能存储的信息就越接近全局attention时的状态。

AdamW优化(LAMB)

AdamW相比于adam的优化在于将权重衰减项从梯度的计算中拿出来直接加在了最后的权重更新步骤上，从而提升优化效果。新增的截断函数的主要目的是为了防止batch_size太大的时候导致优化过程中动量出现极端值影响bp。

这里做的优化是新增了一个 \(\phi\) 截断函数，主要目的是为了防止batch_size太大的时候导致优化过程中动量出现极端值影响bp。这个方法可以将batch_size增大4倍从而加速训练。

3D并行优化

张量并行优化

序列并行(SP)主要有两个目的：平摊LayerNorm和Dropout的计算开销，以及平摊Activation占用显存。另外，对矩阵做切片后再进行矩阵乘法，计算效率要也比两个超大的矩阵乘法要高。

AllGather优化

序列并行(SP)后，在进行张量并行(TP)前需要在fp的时候需要先通过gather把之前层的切片从其他节点copy汇聚过来。为了尽量隐藏通信，可以优化成每通信完成一个切片后，进行这个切片的MLP列切分计算，同时直接把gather结果送给attention并行计算，最后再把切片计算结果concat到一起。

Reduce-Scatter优化

这块是需要把汇聚计算完成的tensor在重新进行切分发送到序列并行的节点里，这里是把MLP的第二次行切分和attention结果加和给merge到了一起，完成一个切片的计算后就发送出去，同步进行下一个切片的计算使计算和通信异步进行。

流水线优化

在warm-up/cool-down过程里，都是必须等通信完成才能进行计算的。为了缩短等待时间，可以将allGather的recv/send拆分开，recv优先级高于send，recv后就能直接开始计算，不需要等send的长尾，从而缩短等待时间。

在稳定状态的时候应该和megatron一样，通信都会和计算异步。实际情况里通信一般都会被隐藏掉。

数据加载优化

数据加载优化的主要思想包括在bp完同步梯度的时候，所有前向相关的数据就没用了，就可以直接释放回池预加载下一轮fp需要的embed。另外，避免单机内多张卡重复读相同的冗余数据，可以先在内存里去好重再copy到显存。

网络通信优化

网络通信优化的内容待补充。

集群容错

错误检测

错误检测的主要思想包括每个worker定期上报心跳给中心节点，确保当前状态正常；状态异常时的自动化诊断，完成诊断后上报给中心节点；中心节点向k8s申请失败节点的拉黑和重分配替换。

状态恢复

状态恢复的主要优化手段包括checkpoint保存和checkpoint读取。在同一数据并行组里的卡，只选一个GPU对应的训练线程读hdfs后写内存，然后通过broadcast给这个数据并行组里的其他卡，可以降低hdfs的读取压力。

状态监控

状态监控主要基于cuda_event的timeline可视化，根据DP来画出卡和卡的数据流依赖关系。

参考

megascale: https://arxiv.org/abs/2402.15627

Palm(并行attention): https://public.agent-matrix.com/publish/shared/Paper/Palm.pdf

滑动窗口注意力解释: https://zhuanlan.zhihu.com/p/223430086