DeepSpeed-Ulysses

attention 前后各一次 A2A 怎么转置张量维度、为何并行度受 head 数约束、USP 怎么混合

核心要点：

• attention 前后各一次 a2a，把张量在序列切分↔head 切分间转置
• 本地 attention 阶段零通信（每卡持完整序列、部分 head）
• per-rank 通信 $O(S \cdot d / N)$，随 $N$ 减少
• 并行度上限 $N \le h$（head 数约束）
• USP 混合：group 内 Ulysses,group 间 Ring

名词定义

名词	定义
Head（注意力头）	多头注意力把 hidden 维切成 $h$ 个独立子空间，每个子空间各算一份 attention，即一个 head；head 间相互独立是 Ulysses 沿 head 维切分的前提。详见 多头注意力。注：本篇 $h$ 指 head 数（该篇 $h$ 指 hidden 维、head 数记为 $n_{\text{head}}$）
Head 切分	attention 计算时把 $h$ 个 head 分到 $N$ 卡，每卡持 $h/N$ 个 head 的完整序列
Pre/Post-attention a2a	Ulysses 在 attention 入口/出口各一次 All-to-All，在序列切分与 head 切分间转置

CP / Ulysses / USP 等共享名词见 7.1 总览 名词定义。

Ulysses 怎么用 a2a 替代 ring?

Ulysses 不让 K/V 在 rank 间流动，而是用 a2a 在 attention 入口把张量从「序列切分」重排为「head 切分」，attention 内部回到 dense 全局计算，末尾再 a2a 重排回序列切分[1]。

设 CP 度 $N$，序列长 $S$, head 数 $h$，每 head 维度 $d_{\text{head}}$。各 rank 初始持 $[b, S/N, h, d_{\text{head}}]$ 的 Q/K/V （沿序列切分），算法 4 步：

1. Pre-attention a2a：沿 head 维 a2a，把张量从 $[b, S/N, h, d_{\text{head}}]$ 重排为 $[b, S, h/N, d_{\text{head}}]$ — 每 rank 持完整序列长度的 $h/N$ 个 head
2. 本地 attention：每 rank 在分到的 $h/N$ 个 head 上做 dense 全局 attention （query 看完整序列，无需通信）
3. Post-attention a2a：沿 head 维再做一次 a2a，把输出重排回 $[b, S/N, h, d_{\text{head}}]$
4. FFN / output projection：在序列切分布局下继续后续层

attention 模块只有 2 次 a2a (pre + post)，不再有 ring 旋转。

Ulysses 每卡通信多少？为什么随 N 减少？

Ulysses 的 per-rank 通信复杂度 $O(S \cdot d_{\text{model}} / N)$，随 $N$ 减少。每 rank 本地数据量 $M_{\text{local}} = b \cdot (S/N) \cdot h \cdot d_{\text{head}} \cdot s_{\text{dtype}}$，单次 a2a 中每 rank 发送 $(N-1)/N$ 份本地数据：

$$\begin{equation} M_{\text{ulysses-a2a}}^{\text{per-rank}} = \frac{N-1}{N^2} \cdot b \cdot S \cdot h \cdot d_{\text{head}} \cdot s_{\text{dtype}} \label{eq:par-cp-ulysses-comm} \end{equation}$$

两次 a2a 合计 per-rank 通信 $\approx 2 b \cdot S \cdot d_{\text{model}} / N \cdot s_{\text{dtype}}$。这与 Ring 形成对照：

• Ring per-rank: $O(S \cdot d_{\text{model}})$，不随 $N$ 变 (见 7.2 Ring Attention)
• Ulysses per-rank: $O(S \cdot d_{\text{model}} / N)$，随 $N$ 减少

通信结构上 a2a 是 1 步内全 rank 同时收发，带宽需求集中；ring 是 $N-1$ 步顺序，每步低带宽但延迟累积。

Ulysses 有什么限制？

Ulysses 的 CP 度被 head 数卡住，还有三条约束：

• head 数约束：$N \leq h$。step 1 需把 $h$ 个 head 均分到 $N$ rank，否则不能整除。GQA / MQA 把 KV head 压到 8 时，Ulysses 的 CP 度被限制
• Causal mask 自然支持：attention 内部是单 rank 完整序列 dense 计算，causal 退化为本地三角 mask
• 互联拓扑要求：a2a 在 $N$ rank 间形成全连接，对 fat-tree / NVSwitch 友好，对纯 ring 拓扑较差

原论文在 64 GPU 上比当时 SOTA 快 2.5×，支持 4× 长度[1]。

Ring 还是 Ulysses？能组合吗？

短序列 / head 少 / ring 互联用 Ring，长序列 / head 多 / 全连接互联用 Ulysses；两条路线可组合。

维度	Ring Attention	DeepSpeed-Ulysses
通信原语	P2P ring (顺序 $N-1$ 步)	a2a （attention 前后各 1 步）
Per-rank 总通信复杂度	$O(S \cdot d)$ 不随 $N$ 变	$O(S \cdot d / N)$ 随 $N$ 减少
并行度上限	无固有限制	$N \leq h$ （head 数约束）
Attention 计算分布	每 rank 算 $(S/N) \times S$ 子块	每 rank 算完整 $S \times S$，仅 $h/N$ 个 head
Causal mask 处理	需 balanced split 或 striped	单 rank 完整 attention，自然处理
拓扑友好度	Ring / torus / NVSwitch 都 OK	Fat-tree / NVSwitch；纯 ring 较差
与 TP 关系	维度正交	与 TP 沿 head 维竞争

@tbl-par-cp-uly-vs-ring Ring Attention 与 DeepSpeed-Ulysses 对比

USP (Unified Sequence Parallelism) 把两者结合[2]：在小 group 内用 Ulysses （head 切分），group 间用 Ring （序列切分），同时绕开 head 数上限和 ring 延迟累积。这在分层拓扑上贴合度高：group 内走高带宽域 (NVLink) 的 a2a, group 间走低带宽域容忍的 ring。

Ring vs Ulysses 定量性能对比

USP 论文在 8×A100 (NVLink) 和 8×L20 (PCIe) 上做了系统实测[2]:

平台	序列长度	纯 Ulysses	纯 Ring	USP 混合 （最优配置）
A100 NVLink	32K	136.4 iters/s	~132	133.0 (Ulysses=4)
A100 NVLink	128K	2.79 iters/s	—	2.77 (Ulysses=4)
L20 PCIe	32K	14.3	40.5	62.8 (Ulysses=2, lb-ring)
L20 PCIe	128K	2.59	3.40	5.48 (Ulysses=2, lb-ring)

@tbl-par-cp-uly-bench USP Ring vs Ulysses 实测 （8 GPU 单节点）

解读：

• NVLink 上 Ulysses 赢：All-to-All 利用全对分带宽，Ring 的 overlap 反而降低计算密度
• PCIe 上 Ring+混合赢：PCIe 带宽不足以支撑 a2a, Ring 的 P2P 通信量恒定更优
• 互联带宽比 (~2.5× NVLink vs IB) 直接决定 CP 策略选择

LoongTrain 在 64 GPU/MHA/1M 场景下实现 vs Megatron-CP 2.88× MFU、vs Ulysses 1.49× 加速[3]。

新方向：Untied Ulysses (ICML 2026) 把分片粒度推到 head 级流水线 (headwise chunking)，中间 tensor 内存降低 87.5%，单 8×H100 节点支持 5M tokens[4]。

除了 USP，还有哪些多维 / 分层 ring 在降通信？

USP 与 LoongTrain 靠「混入 a2a 或加并行维度」降 ring 开销；另一条路线保留 ring 结构，在环内部减少通信量或提高链路利用率：BurstAttention 优化 backward, WallFacer 压跨节点通信占比，TokenRing 提链路利用率。

BurstAttention 用两级结构降通信[5]。集群级 (GAO) 用 online softmax 消除中间矩阵的存储与传输，设备级 (LAO) 把中间矩阵 tile 化留在 SRAM 避免 HBM 往返，double buffer 让计算与下一轮 K/V 传输重叠。它最突出的优势在反向：复用前向保存的 log-sum-exp, backward 只需传 Q / dQ / dO 三份，而 Ring Attention 不复用需传约六份，通信量差约 2×。128K / 32×A100 实测比 TP+FlashAttention 快 1.37×、通信省 40%，延伸工作 BurstEngine 把这套推到 1M+ token[5]。

WallFacer 把 attention 看作 n-body 相互作用，引入一个 team 维度 $C$ 把全局 ring 拆成 $P/C$ 个子 ring 并行[6]。子 ring 的 P2P 限制在节点内 （高带宽），跨节点通信只在初始化阶段用一次 collective 完成，迭代轮数从 $P-1$ 降到约 $P/C^2-1$。算例 ($C=4$, $P=64$, LLaMA-30B) 通信量从 1.625 GB 降到 0.558 GB （省约 66%）；512K / 32×A100 实测 GPT 吞吐提升 77%、DiT 提升 114%。

TokenRing 针对链路利用率[7]。随序列变长，计算量按 $N^{-2}$ 下降而通信量按 $N^{-1}$ 下降，单向 ring 让另一半链路闲置。TokenRing 在 full-mesh 拓扑上双向传输——query 顺时针、block 中间输出 (output + lse) 逆时针——两类数据并发并与计算重叠，4×A10 / 24K 序列延迟约 3.5 ms、Ring Attention 约 7.6 ms （降幅约 54%）。收益依赖 full-mesh 互联，受限拓扑下部分折损。

方案	切入点	关键机制	实测要点
USP / LoongTrain	换原语 / 加维度	group 内 a2a + group 间 ring （见上文）	LoongTrain vs Megatron-CP 2.88×
BurstAttention	优化 backward	两级 GAO+LAO，复用 lse 省反向通信	反向通信约 Ring 的 1/2
WallFacer	压跨节点占比	n-body 子 ring,P2P 限节点内	通信量省约 66%
TokenRing	提链路利用率	双向 ring,full-mesh 双向并发	24K 延迟降约 54%

@tbl-par-cp-uly-multidim 多维/分层 ring 变体的切入点对比

这些工作与 USP / LoongTrain 正交：后者换通信原语或加并行维度，前者在保留 ring 的前提下优化某一环节，可叠加使用。

Takeaway

知识点	核心结论
a2a 转置	attention 前后各 1 次 a2a，序列切分↔head 切分
本地 attention	每卡持完整序列、$h/N$ head,attention 内零通信
通信复杂度	$O(S \cdot d / N)$，随 $N$ 减少（与 Ring 不变相反）
head 数上限	$N \le h$,GQA/MQA 下受限
USP 混合	group 内 Ulysses + group 间 Ring，贴合分层拓扑
多维 ring	BurstAttention/WallFacer/TokenRing 在 ring 内部降通信，与 USP 正交

@tbl-par-cp-uly-takeaway Ulysses 与混合并行核心知识点

参考资料

1. Jacobs et al., DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models, arXiv:2309.14509, 2023. https://arxiv.org/abs/2309.14509
2. Fang et al., USP: A Unified Sequence Parallelism Approach for Long Context Generative AI, arXiv:2405.07719, 2024. https://arxiv.org/abs/2405.07719
3. LoongTrain: Efficiently Training Long Sequence LLMs with Head-Context Parallelism, arXiv:2406.18485, 2024. https://arxiv.org/abs/2406.18485
4. Ghadia et al., Untied Ulysses: Rethinking Sequence Parallelism with Headwise Chunking, ICML 2026, arXiv:2602.21196. https://arxiv.org/abs/2602.21196
5. Sun et al., BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences, arXiv:2403.09347, 2024. https://arxiv.org/abs/2403.09347
6. WallFacer: Guiding Transformer Model Training Out of the Long-Context Dark Forest with N-body Problem, arXiv:2407.00611v2, 2024 （v3 后更名为 StarTrail: Concentric Ring Sequence Parallelism）. https://arxiv.org/abs/2407.00611v2
7. TokenRing: An Efficient Parallelism Framework for Infinite-Context LLMs via Bidirectional Communication, arXiv:2412.20501, 2024. https://arxiv.org/abs/2412.20501