【大模型LLM面试合集】分布式训练_序列并行

5.序列并行

1.序列并行（Colossal-AI）

Colossal-AI 发表的论文：Sequence Parallelism: Long Sequence Training from System Perspective， 主要是解决模型的输入长度(sequence length)限制。

Colossal-AI 序列并行诞生的背景是 self-attention 的内存需求是输入长度（sequence length）的2次方。其复杂度为 $O(n^2)$，其中，n 是序列长度。换言之，长序列数据将增加中间activation内存使用量，从而限制设备的训练能力。

而现有的工作侧重于从算法的角度降低时间和空间复杂度。因此，作者提出了序列并行，这是一种内存高效的并行方法，可以帮助我们打破输入序列长度限制，并在 GPU 上有效地训练更长的序列；同时，该方法与大多数现有的并行技术兼容（例如：数据并行、流水线并行和张量并行）。

更重要的是，不再需要单个设备来保存整个序列。 即在稀疏注意力的情况下，我们的序列并行使我们能够训练具有无限长序列的 Transformer。

具体来说，将输入序列分割成多个块，并将每个块输入到其相应的设备（即 GPU）中。为了计算注意力输出，我们将环状通信与自注意力计算相结合，并提出了环自注意力（RSA）如下图所示。

实验表明，当按批量大小和序列长度进行缩放时，序列并行表现良好。

当扩展到 64 个 NVIDIA P100 GPU 时，与张量并相比，该法分别实现了 13.7 倍和 3.0 倍的最大批量大小和序列长度。

通过稀疏注意力，序列可以处理具有超过 114K 个 Token 的序列，这比现有的在单个设备上保存整个序列的稀疏注意力运行长度超过 27 倍。

除此之外，与张量并行和流水线并行不同，序列并行不受超参数（例如： 注意力头数、层数）限制。 因此，只要序列长度能被序列并行大小整除，我们的序列并行就可以使用。

2.序列并行

Megatron-LM 发表的论文：Reducing Activation Recomputation in Large Transformer Models， 主要是减少模型显存。

Megatron-LM 的初衷是考虑通过其他方式分摊张量并行中无法分摊的显存，因此提出了序列并行的方法。

虽然 Megatron-LM 引用了 Colossal-AI 的序列并行的这篇文章，但是这两者其实并不是一个东西。

Megatron-LM 只是借用了 Colossal-AI 把 Sequence 这个维度进行平均划分的思想。在 张量的基础上，将 Transformer 层中的 LayerNorm 以及 Dropout 的输入按输入长度（Sequence Length）维度进行了切分，使得各个设备上面只需要做一部分的 Dropout 和 LayerNorm 即可。

这样做的好处有：

1. LayerNorm 和 Dropout 的计算被平摊到了各个设备上，减少了计算资源的浪费；
2. LayerNorm 和 Dropout 所产生的激活值也被平摊到了各个设备上，进一步降低了显存开销。

在 Megatron-LM 序列并行的这篇论文中，首先分析了 Transformer 模型运行时的显存占用情况。

假设输入长度为 s ，batch size为 b ，hidden dim为 h ，attention head数量为 a ，则每一层 Transformer（上图的灰色区域）的显存占用：

$$Activationsmemoryperlayer=$$