大模型训练之序列并行双雄：DeepSpeed Ulysses和Ring-Attention

2024-05-07 14:05

©作者 | 方佳瑞

单位 | 腾讯

研究方向 | 机器学习系统

随着 Gemini 1M context length 和 Sora 出世，如何训练超长上下文的大模型引起了大家广泛关注。

本文对比两种目前炙手可热长文本训练方法 DeepSpeed Ulysess [1] 和 Ring-Attention [2]。2023 年末，二者几乎同时出现，但是设计方法大相径庭，可谓一时瑜亮。

DeepSpeed Ulysess：切分 Q、K、V 序列维度，核心卖点保持通信复杂度低，和 GPU 数无关，和序列长度呈线性关系。
Ring-Attention：切分 Q、K、V 序列维度，核心卖点是通信和计算重叠。

下面，我用 FlashAttention Style 的示意图来对比二者区别。图中 N 表示序列长度，d 表示 hidden_size=(hc * hs)，hc = head_cnt，hs=head_size，P 表示 GPU 数目（图中 P=4）。红色虚线表示通信，黑色虚线表示计算。

DS-Ulysses

DS-Ulysses 对 Q、K、V 沿着 N 维度切分成 P 份，三个分布式矩阵通过 All2All 变成沿 d 维度切分了。参见我之前的文章 [3]，All2All 等价于一个分布式转置操作。之后就是正常的 softmax(QK^T)V 计算，可以用 FlashAttention 加速，得到结果再通过 All2All 转置回来。

因为 All2All 最有通信量是 O(n)，n 是 message size，所以 DS-Ulysses 通信量位 O(Nxd)，和 P 没关系。所以可以扩展到很多 GPU 上。Ulysses 可以和 ZeRO 正交使用，ZeRO 可以进一步切分 Q、K、V，减少显存消耗。

Ulysses 也有明显缺点，就是转置后切分维度 d/P，我们希望 d/P=hc/P * head_size，即对 head_cnt 所在维度切分，这样 Attention 的计算都在一张卡上完成，从而可以使用 FlashAttention 等单卡优化。但是如果遇到 GQA 或者 MQA 情况，K、V 的 head_cnt 很小，导致 GPU 数目 P 也不能变得很大。

Ring-Attention

Ring-Attention 就是FlashAttention（FA）的分布式版本，利用了 online softmax 这个大杀器。FlashAttention 文章一搜一大把，我也解读过 [4]。这里推荐朱小霖的 Ring-Attention 文章，里面有一个非常好开源实现，在原始 Ring 基础上做了很多改进。

https://zhuanlan.zhihu.com/p/683714620

https://github.com/zhuzilin/ring-flash-attention

Ring-Attention 采用 FA 的双循环计算模式，外层循环循环遍历 Q，内层循环遍历 K、V，使用 online softmax 增量更新最终结果，这和 FA 一模一样。当 K、V 计算穿越下图虚线部分时候，需要 P2P 通信，向相邻的 GPU 卡通信。通信和计算可以重叠起来。下图只画了一个 head 的 Attention 计算，可以并行做 head_cnt 个这样的计算。

Ring-Attention 的分块大小 (下图中的参数 c) 是可调节的。红色箭头表示的 fp16 格式 KVCache 的 P2P 通信量是 bytes。QKV 分块计算量是 FLOPS。所以只要满足计算量大于通信量，计算通信可以重叠起来，从而让通信开销消失。另外 K、V 的计算结果 intermediate tensor 只需要 c x c 大小部分，内存消耗很少。

Ring-Attention 也有很多缺陷。比如 Self-Atention 计算有效部分一般是一个下三角，所以均匀切分 Q 的话，负载是不均衡的。这个问题 @朱小霖的实现做了优化。另外处理变长序列也不容易，这在 SFT 任务中比较常见。

二者比较

通信量：Ulysses 完胜。

DS-Ulysses 三次 All2All 通信量 3xO(Nxd)。
Ring-Attention ：N/P/c x (P-1)/PxO(Nxd)=O(N^2xd/(Pxc))，外层循环每个GPU需要N/P/c次迭代，内层循环每个GPU收发(P-1)/P x O(Nxd)数据。通信会随着序列长度增长而平方增长。所以非常依赖和计算重叠。

通信方式：Ring-Attention 更鲁棒。