序列并行

SP

Sam Ade Jacobs et al., DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models, 2023.

这篇处理的是 SP，即 Sequence parallelism。

DP 将一个 batch 中的不同 samples 分给不同 GPU 并行处理，但每张 GPU 仍然需要处理完整的单条 sample。对于长序列训练，即使 micro-batch size 已经是 1，单条 sample 仍可能因为 sequence length 太长而无法高效处理。

SP 针对的就是长序列样本的并发。很自然地，只需要把长序列切开平分给每张 GPU 不就行了嘛！假如说有四张 GPU，一条 128K token 的长序列，每张分别负责 32K token。但一个问题是，global attention 这个操作要求每个 token 都 attend 整个序列中的其他 token。

DeepSpeed-Ulysses 通过两次 All-to-All，既保持了 sequence sharding，又让分布在不同 GPU 上的序列 token 做完整的 global attention。

第一次 All-to-All 很巧妙，它把数据按 sequence 切分转换成按注意力头切分。比如说，原来 GPU0 负责 0-32K token 的 H0-3 四个注意力头的计算，转换成负责 0-128K token 的 H0 单个注意力头的计算；其他三张 GPU 也做类似的转换。注意到转换前后每张 GPU 持有的 QKV 数据量没有变化！既然每张 GPU 现在都有了整条序列，而注意力头之间的计算是独立的，直接算 attention 就可以了。第二次 All-to-All 就把数据交换回来，恢复原来的 sequence partition，以便 output projection、MLP、LayerNorm 这些操作能正常地按照 SP 进行。

之前的 SP 方法比如 Megatron-SP，主要使用 All-Gather/Reduce-Scatter 在 SP 和 TP 之间转换，通信量随序列长度线性增长。DeepSpeed-Ulysses 则利用 sequence-to-head 的 trick，将通信量降低到 $O(Nh/P)$，其中 $h$ 是 hidden size，$P$ 是 SP degree。所以如果序列长度和 GPU 数量同比例增加，通信量就可以保持近似不变，从而实现更好的 scalability。