上一篇 TP 写完之后，我对 Megatron 的 Column Parallel + Row Parallel 已经比较清楚了：前一层按输出特征切分，后一层按输入特征切分，两层之间不需要 AllGather，只在最后做一次 AllReduce 把部分和加起来。
但继续看 MIT 6.S191 Lecture 9 关于 Sequence Parallelism 的部分时，我发现一个之前没注意到的问题。
Row Parallel 最后那一步 AllReduce 之后，每张 GPU 拿到的是完全相同的完整输出 \(Y \in \mathbb{R}^{S \times H}\)。紧接着是 LayerNorm——而 Normalization 那篇已经写清楚了，LayerNorm 对每个 Token 的 Hidden 维度独立计算统计量，Token 之间互不依赖。
既然每个 Token 的 LayerNorm 都是独立的，为什么还要让两张 GPU 各自对全部 \(S\) 个 Token 做相同的计算？
再往后看 Context Parallelism，问题变得更有意思：它和 SP 一样，也是沿着 Sequence 维度切分。切的维度一样，为什么还需要另起一个名字？
TP 做了计算分片，Activation 为什么还有冗余#
用 TP 那篇的两层 MLP 例子。TP = 2，Sequence Length = 8，Hidden Dimension = 4。
第一层 Column Parallel：
$$ W_1 = [W_1^{(0)} \mid W_1^{(1)}] $$每张 GPU 得到不同的输出特征分片。ReLU 是逐元素算子，在各卡本地执行。
第二层 Row Parallel：
$$ Y = Z_0 + Z_1, \quad Z_i = H_i W_2^{(i)} $$AllReduce(SUM) 之后，两张 GPU 都持有完整的 \(Y \in \mathbb{R}^{8 \times 4}\)。
到这里我想了一会儿。TP 并不是在所有阶段都产生相同的 Activation。Column Parallel 的输出在两卡上是不同的——GPU 0 持有前 4 列输出，GPU 1 持有后 4 列。冗余发生在一个特定的位置：Row Parallel 通过 AllReduce 把部分和归约之后，两卡拿到了完全相同的完整 Tensor。
这里还要区分三种显存占用。参数在 TP 中是持久分片的（Column 切列，Row 切行），每卡只存自己的一份。梯度与参数一一对应，也是分片的。但 Activation 不一样——AllReduce 之后产生的 \(Y\) 是完整复制的，两卡各存一份 \(8 \times 4\) 的完整 Tensor。
接下来这个完整的 \(Y\) 被送进 LayerNorm。两张 GPU 对相同的 8 个 Token 做完全相同的计算，产生相同的输出。计算冗余，Activation 存储也冗余。
那为什么不让 GPU 0 只负责 Token 0–3 的 LayerNorm，GPU 1 负责 Token 4–7？
SP 的做法：在 ReduceScatter 和 AllGather 之间插入计算#
通信原语那篇介绍过 ReduceScatter：它先跨 GPU 做 Reduce(SUM)，再把归约结果的不同分片分发给不同 GPU。也介绍过一个关键等式：
$$ \operatorname{AllReduce(SUM)} = \operatorname{AllGather} \circ \operatorname{ReduceScatter(SUM)} $$这个等式的含义是：先 ReduceScatter 再 AllGather，最终每张 GPU 得到的结果和直接做 AllReduce 相同。
在普通 TP 中，Row Parallel 的部分和通过 AllReduce 一步到位变成完整的 \(Y\)。SP 做的事情是把这一步拆开——先 ReduceScatter，计算完 LayerNorm 之后，再 AllGather。
普通 TP：Row Parallel → AllReduce → 两卡各持完整 \(Y\) \([8,4]\) → LayerNorm
TP + SP：Row Parallel → ReduceScatter → GPU 0 持有 \(Y[0{:}4]\) \([4,4]\)，GPU 1 持有 \(Y[4{:}8]\) \([4,4]\) → LayerNorm → AllGather
ReduceScatter 之后，每卡只持有一半 Token 的 Activation，但每个 Token 的完整 Hidden 向量都在本地。LayerNorm 需要的恰好就是这个——对每个 Token 沿 Hidden 维度算 \(\mu\) 和 \(\sigma\)。
这里有一个很容易犯的错误。ReduceScatter 不是简单地把两卡的局部贡献直接 Scatter 出去，它先做了 Reduce。如果跳过 Reduce 直接把 Row Parallel 的部分和 \(Z_0\) 送进 LayerNorm：
$$ \operatorname{LN}(Z_0 + Z_1) \neq \operatorname{LN}(Z_0) + \operatorname{LN}(Z_1) $$LayerNorm 包含非线性操作（除以标准差），对加法不满足分配律。GPU 0 上的 \(Z_0\) 不是完整的输出，它只是部分和，基于它算出的均值和方差是错的。
但 ReduceScatter 之后的情况不同。这时候 GPU 0 上的 \(Y[0{:}4]\) 已经是完整归约的结果——Token 0 的所有 4 个 Hidden 维度都在 GPU 0 上，数值正确。对它做 LayerNorm 和对完整 \(Y\) 的前 4 行做 LayerNorm 是等价的：
$$ \operatorname{LN}(X_0 \oplus_S X_1) = \operatorname{LN}(X_0) \oplus_S \operatorname{LN}(X_1) $$其中 \(\oplus_S\) 表示沿 Sequence 维度拼接。这个等式成立的前提：每个 Token 的完整 Hidden 向量都在本地，LN 的统计范围仅覆盖该 Token 的 Hidden 维度。
SP 没有改变 LayerNorm 的数学定义。它改变的是中间 Activation 的分布方式——从"每卡持有全部 Token"变成"每卡只持有自己负责的 Token"。
LayerNorm 之后，如果后续算子也不需要跨 Token 数据（比如 Residual Add、Dropout、逐元素激活函数），可以继续在分片状态下本地处理。AllGather 不必紧跟在 LayerNorm 后面。只有当后续算子确实需要完整 Sequence 时，才做 AllGather 恢复完整布局。
我觉得这背后有一个很重要的工程思想：不要在每个算子之后机械地恢复完整 Tensor，而应该尽可能延迟通信，直到后续算子真正需要。
关于通信量：标准 Ring Collective 下，ReduceScatter + AllGather 的总传输字节数和 AllReduce 相当。SP 的收益不在于减少通信量，而在于把通信拆散后，中间的计算可以在分片状态下执行，从而减少 Activation 显存。
但这不意味着整个 Transformer 的 Activation 显存能直接降到 \(1/\text{TP}\)。SP 只影响那些可以在 Sequence 分片上独立执行的算子（LayerNorm、Dropout、Residual Add 等）。Attention 和 MLP 的 TP 部分仍然按原来的方式工作。实际的 Peak Memory 变化取决于哪些算子处于 Sequence 分片状态、以及它们在整体显存中占多大比例。
MLP 能按 Sequence 切分，Attention 不行#
到这里我以为事情已经解决了——把 TP 的 AllReduce 拆成 ReduceScatter + AllGather，中间的算子在分片状态下执行，完事。
但 Self-Attention 打破了这个计划。
假设 \(X \in \mathbb{R}^{8 \times H}\)，按 Sequence 分给两张 GPU：GPU 0 持有 Token 0–3，GPU 1 持有 Token 4–7。
MLP 没有问题。\(y_i = \text{MLP}(x_i)\) 只需要当前 Token 的 Hidden 向量和模型权重，不需要其他 Token 的任何信息。当权重在本地可用时，每个 Token 的 MLP 可以独立执行。
Self-Attention 就不一样了。
$$ o_i = \sum_{j \in \mathcal{A}(i)} \alpha_{ij} v_j $$其中 \(\mathcal{A}(i)\) 是 Attention Mask 允许 Query \(i\) 关注的位置集合。Causal Attention 那篇已经分析过：在 Causal Mask 下，Token \(i\) 可以关注位置 \(0\) 到 \(i\)。
考虑 GPU 1 上的 Token 6。它需要关注 Token 0–6 的 Key 和 Value。其中 Token 0–3 的 KV 位于 GPU 0。GPU 1 不能只靠自己本地的 KV 完成 Token 6 的 Attention 计算。
反过来看 GPU 0。它负责 Token 0–3 的 Query。在 Causal Mask 下，Token 3 最多关注 Token 0–3，全部在本地。GPU 0 在 Forward 中不需要 GPU 1 的 KV。这种非对称的依赖关系在 Causal Attention 文章末尾已经讨论过。
所以问题不在于"能不能按 Sequence 切分"——任何 Tensor 都可以按任意维度切。问题在于某些算子存在跨 Token 的数据依赖。
SP（Megatron 风格）回避了这个问题。它在 Attention 部分继续使用 TP 的原有布局（Column/Row Parallel 的头分片方式），Sequence Sharding 只覆盖 LayerNorm、Dropout、Residual 等不需要跨 Token 数据的算子。
但如果训练的 Sequence 足够长，即使 Attention 的计算被 TP 分了头，每张 GPU 仍然要处理所有 \(S\) 个 Token 的 Query 与 KV。Activation 的显存随 \(S\) 增长——这正是 Context Parallelism 要解决的问题。
CP 的做法：固定本地 Query，跨 GPU 通信 KV#
Context Parallelism 更系统地沿 Sequence 切分：输入和 Activation 按 Sequence 分片，每张 GPU 只负责自己那部分 Token 的计算，包括 Attention。
先看最直观的方案。每张 GPU 持有自己的 Query Shard，通过 AllGather 收集完整的 KV，然后执行 Attention：
$$ O_{\text{local}} = \operatorname{Attention}(Q_{\text{local}}, K_{\text{global}}, V_{\text{global}}) $$这里 \(K_{\text{global}}\) 是当前 Query 在 Mask 约束下需要访问的 Key 集合。为什么不需要 AllGather Q？因为每张 GPU 只负责计算自己本地 Query 对应的 Attention Output。GPU 1 需要远端的 KV，但它不需要替 GPU 0 计算 Query——GPU 0 自己会做。
AllGather KV 在数学上是正确的。但每张 GPU 临时持有完整的 KV 会带来显存压力。Sequence 越长，KV 的 \(S \times d\) 就越大。我们希望避免一次性把所有远端 KV 聚合到本地。
一个更高效的方式是 KV Chunk Exchange：每张 GPU 固定自己的 Query，逐步接收所需的 KV Block，边收边算。
以 Causal Attention、2 GPU 为例：
- 每张 GPU 先用本地 KV 计算一轮 Attention，得到初始的 Softmax 状态。
- GPU 之间 P2P 交换 KV Block。GPU 1 收到 GPU 0 的 KV 后，用 Online Softmax 更新状态。GPU 0 收到 GPU 1 的 KV 后，因为 Causal Mask 下 Token 0–3 不可见 Token 4–7，可以跳过。
- 处理完所有需要的 KV Block 后，各 GPU 用最终状态算出完整的 Attention Output。
KV Block 处理完就可以释放，不需要永久驻留本地。
这里的示意是一个简化模型。实际的 Ring Attention / Megatron CP 在通信调度、计算重叠、Causal Mask 优化方面会更复杂。比如多 GPU 时 KV Block 可能沿 Ring 方向逐步传递，通信和计算可以流水线重叠。不同实现的 KV 交换顺序和 Block 划分策略也不完全相同。
还有一点容易忽略：Forward 的 KV 通信只是故事的一半。Backward 时，KV 在 Forward 中可能被多个远端 Query 使用，因此反向传播需要聚合来自不同 GPU 的梯度贡献。Forward 的 KV 交换方向和 Backward 的梯度聚合方向不一定对称。
CP 为什么与 FlashAttention 天然互补#
读到这里我有一个疑问：CP 的分块 KV 计算，怎么保证 Softmax 的数值正确？
普通 Softmax 需要知道全局最大值来做数值稳定的归一化。如果在每个 KV Block 上独立做 Softmax，各 Block 的局部最大值不同，分母也不同——不能简单把各 Block 的 Softmax Output 加起来。
$$ \operatorname{softmax}_{\text{global}} \neq \sum_b \operatorname{softmax}_{\text{block}_b} $$FlashAttention 那篇介绍过 Online Softmax 的做法——不需要预先看完所有 Score 就能增量计算 Softmax。它维护三个状态：
- \(m\)：当前已处理的所有 Score 中的最大值
- \(l\)：以 \(m\) 为基准的缩放指数和
- \(p\)：缩放后的 Value 加权累积
每收到一个新的 KV Block，用新 Block 的 Score 更新 \(m\)，根据新旧 \(m\) 的差异重新缩放 \(l\) 和 \(p\)，然后累加新 Block 的贡献。处理完所有可见 KV 后：
$$ O = \frac{p}{l} $$逐 Query 行做除法即可。
这和 CP 的需求天然契合：
- Q 留在本地，不需要移动。
- KV Block 逐个到来（不管是从本地还是从远端）。
- 每到一个 Block，更新 Online Softmax 状态。
- 不需要保留完整的 \(S \times S\) Score Matrix。
- 不需要同时在本地驻留所有远端 KV。
- 处理完最后一个可见 Block 后，得到数学上完整的 Attention Output。
FlashAttention 原本是单 GPU 上的优化，目标是减少 Attention 中间矩阵的 HBM 读写。CP 是多 GPU 上的并行策略，目标是把长 Sequence 的 Activation 和计算分摊到多张卡上。两者在不同层面解决不同的问题，但 Online Softmax 这个数学工具同时服务于两者。
需要强调的是：CP 不会把精确 Self-Attention 的总计算复杂度从 \(O(S^2 d)\) 降低到 \(O(Sd)\)。 它改变的是计算在多 GPU 间的分配方式。总的浮点运算量不变，系统性能还受到 KV 通信延迟、负载均衡、GPU 互连带宽以及通信与计算重叠效率的影响。
SP 和 CP 都切 Sequence，为什么两个名字#
这是我学这部分内容时最大的困惑之一。
先坦白一件事：Context 不是一个区别于 Sequence Length 的新 Tensor 维度。 Transformer 的输入通常是 \(X \in \mathbb{R}^{B \times S \times H}\)，被切分的维度就是 \(S\)。“Context” 在 NLP 中通常强调 Token 可以利用的上下文窗口，但它并不是一个与 \(S\) 独立的 Axis。
SP 和 CP 确实都在沿 Sequence 维度做文章。区分它们的不是"切哪个维度"，而是"在哪些算子上切"以及"怎么处理跨 Token 依赖"。
在 Megatron 的语境下：
| SP（Megatron 风格） | CP（Megatron 风格） | |
|---|---|---|
| 与 TP 的关系 | 作为 TP 的补充，在 TP 不覆盖的算子上消除冗余 | 更独立的并行维度，可与 TP/DP/PP 组合 |
| 覆盖的算子 | LayerNorm、Dropout、Residual Add 等 | 系统性覆盖包括 Attention 在内的整个 Transformer Block |
| 核心通信 | ReduceScatter / AllGather（替代 AllReduce） | 跨 GPU KV 通信（P2P / Ring / AllGather） |
| Attention 处理 | 继续使用 TP 的头分片方式 | 每 GPU 只计算本地 Query，通过 KV 通信满足依赖 |
| 主要动机 | 消除 TP 某些阶段的冗余 Activation | 支持长 Context 训练，分摊 Sequence 维度的显存与计算 |
这个区分是 Megatron 的术语约定。其他研究或框架中，“Sequence Parallelism” 的定义未必完全相同——比如有些工作用 “Sequence Parallelism” 泛指一切沿 Sequence 维度的并行策略，包括 Attention 内部的分块通信。阅读不同文献时需要注意上下文。
算子是否可以直接处理 Sequence 分片，取决于它的数据依赖。并行策略的名称，只是对不同计算与通信组织方式的约定。
写在最后#
从 TP 到 SP 再到 CP，我觉得有一条贯穿始终的线索。
TP 那篇的结论是：Column + Row 的巧妙之处在于让上一层的分片输出直接成为下一层的分片输入，避免中间不必要的通信。
SP 进一步提醒我：不要因为上一个算子需要归约，就认为必须把完整结果复制到所有 GPU 再执行后续计算。如果后续算子本身不需要跨 Token 数据，保持分片状态就够了。
CP 再进一步：即使一个算子（Self-Attention）确实存在跨 Token 依赖，也未必需要把所有远端数据一次性聚合到本地。重新组织计算顺序，通过分块通信和 Online Softmax 的增量更新，可以在不物化完整 KV 的情况下满足数学上的完整依赖。
一个 Tensor 在多 GPU 上可以采用不同的分片布局。真正决定是否需要通信的，不是"数据在哪里"，而是"下一个算子需要访问什么"。
这大概是分布式训练中反复出现的一个设计原则：不要机械地恢复完整数据，而是让数据布局去适应计算的需求。
参考资料#
- FSDP 已经能分片参数，为什么还需要 Tensor Parallelism？
- 从 AllReduce 到 Ring：分布式训练中的通信原语
- FlashAttention 的核心思想：为什么 Attention Matrix 根本不需要存下来？
- 为什么 Self-Attention 不能偷看未来？从 Causal Mask 到自回归训练
- Norm 到底归一化了谁？从 BatchNorm、LayerNorm 到 RMSNorm
- 从 DDP 到 ZeRO：为什么分布式训练要对参数、梯度和优化器状态分片？
- MIT 6.S191: Introduction to Deep Learning — Lecture 9: Secrets to Massively Parallel Training. https://introtodeeplearning.com/
- Shoeybi et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2019.
- Korthikanti et al. Reducing Activation Recomputation in Large Transformer Models. 2022.
- Liu et al. Ring Attention with Blockwise Transformers for Near-Infinite Context. 2023.
- Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. 2022.
- NVIDIA Megatron Core — Context Parallel, Parallelism Strategies.