ulysses-parallel

本文最后更新于 2026年7月3日

最近重新看了一下 DeepSpeed-Ulysses，TLDR：

attention 之前，每张卡拿的是一段 sequence、完整 hidden；attention 的时候，每张卡换成完整 sequence、一部分 head。

也就是说，Ulysses Parallelism 做的事情，本质上就是在 Transformer block 内部找了一个很自然的分界点：MLP、Linear 这类 token-wise 的计算适合沿着 sequence 切；但是 Attention 需要看到完整上下文，所以在进 Attention 之前，要把数据重新洗牌成「每张卡负责一部分 head」的形态。

背景：为什么要序列并行

长上下文训练或者多模态训练里，sequence length 往往会变得很长。假设输入张量是，其中是序列长度，是 hidden size。

如果每张卡都拿完整的，那 attention 里的激活、KV、中间 attention score 都会跟着长序列一起膨胀。最直接的想法就是：既然 batch 可以切，sequence 当然也可以切。把按卡数切开，每张卡只拿个 token，这样每张卡上的激活显存自然就降下来了。

问题在于，Attention 不像 MLP。

MLP 对每个 token 是独立计算的，local token 只需要 local hidden state 就能算；但是 self-attention 里，每个 token 都要和完整序列上的 K/V 交互。如果只是朴素地把 sequence 切掉，那么每张卡只能看到局部上下文，结果肯定不对。

所以 Ulysses 要解决的问题可以概括成一句话：

如何让大部分计算保持 sequence parallel，同时在 Attention 处又能算出数值等价的 full attention？

具体介绍

Ulysses 的做法是两个 All-To-All。

第一个 All-To-All 发生在 QKV projection 之后、Attention 之前，把数据从：

1	`每张卡持有：N/P 个 token，全部 attention heads`

换成：

1	`每张卡持有：完整 N 个 token，1/P 的 attention heads`

这样每张卡就可以在自己负责的 head 上独立算 full attention。

第二个 All-To-All 发生在 Attention 输出之后，把数据再换回：

1	`每张卡持有：N/P 个 token，全部 hidden`

这样后面的、MLP、下一层 Transformer block 又可以继续按 sequence parallel 的方式执行。

注意这里有一个非常关键的点：Ulysses 不是把单个 head 内部的 head dimension 再切碎，而是按 head 切。也就是说，假设 attention head 数量是 num_heads，并行度是，通常需要满足 num_heads % P == 0。如果一个 head 自己的 hidden 维度被切掉，那的数值就不完整了。

为了讲起来简单，下面先假设正好等于 head 数量。实际工程里只要 head 数能被整除即可。

第一步：先沿 sequence 切输入

一开始，输入沿着 sequence 维度切分。每个 rank 上拿到一个 local block：

1	`Local X: [N/P, d]`

在 Ulysses 里，单独使用时参数通常没有像 Tensor Parallelism 那样按 hidden 维度切开，所以每张卡上都有完整的。于是每张卡可以直接用自己的 local token 去做 QKV projection：

1	`Local Q/K/V: [N/P, d]`

这个时候的 Local Q/K/V 有一个特点：sequence 维度是不完整的，只有；但是 hidden/head 维度是完整的，包含所有 attention heads。

到这里为止，MLP 类似的 token-wise 计算都很舒服，但 Attention 还不行。因为每个 head 需要完整序列上的 Q/K/V，当前每张卡只拿到了序列的一小段。

第二步：对 head 维度做 All-To-All

现在要做第一次数据重排。我们把每个 rank 上的 Local Q/K/V 按 head 维度切成份：

1	`[N/P, d] -> P 个 [N/P, d/P]`

然后对这些小块做 All-To-All。直观理解就是：每张卡把自己手里不同 head 的那一小段 sequence 分别发给对应的 rank；同时也从其他 rank 收到同一个 head 在其他 sequence 分片上的数据。

以 rank 0 为例，All-To-All 之前，它只有自己那段 sequence 上所有 head 的 Q/K/V；All-To-All 之后，它会收到所有 rank 发来的 head 0 分片。把这些分片沿 sequence 维度 concat 起来，就得到了：

1	`Qh/Kh/Vh: [N, d/P]`

也就是某一组 attention heads 在完整 sequence 上的 Q/K/V。

这个地方就是 Ulysses 最妙的地方。第一次 All-To-All 做完以后，数据从「按 sequence 分片」变成了「按 head 分片」。而 Attention 恰好天然可以按 head 并行，因为不同 head 之间本来就是相互独立的。

第三步：每张卡独立算 full attention

现在每个 rank 都有自己负责的 head 在完整 sequence 上的 Q/K/V，所以可以直接执行标准的 attention：

1	`Oh = softmax(Qh @ Kh.T) @ Vh`

输出形状仍然是：

1	`Oh: [N, d/P]`

这一步和普通 MHA 没有本质区别，只是每张卡只算一部分 head。换句话说，Ulysses 并没有近似 attention，也没有把 attention 改成局部 attention；它只是通过通信把数据换了个摆法，然后在每张卡上算数值等价的 full attention。

第四步：再做一次 All-To-All 换回来

Attention 算完后，每个 rank 手里是：

1	`Oh: [N, d/P]`

但后面的和 MLP 更希望看到：

1	`Local O: [N/P, d]`

也就是每张卡重新只负责一段 sequence，但是 hidden 维度完整。于是这里需要第二次 All-To-All。

这一次的通信方向可以理解为第一次的反操作：每个 rank 把自己完整 sequence 上的一部分 head，按 sequence 再切回份，然后发给对应 rank。通信结束后，同一个 sequence 分片上的不同 head 会重新聚到一张卡上。

concat 之后，每个 rank 得到：

1	`Local O: [N/P, d]`

这就回到了 Transformer block 一开始的 local sequence 形态。后面的 output projection 和 MLP 都可以直接在 local token 上继续算：

1 2	`Local O @ Wo -> [N/P, d] MLP(Local O) -> [N/P, d]`

通信量为什么是 Ulysses 的卖点

Ulysses 的核心卖点之一是通信量比较低。forward 里主要有两次通信：

Attention 前，对 Q/K/V 做一次 All-To-All，消息量大约是
Attention 后，对 O 做一次 All-To-All，消息量大约是

所以总消息量可以粗略看成。但 All-To-All 的一个重要性质是，当总消息被分散到个 rank 上时，每条链路上的通信量大约是。

这就带来一个很有意思的结论：如果 sequence length 变长，同时并行度也能跟着变大，那么单链路通信压力可以不随着线性爆炸。

相比之下，Megatron-LM 里的 Sequence Parallelism 通常要和 Tensor Parallelism 绑定，通信模式也更多依赖 All-Gather / Reduce-Scatter。Ulysses 这个「两次 All-To-All 完成形态转换」的设计，在长序列场景下就显得非常直接。

当然，漂亮归漂亮，限制也很明显。

第一个限制是 head 数。因为 Ulysses 是按 head 切的，所以并行度会受 num_heads 限制。尤其是 GQA/MQA 场景里，K/V head 数更少，这个限制会更明显。

第二个限制是 All-To-All 本身。All-To-All 对网络拓扑、NCCL 实现、跨机带宽都比较敏感。如果机器之间互联不够好，理论上的通信量优势不一定能完全变成实际速度优势。

总结

Ulysses Parallelism 可以用一句话概括：

用两次 All-To-All，在 sequence parallel 和 head parallel 之间来回切换。

进 Attention 前，把 [N/P, all heads] 变成 [N, partial heads]，让每张卡都能在自己负责的 head 上算 full attention；Attention 后，再把 [N, partial heads] 变回 [N/P, all heads]，让后面的 projection 和 MLP 继续沿 sequence 并行。

所以 Ulysses 的重点不是发明了一个新的 attention，而是非常巧妙地利用 MHA 本身「head 之间独立」这个结构，把长序列训练里最麻烦的 full context 需求变成了一次数据重排问题。

本文参考 https://zhuanlan.zhihu.com/p/5750410146

Bertrand定理下一篇