Context Parallel 与 Ulysses：序列切到多卡后，
attention 怎么算？两者通信差在哪？
先弄清楚：每个概念是什么
- Context parallel（CP）
- activation 随 sequence length \(s\) 线性涨，所以把所有 activation 沿 sequence 切成 cp 份，每卡只管一段 token。LayerNorm、MLP 这些逐 token 的算子各算各的，只有 attention 需要看别人的 token。例：s = 128K、cp = 8 → 每卡 16K 个 token。
- Ring attention（Megatron CP 默认）
- 每卡的 Q 块不动，K、V 块沿环传。每一步用自己的 Q 和手上这块 KV 算一次，同时把它发给下一张卡，cp 步后每卡都见过全部 KV。
- Online softmax 合并
- 每块 KV 算出局部输出 \(o_j\) 和 log-sum-exp \(\mathrm{lse}_j\)，按 LSE 重新加权相加：\(\mathrm{lse} = \log(e^{\mathrm{lse}_1} + e^{\mathrm{lse}_2})\)，\(o = e^{\mathrm{lse}_1-\mathrm{lse}}\,o_1 + e^{\mathrm{lse}_2-\mathrm{lse}}\,o_2\)。结果和一次看完整个序列一样，同 FlashAttention 的分块。
- DeepSpeed-Ulysses
- 对 Q、K、V 做 all-to-all，把"按 sequence 切"换成"按 head 切"。每卡拿到完整序列、\(a/\text{cp}\) 个 head，本地算普通 attention，再 all-to-all 换回按 sequence 切。例：32 个 head、cp = 8 → 每卡 4 个 head × 128K 个 token。
看图：ring attention 怎么转，causal 怎么配平
30 秒回答
Activation memory grows with sequence length, so context parallelism shards every activation along the sequence across cp GPUs. Only attention needs other GPUs' tokens. Ring attention, Megatron's default, keeps each query block local and passes KV blocks around a ring with point-to-point sends, overlapping the transfer with compute and merging partial results with online softmax. DeepSpeed-Ulysses instead uses all-to-all to switch from sequence sharding to head sharding, computes full-sequence attention on a/cp heads per GPU, then switches back. Ulysses moves less data per GPU, about 4sh/cp versus 2sh for ring with standard multi-head attention, but attention has to wait for the all-to-all, and cp cannot exceed the head count.
推导 / 关键公式
每卡每层通信（forward，元素数）
Ring：收 \(\text{cp}-1\) 块 KV，每块 \(2\,(s/\text{cp})\,h_{kv}\)；\(h_{kv}\) 是 KV head 数 × head dim：
\[ V_{\text{ring}} = 2\,s\,h_{kv}\cdot\frac{\text{cp}-1}{\text{cp}} \]Ulysses：Q、K、V、O 各做一次 all-to-all，每卡把本地数据的 \((\text{cp}-1)/\text{cp}\) 发出去：
\[ V_{\text{Ulysses}} = \frac{s}{\text{cp}}\,(2h + 2h_{kv})\cdot\frac{\text{cp}-1}{\text{cp}} \]MHA（\(h_{kv} = h\)）时 ring \(\approx 2sh\)，不随 cp 变小；Ulysses \(\approx 4sh/\text{cp}\)。
为什么要一头一尾配对
causal mask 下越靠后的 token 要看的 KV 越多。cp = 4 按顺序切，计算量是 0.5 : 1.5 : 2.5 : 3.5（对角块只算一半），最慢的卡是平均的 1.75 倍。切成 2·cp = 8 块、卡 \(i\) 拿第 \(i\) 和第 \(7-i\) 块后，每卡都是 \((i + 0.5) + (7.5 - i) = 8\) 个小块。
对比与手算例子
\(s = 128\text{K}\)、\(h = 4096\)、cp = 8、bf16，每卡每层 forward：
| Ring（p2p） | Ulysses（all-to-all） | |
|---|---|---|
| MHA | 1792 MiB | 448 MiB |
| GQA，8 个 KV head（\(h_{kv}/h = 1/4\)） | 448 MiB | 280 MiB |
| 能否 overlap | 能：算当前块时传下一块 | 难：attention 要等 all-to-all 完成 |
| 限制 | 每卡 chunk 要够长 | cp ≤ head 数，all-to-all 要高 bandwidth |
必答点
- 为什么：activation \(\propto s\)；CP 沿 sequence 切全部 activation，只有 attention 要跨卡
- Ring：Q 不动、KV 绕环 p2p，online softmax（LSE）合并，通信和计算 overlap
- Ulysses：两次 all-to-all，sequence 切 ↔ head 切；cp ≤ head 数
- 通信：ring \(\approx 2sh_{kv}\) 不随 cp 降；Ulysses \(\approx (2h+2h_{kv})\,s/\text{cp}\)
- causal：切 2·cp 块，卡 \(i\) 拿 \(i\) 和 \(2\text{cp}-1-i\)
追问
CP 和 Megatron SP 有什么区别？
SP 只切 TP 管不到的 LayerNorm、dropout 的 activation，依附 TP。CP 切全部 activation（包括 attention），可以和 TP + SP 同时开。
Ring 的通信什么时候能被完全盖住？
一块的计算 \(4hc^2/F\) 要不小于传输 \(4ch/B\)（bf16 的 K、V 块），即每卡 block 长度 \(c \ge F/B\)（\(F\) 是 FLOPS，\(B\) 是 bandwidth）：A100 走 NVLink 约 1K token，走 InfiniBand 约 24.5K。
Megatron 里两种能一起用吗？
能。--cp-comm-type 有 p2p（ring）、a2a（Ulysses 式）、allgather（先 all-gather 全部 KV，不能 overlap）和 a2a+p2p：机内用 a2a 走 NVLink，机间用 p2p 走 InfiniBand。
来源Megatron Core: Context Parallelism · Liu et al. 2023（Ring Attention）§3、Table 2（pp.3–5）、§6（p.9）· Jacobs et al. 2023（DeepSpeed-Ulysses）§1（p.2）、§3.1–3.2（pp.4–5）· Transformer Engine cp_comm_type docstring、context_parallel.py（2·cp 块分配）· Megatron-LM arguments.py · DeepSpeed sequence/layer.py（head 数整除检查）