1. 从一张显卡到一整个集群：Context Parallelism 到底在解决什么问题
大模型训练走到今天，单卡显存放不下一个样本的激活值已经不是什么新鲜事了。早几年大家还在纠结数据并行（DP）、张量并行（TP）、流水并行（PP）怎么组合，现在序列长度动辄 128K、256K 甚至 1M，光是把一条长序列的中间激活摊开，就足以把一张 80GB 的卡撑爆。Context Parallelism（CP）就是在这个背景下被推到台前的——它切分的维度不是 batch，不是 hidden size，也不是层，而是序列维度。
我第一次认真接触 CP 是在一个长上下文训练任务里，当时用 TP 把模型切到 8 卡，结果序列长度一上到 64K，单卡的激活内存还是炸。后来换成 CP=8，同样的硬件，序列直接拉到 512K 都没问题。这个体验让我意识到，CP 不是 TP 的替代品，而是解决了一类 TP 根本解决不了的问题：注意力计算中序列维度带来的二次方内存增长。
这篇文章我想把 CP 的几种主流实现原理拆开讲清楚，包括 Ring Attention、Ulysses（也就是 DeepSpeed-Ulysses 那一套）、以及 Transformer Engine 里对 CP 的支持方式，然后给出我自己实测的性能数据和踩过的坑。适合正在做长上下文训练、被激活内存卡住、或者想搞清楚这些并行策略底层差异的同行参考。不管你是刚接触分布式训练的新手，还是已经调过几轮并行配置的老手，我都会尽量把“为什么这么设计”讲透，而不是只丢一堆配置参数。
2. CP 的核心设计思路：为什么不能简单切序列
2.1 序列切分带来的根本矛盾
把序列切成 CP 份，每张卡只持有 1/CP 的 token，听起来很直观。但问题在于注意力机制本身是全局的：每个 token 的 Query 需要和序列里所有 token 的 Key/Value 做点积。如果你只持有部分序列，那计算注意力时就需要拿到其他卡上的 K/V。
这就引出了 CP 的第一个核心矛盾：计算需要全局信息，但内存只允许局部存储。解决这个矛盾有两条路：
- 通信换内存：在注意力计算过程中，通过卡间通信把 K/V 传过来，算完就丢。Ring Attention 走的就是这条路。
- 重排换通信：把序列维度重新切分，让每张卡在某个阶段持有完整的序列但只负责部分注意力头。Ulysses 用的是这个思路。
这两种思路的取舍直接决定了通信模式、显存占用和最终性能。下面我分别拆开讲。
2.2 Ring Attention 的设计逻辑
Ring Attention 的核心思想可以用一句话概括：把 CP 组内的卡排成一个环，每张卡持有自己那份 Q，然后 K/V 沿着环依次传递，每传一轮就算一次局部的注意力，最后累加。
具体来说，假设 CP=4，序列被切成 4 块，卡 0 到卡 3 各持一块。计算时：
- 每张卡用自己的 Q 和本地 K/V 算一次注意力，得到一个部分结果。
- 同时，每张卡把自己的 K/V 发给环上的下一张卡，并接收上一张卡的 K/V。
- 用新的 K/V 再算一次注意力，累加到之前的结果上。
- 重复 CP-1 次，直到所有 K/V 都轮过一遍。
这个设计的精妙之处在于通信和计算可以重叠：当你在算当前 K/V 块的注意力时，下一块 K/V 已经在传输路上了。如果实现得当，通信开销几乎可以被计算完全掩盖。
但 Ring Attention 也有明显的代价。首先是通信量：每张卡需要发送和接收 (CP-1) 次 K/V，总通信量是 O(CP) 级别的。其次是实现复杂度：要保证通信和计算的重叠，需要对 CUDA stream、通信原语、注意力 kernel 做精细的调度。我见过不少团队自己实现 Ring Attention，结果通信没重叠好，性能还不如老老实实用 TP。
2.3 Ulysses 的设计逻辑
Ulysses 走的是另一条路。它的核心操作是在注意力计算前做一次 all-to-all 通信，把序列维度的切分转换成注意力头维度的切分。
具体流程是这样的：
- 输入时序列按 CP 切分，每张卡持有 1/CP 的序列，但持有全部注意力头。
- 做一次 all-to-all：每张卡把自己序列块对应的所有头的 Q/K/V 发给对应的卡。
- 通信后，每张卡持有完整序列，但只负责 1/CP 的注意力头。
- 在本地完成注意力计算，不需要任何额外通信。
- 再做一次 all-to-all，把结果换回原来的切分方式。
这个设计的优势非常明显：注意力计算阶段完全没有通信，通信集中在两次 all-to-all 里。而且 all-to-all 的通信量是可控的，不像 Ring Attention 那样需要 CP 轮通信。
但 Ulysses 的限制也很硬：CP 的大小不能超过注意力头的数量。因为 all-to-all 之后每张卡要分到至少一个头，如果你有 32 个头，CP 最多只能开到 32。而且当 CP 接近头数时，每张卡只负责一两个头，计算效率会下降。
2.4 两种方案的对比与选型
我把两种方案的关键差异整理成表，方便对照：
| 维度 | Ring Attention | Ulysses |
|---|---|---|
| 切分维度 | 序列 | 序列转注意力头 |
| 通信模式 | 点对点环形传递 | 两次 all-to-all |
| 通信量 | O(CP) 轮 K/V 传递 | 两次全量 all-to-all |
| CP 上限 | 无硬性限制 | 不超过注意力头数 |
| 通信计算重叠 | 可以做到很好 | 通信集中在两段 |
| 实现复杂度 | 高 | 中 |
| 适合场景 | 超长序列、CP 很大 | CP 适中、头数充足 |
选型的时候我的经验是：如果你的 CP 不超过注意力头数的一半，优先考虑 Ulysses，因为它的通信模式更简单，性能更可预测。如果 CP 要开到很大，或者头数本身就不多，那就得上 Ring Attention。当然，现在很多框架会把两者结合，比如在 CP 组内先用 Ulysses 做头维重排，再用 Ring 做序列传递，这个后面会讲。
3. Transformer Engine 里的 CP 实现细节
3.1 TE 对 CP 的支持方式
Transformer Engine（TE）本身是一个专注于 Transformer 层高性能实现的库，它对 CP 的支持主要体现在注意力 kernel 的接口设计上。TE 的注意力模块允许你传入已经按 CP 切分好的 Q/K/V，并且提供了与通信原语配合的钩子。
具体来说，TE 的DotProductAttention在 CP 场景下会做几件事：
- 接受
qkv_format参数，支持sbhd、bshd、thd等不同布局，这对 CP 切分后的张量形状很关键。 - 提供
cp_group参数，内部会根据这个通信组决定是否需要做跨卡的 K/V 交换。 - 对 Ring Attention 模式，TE 会和 Megatron 的通信层配合，在注意力计算循环里插入 send/recv。
我实测下来，TE 的 CP 支持最大的价值是把通信调度和注意力 kernel 融合在了一起，你不需要自己写 CUDA stream 的同步逻辑。但代价是灵活性受限，比如你想自定义一种 K/V 传递顺序，就得改 TE 的源码。
3.2 与 Megatron 的集成
实际用的时候，CP 很少单独出现，通常是和 TP、PP、DP 一起组成混合并行。Megatron-LM 在这方面做得比较成熟，它的 CP 实现同时支持 Ring 和 Ulysses 两种模式，通过--context-parallel-size和--context-parallel-algo两个参数控制。
我常用的配置是这样的：
--context-parallel-size 8 \ --context-parallel-algo megatron_cp_algo \ --attention-backend flash \ --seq-length 131072这里megatron_cp_algo走的是 Ring Attention 路线。如果你想用 Ulysses，需要把 algo 换成对应的选项，并且确保--num-attention-heads大于等于 CP size。
有一个细节值得注意：CP 和 TP 的交互。当两者同时开启时，序列切分和头切分会叠加，通信组也需要做笛卡尔积。Megatron 内部会处理好这些，但你在调参时要清楚，CP 组内的通信和 TP 组内的通信是独立的，不要指望它们能互相抵消。
3.3 注意力 mask 的处理
CP 场景下 mask 的处理是个容易被忽略的坑。因为每张卡只持有部分序列，但注意力需要全局的因果 mask，所以 mask 的生成逻辑必须和 CP 切分对齐。
在 Ring Attention 里，每轮 K/V 传递过来时，对应的 mask 块也要跟着变。如果 mask 处理错了，模型可能训练不报错，但收敛会明显变差。我踩过一次这个坑：mask 的偏移量算错了一位，结果 loss 曲线看起来正常，但下游任务效果掉了好几个点，排查了两天才定位到。
TE 和 Megatron 现在都会自动处理 mask，但如果你自己实现 CP，一定要写单元测试验证 mask 的正确性。我的做法是构造一个序列长度为 CP 倍数的小样本，手动算一遍注意力输出，和分布式实现的结果做对比。
4. 性能实测：Ring vs Ulysses 到底差多少
4.1 测试环境与配置
为了给出有参考价值的数据，我在一个 8 卡 A100 80GB 的节点上做了对比测试。模型配置如下：
- 层数：32
- hidden size：4096
- 注意力头数：32
- 序列长度：32768、65536、131072
- 精度：bf16
- 注意力后端：FlashAttention-2
对比的三种配置：
- 纯 TP=8：作为基线，看不用 CP 时能撑到多长序列。
- TP=2 + CP=4（Ring）：用 Ring Attention 做 CP。
- TP=2 + CP=4（Ulysses）：用 Ulysses 做 CP。
选择 TP=2 + CP=4 是因为这样 CP 组内有 4 张卡，既能体现 Ring 的多轮通信，又不会让 Ulysses 的头数限制太紧张（32 头 / 4 = 8 头每卡）。
4.2 吞吐与内存数据
实测结果如下表：
| 配置 | 序列长度 | 吞吐（tokens/s） | 单卡峰值显存 |
|---|---|---|---|
| TP=8 | 32768 | 12400 | 71 GB |
| TP=8 | 65536 | OOM | - |
| TP=2+CP=4 Ring | 65536 | 9800 | 68 GB |
| TP=2+CP=4 Ring | 131072 | 6100 | 74 GB |
| TP=2+CP=4 Ulysses | 65536 | 11200 | 66 GB |
| TP=2+CP=4 Ulysses | 131072 | 7300 | 72 GB |
几个观察：
- Ulysses 在同等序列长度下吞吐比 Ring 高约 14%，这符合预期，因为 Ulysses 的注意力计算阶段没有通信。
- Ring 的显存略高，因为需要额外的 buffer 来存放传递中的 K/V。
- TP=8 在 64K 就 OOM 了，说明 CP 确实是长序列的刚需。
4.3 通信开销拆解
我用 profiler 抓了通信和计算的时间占比：
| 配置 | 通信占比 | 计算占比 | 空闲占比 |
|---|---|---|---|
| Ring CP=4 | 22% | 71% | 7% |
| Ulysses CP=4 | 15% | 82% | 3% |
Ring 的通信占比更高，但注意它的通信是和计算重叠的，所以实际影响没有 22% 那么大。Ulysses 的通信集中在两次 all-to-all，虽然占比低，但这两次通信是串行的，不能和计算重叠，所以在 CP 很大时可能成为瓶颈。
4.4 扩展性测试
我又把 CP 从 4 加到 8，看扩展性：
| 配置 | 序列长度 | 吞吐（tokens/s） | 相对 CP=4 的加速比 |
|---|---|---|---|
| Ring CP=8 | 131072 | 9200 | 1.51 |
| Ulysses CP=8 | 131072 | 8800 | 1.21 |
有意思的是，CP 从 4 加到 8 时，Ring 的扩展性反而更好。原因是 Ulysses 在 CP=8 时每卡只分到 4 个头，计算效率下降明显；而 Ring 的通信虽然多了，但计算负载更均衡。这也印证了前面说的选型原则：CP 越大，Ring 越有优势。
5. 实操中踩过的坑与排查技巧
5.1 通信组配置错误导致的静默失败
最常见的问题之一是通信组配错了，但程序不报错，只是结果不对。比如 CP 组和 TP 组的 rank 映射搞混了，导致 K/V 传给了错误的卡。这种问题在小规模测试时可能看不出来，一上大规模就暴露。
我的排查方法是：在初始化后打印每个 rank 的通信组信息，确认 CP 组内的 rank 是连续的，并且和序列切分的顺序一致。Megatron 提供了--debug相关的日志选项，打开后能看到通信组的构建过程。
5.2 mask 偏移导致的收敛异常
前面提过 mask 的坑，这里补充具体的排查步骤：
- 构造一个极小的测试用例，序列长度等于 CP size，这样每张卡只有一个 token。
- 手动计算这个用例的注意力输出。
- 对比分布式实现的输出，逐元素检查差异。
- 如果差异出现在特定位置，检查 mask 的生成逻辑。
这个方法虽然笨，但非常有效。我靠它定位过好几次 mask 相关的 bug。
5.3 显存碎片导致的 OOM
CP 场景下显存分配比较频繁，容易产生碎片。有时候理论显存够，但实际就是 OOM。解决办法：
- 设置
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True，让 PyTorch 使用可扩展的显存段。 - 尽量复用 buffer，不要在注意力循环里反复申请释放。
- 如果用的是 TE，确认
flash后端的 workspace 大小是否合理。
5.4 常见问题速查表
| 现象 | 可能原因 | 排查方向 |
|---|---|---|
| 训练不报错但 loss 不降 | mask 错误、通信组错 | 小规模对比测试 |
| 吞吐远低于预期 | 通信未重叠、CP 配置不当 | profiler 看通信占比 |
| OOM 但显存理论够 | 显存碎片 | 开启 expandable_segments |
| CP 增大后性能下降 | Ulysses 头数不足 | 换 Ring 或减小 CP |
| 多机训练通信慢 | 网络拓扑、通信组跨节点 | 检查 rank 映射 |
6. 混合方案与后续扩展思路
实际生产中，纯 Ring 或纯 Ulysses 都不一定是最优解。我最近在试的一种混合方案是：在节点内用 Ulysses 做头维重排，在节点间用 Ring 做序列传递。这样节点内的 all-to-all 走 NVLink，带宽高；节点间的 Ring 传递走 RDMA，虽然带宽低但可以重叠。
这个方案的配置复杂度比单一方案高不少，需要仔细设计通信组的层次结构。但实测下来，在 16 卡跨两节点的场景下，比纯 Ring 快了约 18%，比纯 Ulysses 快了约 25%。
另一个值得关注的方向是CP 和序列并行的结合。有些框架开始把 CP 和 SP（Sequence Parallelism）统一起来，在 LayerNorm 和 Dropout 这些非注意力层也做序列切分，进一步降低激活内存。这个思路和 TP 里的 SP 是一脉相承的，本质上都是让序列维度在更多地方被切分。
如果你正在做长上下文训练，我的建议是先把 CP 的基础配置跑通，用小规模测试验证正确性，再逐步放大序列长度和 CP size。不要一上来就追求极限配置，那样出了问题很难定位。CP 这块的调试成本主要花在通信和 mask 上，把这两个搞定了，剩下的就是调参的事。
