Sequence Parallel（SP）与 Context Parallel（CP）都会把 [sequence, batch, hidden] activation 沿 token 维切到多张 GPU，因此很容易被当成同一种技术。真正的区别不在“切哪个维度”，而在分片保持到哪里：SP 主要消除 Tensor Parallel ranks 之间重复保存的 LayerNorm、dropout、residual activations；CP 则让网络输入和几乎全部 activations 都保持 sequence-sharded，并专门解决 Attention 对全局上下文的依赖。

可以先记住一句话：

SP：在 TP block 的边界省 activation，进入需要完整 token 范围的计算前再聚合
CP：整个网络都只持有本地 token，Attention 内显式交换远端上下文

两者都不是“把 sequence 平均切开就结束”。要判断实现是否正确，必须追踪每个 tensor 的 logical placement、collective 和 global token position。

从 Transformer Activation 的 Shape 开始

设一个 Transformer block 的输入为：

\[X\in\mathbb{R}^{S\times B\times H}\]

其中 $S$ 是 sequence length，$B$ 是 micro-batch size，$H$ 是 hidden size。若使用 $p$ 张卡沿 sequence 分片，每张卡理想上只保存：

\[X_r\in\mathbb{R}^{(S/p)\times B\times H}\]

单份 activation 的理论显存从：

\[M_X=S\times B\times H\times bytes(dtype)\]

下降为：

\[M_{X_r}\approx\frac{M_X}{p}\]

但这个除以 $p$ 只有在 tensor 整个生命周期都保持分片时才成立。若下一层立刻 AllGather 回完整 $S$，峰值显存仍可能出现完整 activation；若 Attention 还 materialize 大型临时 tensor，真正峰值也不只由 block input 决定。

为什么 Tensor Parallel 仍会复制大量 Activation

Megatron-style Tensor Parallel 把 MLP 与 Attention 的线性层沿 hidden/features 或 heads 切开。Column-parallel linear 产生 feature-sharded 输出，row-parallel linear 产生 partial output，传统实现用 AllReduce 恢复 replicated residual：

replicated residual X
  → column-parallel linear
  → feature-sharded hidden
  → row-parallel linear
  → partial [S,B,H] on every rank
  → AllReduce
  → replicated [S,B,H] on every rank

AllReduce 后，每个 TP rank 都持有相同的完整 residual。LayerNorm、dropout 和 residual add 虽然计算量不大，却让 $S\times B\times H$ activation 在 TP group 内复制 $p_{TP}$ 份。

TP 已经切开 parameters 和大 GEMM，并不等于所有 activations 都自然分片。Sequence Parallel 正是针对这段 replicated 区域。

Sequence Parallel：把 AllReduce 拆成 ReduceScatter 与 AllGather

从通信语义看：

\[\operatorname{AllReduce} = \operatorname{AllGather} \circ \operatorname{ReduceScatter}\]

传统 row-parallel output 把 partial results AllReduce 后，所有 ranks 立刻拿到完整 [S,B,H]。SP 改成 ReduceScatter：先完成跨 rank 求和，同时让每个 rank 只保留一段 sequence。

row-parallel partial outputs
  → ReduceScatter(sequence)
  → sequence-sharded residual
  → local dropout / residual / LayerNorm
  → AllGather(sequence)
  → next column-parallel linear

本来 AllReduce 已经包含等价的 reduce 与 gather 通信。SP 并非额外凭空增加一整套数据搬运，而是推迟 AllGather，让两个 collectives 之间的 activation 保持 sequence-sharded，从而减少驻留显存。

这段区域里的操作必须对 token 独立：

  • LayerNorm 通常沿 hidden dimension 归一化，每个 token 可独立计算；
  • elementwise activation、bias、dropout 可对本地 tokens 执行；
  • residual add 只要两侧使用相同 token shard，也可本地完成；
  • token-wise MLP 的非线性本身不需要看其他 tokens。

若某个操作沿 sequence 求统计量，或者需要任意 token 与其他 token 交互，就不能直接把它留在普通 SP 本地区域。

SP 的 Forward 与 Backward Placement

设 TP size 为 $p$，SP 与 TP 使用同一个 process group。一个简化 block 的 Forward placement 是：

X: sequence-sharded
  → local LayerNorm
  → AllGather sequence
X_full: replicated over sequence, or enters TP-specific layout
  → column/row-parallel Attention or MLP
  → ReduceScatter sequence
Y: sequence-sharded
  → local dropout + residual

Backward 按相反方向恢复 tensor placement。Forward 的 AllGather 对应 Backward 的 ReduceScatter；Forward 的 ReduceScatter 对应 Backward 的 AllGather。实现还可能把 collective 与 GEMM 分块重叠，但 logical result 不变。

最常见的错误不是 shape 不匹配，而是 shape 相同、语义错误。例如每个 rank 都有 [S/p,B,H]，但：

  • 它可能是不同 token ranges 的有效 shards；
  • 可能是对同一 token range 的 partial sums；
  • 也可能是错误地重复保存同一段 tokens。

因此日志和 API 不能只记录 shape，还要记录 global token 映射（连续分片可简写为 global_offset）、shard dimension、是否 partial 以及所属 process group。

SP 为什么没有解决长上下文 Attention

Self-Attention 对每个 Query token $i$ 需要读取所有允许位置的 Key/Value：

\[O_i = \operatorname{softmax} \left( \frac{Q_iK^T}{\sqrt d}+M_i \right)V\]

即使 rank $r$ 只负责一段 $Q_r$，它仍需要全局 $K,V$ 才能得到精确结果。普通 SP 在 Attention 前把所需 sequence 数据聚合回来，然后仍按 TP heads/features 计算；它只是让 Attention 之外的一段 activation 生命周期被分片。

所以 SP 主要解决“TP ranks 为什么重复保存 token-local activation”，并没有让单个 Attention head 的上下文天然跨设备扩展。真正把 Attention 的 sequence 也切开的，是 Context Parallel 或 Ring Attention 一类算法。

Context Parallel：让本地 Query 看见全局 KV

CP 把输入和整网 activations 沿 sequence 分给独立的 CP ranks：

rank 0 owns tokens 0..S/p-1
rank 1 owns tokens S/p..2S/p-1
...
rank p-1 owns the final token range

Linear、LayerNorm、MLP 等不跨 token 的模块可以直接处理本地 shard。Attention 必须额外通信，使本地 Queries 与全局 Keys/Values 发生交互。

最直观的实现是 AllGather KV：

local Q_r, K_r, V_r
  → AllGather K/V across CP group
  → local Q_r attends to full K/V
  → keep only output O_r for local queries

它简单，但每张卡会在 Attention 内临时持有完整 KV，通信与峰值显存可能较高。GQA/MQA 的 KV heads 较少时，这条路径可能仍有吸引力；MHA 或超长上下文下则需要更流式的办法。

P2P Ring：让 KV Blocks 依次经过每个 Query Shard

Ring-style CP 固定本地 $Q_r$，把 $K,V$ blocks 沿环传递。每收到一块，就计算本地 Query 对该块的 attention contribution，同时发送上一块、接收下一块：

step 0: Q_r × local KV_r
step 1: Q_r × KV_(r-1)
step 2: Q_r × KV_(r-2)
...

不能分别对每个 KV block 做 Softmax 后再把 outputs 相加，因为 Softmax denominator 是全局的。Blockwise exact attention 要维护 online statistics。若第 $j$ 块得到 local maximum $m_j$、exponential sum $l_j$ 和未归一化 weighted value $u_j$，合并到 running state 时需要按新的最大值重新缩放：

\[m'=\max(m,m_j)\] \[l'=e^{m-m'}l+e^{m_j-m'}l_j\] \[u'=e^{m-m'}u+e^{m_j-m'}u_j\]

最终：

\[O=\frac{u}{l}\]

这些统计量都按 Query row 独立维护。初始化为 $m=-\infty,l=0,u=0$；若某个 KV block 对该 row 完全被 mask，应直接跳过，不能先对一整行 $-\infty$ 做减最大值，否则 $-\infty-(-\infty)$ 会产生 NaN。首个有效 block 可以直接初始化 running state，之后再应用上述合并式；最终只有 $l>0$ 的 row 才能计算 $u/l$。若一个 row 完全没有合法 Key，必须走 attention backend 约定的空行处理，不能无条件除以零。这是 online softmax 在 causal/padding mask 下的数值边界，不改变合法位置上的 attention 定义。

这样无需 materialize 完整 attention matrix，也无需在单卡保存完整 KV。P2P 传输能否隐藏，取决于每个 local attention block 的计算时间是否足以覆盖下一块 KV 的通信。

All-to-All：在 Sequence 与 Head 维之间转置

另一类 CP 做法使用 All-to-All，在 sequence-sharded 与 head-sharded layout 之间转换。可以把它理解为一次分布式转置：

before A2A:
  each rank has all local heads for S/p tokens

after A2A:
  each rank has a subset of heads for full S tokens

每个 rank 获得完整 sequence、较少 heads，于是本地完成这些 heads 的 Attention；输出再通过反向 All-to-All 恢复 sequence shards。DeepSpeed Ulysses 属于这类思路。

它避免环形多步传输，却要求 head 数、CP degree 与 layout 合法匹配。若 GQA 的 KV heads 很少，继续沿 heads 切分可能出现每 rank 不足一个 KV head，需要复制、分组或采用混合路径。

CP 通信方式不是一个固定算法

现代训练栈通常提供多种 Context Parallel communication：

方式 核心动作 优势 主要约束
AllGather 聚合完整 KV 后计算 简单，易验证 临时 KV 与暴露通信较大
P2P/Ring KV blocks 环形流动 可与 block attention 重叠 多步调度、mask 与负载均衡复杂
All-to-All sequence/head layout 转置 通信轮次清晰 受 head 数与 A2A 拓扑限制
A2A+P2P 节点内 A2A、节点间 P2P 等分层组合 贴合分层互联 process groups 与调优更复杂

Megatron Core 当前把这些路径暴露为 p2p、all_gather、a2a 与 a2a+p2p。选择依据应是模型的 Q/KV head 数、sequence length、节点内外拓扑、通信能否重叠和 kernel block shape，而不是只看算法名字。

Causal Mask 会产生负载不均衡

双向 Attention 中，每个 Query shard 大致处理相同数量的有效 QK pairs。Causal Attention 只允许 token 看到自己及更早位置。若简单按连续 sequence 分片：

  • 早期 Query shard 可见的历史较短；
  • 后期 Query shard 要处理更多有效 KV blocks；
  • 不同 ranks 的有效 attention work 明显不均衡。

一种思路是让每个 rank 同时拥有靠前和靠后的 token chunks，用互补位置平衡有效三角区域；另一种是在 block schedule 中跳过完全被 mask 的块，并重新安排通信/计算顺序。

无论使用哪种方式，global position 必须保持正确。把本地 token index 0..S/p-1 直接当成全局位置，会同时破坏 causal mask、RoPE phase 和数据对齐。

Position、Mask 与 RNG 必须跟随 Global Token

Sequence sharding 后至少有三类隐式状态需要显式化。

Global Position

只有 Rank $r$ 持有一段连续 token 时，第 $j$ 个本地 token 才能写成：

\[position=offset_r+j\]

若采用前文的首尾配对或其他非连续分片，应使用显式映射 $position=index_map_r[j]$。例如全局位置 0..7 中，某 rank 持有 [0,1,6,7]，单个 offset_r+j 无法表示后半段。若使用 packed sequences，还要由 segment ID 与边界得到样本内的 position IDs；物理全局下标、样本内 RoPE 位置和 causal 可见性不能混为同一整数。

Attention Mask

Mask 要根据 global query/key positions、segment id、padding 和 causal/bidirectional 规则生成。环形 KV 走到新 rank 后，它携带的不是“当前本地第几个 block”，而是明确的 global key range。

RNG Mapping

Dropout mask 应能由 global token、layer、micro-batch 和 RNG stream 唯一决定。Checkpoint 恢复或 CP degree 改变后，若随机数仅依赖 local rank 与 local tensor shape，数据相同也可能得到不同 mask，导致不可复现甚至破坏训练等价性。

SP、CP 与 TP 怎样组成二维切分

假设 TP=4、CP=2，一份 model replica 需要 8 张 GPU。可以把 ranks 画成二维网格：

             TP position
           0   1   2   3
CP shard 0  0   1   2   3
CP shard 1  4   5   6   7
  • [0,1,2,3] 和 [4,5,6,7] 分别是 TP groups，共同切 layer features/heads；
  • [0,4]、[1,5]、[2,6]、[3,7] 是 CP groups，交换对应 TP position 的上下文；
  • SP 通常复用每行的 TP group，让 token-local activations 在 TP ranks 间 sequence-sharded；
  • CP 则沿列切完整网络的 token ranges。

加入 PP 与 DP 后，总卡数在不考虑 EP 等额外维度时为：

\[W=p_{TP}\times p_{CP}\times p_{PP}\times p_{DP}\]

每种 collective 必须落在正确 group。把 CP KV exchange 错发到 TP group，shape 可能仍合法，却会把不同 heads 而不是不同 token ranges 拼在一起。

SP 与 CP 同时启用时发生什么

CP 先决定每个 CP rank 拥有全局 sequence 的哪一段；在每个 CP shard 内，TP/SP 还可以进一步处理这段局部 sequence。

若全局 $S$ 个 tokens，CP size 为 $c$，TP/SP size 为 $t$，某些 SP-resident activations 的每 rank token 数可能近似为：

\[S_{local}\approx\frac{S}{c\times t}\]

但不能由这个公式推断所有 tensors 都除以 $ct$：

  • TP linear 内部可能暂时 AllGather SP shards；
  • CP attention 会交换或重排 KV/context；
  • parameters 沿 TP 切、在 CP group 内通常复制；
  • gradient synchronization 还要结合 DP×CP replicas；
  • fused kernel 可能要求额外 padding 与 alignment。

正确做法是逐算子标注 placement，而不是给整个模型贴一个“sequence 已除以 $ct$”的标签。

GQA、MQA 与 MLA 怎样改变 CP 成本

CP Attention 的通信主要与传输的 K/V 表示相关。MHA 的 KV heads 与 Query heads 同量，KV payload 较大；GQA/MQA 减少 KV heads，因此 AllGather 或 Ring 中每个 token 的 KV bytes 下降。

若每个 token、每层 KV payload 近似为：

\[V_{KV/token} = 2\times n_{kv}\times d_h\times bytes(dtype)\]

减小 $n_{kv}$ 会直接减少 CP 通信。MLA 则可能传递压缩 latent 和位置相关分量，具体收益取决于训练实现是否能在 Attention 计算前后保持压缩表示，而不是在通信前就展开为完整 K/V。

因此同一个 CP degree 在 MHA、GQA 和 MLA 上的瓶颈可能完全不同。配置必须基于真实 boundary tensors 测量。

Packed Sequence 与变长样本为什么更难

长上下文训练常把多个短样本 pack 到固定 token budget。此时一个 physical sequence tensor 里可能包含多个互相不可见的 segments。CP 分片不能只切连续字节，还要保留：

  • 每个 token 的 segment id 与 global-in-segment position；
  • 跨 CP shard 的 segment boundaries；
  • causal mask 在 segment 处重置；
  • load balance 不能把 padding 节省又变成某些 ranks 的空闲；
  • checkpoint 后 data loader 要恢复同一 packing/cursor 状态。

若不同 CP ranks 的有效 token 数差异大，local MLP、Attention blocks 与通信 payload 都会失衡。仅让 padded tensor shape 整齐，并不代表有效工作均衡。

与 Activation Checkpointing 的取舍

不使用 CP 时，长 sequence OOM 常通过 activation recomputation 缓解：少保存中间 activation，Backward 时重跑 Forward。SP/CP 则通过分片减少每 rank activation。

两者可以组合，但优化目标不同：

  • recomputation 用额外 FLOPs 换显存；
  • SP 主要去掉 TP ranks 间的 activation replication；
  • CP 用额外 context communication 换取全网 sequence 分片；
  • selective recomputation 只重算显存占比高、计算相对便宜的算子。

选择时应比较 step time，而不是只看峰值显存。CP 若使 Attention 通信完全暴露，可能比保留较小 CP、增加 selective recomputation 更慢。

怎样选择 SP 与 CP

可以按问题来源判断：

使用 TP 后，LayerNorm/Residual Activation 复制过多

优先考虑 SP。它复用 TP group，并把 AllReduce 拆成 ReduceScatter/AllGather，通常不改变 Attention 的全局计算方式。

单个 Sequence 的 Attention/Activation 无法放入现有 TP 布局

考虑 CP。它让完整网络沿 sequence 分片，但要为 Attention 选择 KV AllGather、P2P Ring、A2A 或分层通信。

Sequence 不长，主要瓶颈是 Parameter/Optimizer State

SP/CP 不是第一选择。先看 FSDP/ZeRO、TP、PP 与 activation checkpointing。

Head 数少、GQA/MQA 明显

KV AllGather 可能比想象中便宜；A2A 沿 head 切分反而受 head 数限制。应以真实 Q/KV layout 选择。

考虑把 TP 放在节点内，CP 使用拓扑感知的 P2P 或 A2A+P2P 层次结构，并验证通信/计算 overlap。

正确性验证不能只对最终 Loss

小规模 reference 测试应逐层比较：

  1. 无 SP/CP 的单卡或复制基线；
  2. LayerNorm、residual 和 dropout 后的 global tensor；
  3. Q/K/V 的 global position 与 head ownership；
  4. Attention online Softmax 的 max、sum 与 output；
  5. causal/packed mask 的有效 blocks；
  6. input gradient 与 parameter gradient；
  7. optimizer update 后的 global weights；
  8. checkpoint 恢复后的下一步 loss；
  9. 改变 SP/CP degree 后的 reshard parity。

测试 tensor 应刻意包含：不能被 shard size 整除的长度、多个 packed segments、GQA 少 KV heads、dropout、causal 与 bidirectional mask。只用全零 mask、无 dropout、规则长度，会漏掉最危险的 placement 问题。

性能分析要看哪些量

至少记录：

  • 每 rank peak activation memory 与 temporary KV memory；
  • SP AllGather/ReduceScatter bytes、等待时间与 overlap；
  • CP 每层的通信类型、payload 与 exposed latency；
  • local attention block shape 与 kernel efficiency；
  • causal/packed sequence 的有效 work imbalance；
  • TP、CP collectives 是否争用同一链路和 SM；
  • recomputation FLOPs 与节省的 activation bytes；
  • tokens/s、MFU、step time P50/P99 与 loss parity。

“通信 kernel 与 GEMM 在 timeline 上重叠”不等于通信被免费隐藏。两者可能争抢 HBM bandwidth、NVLink 或 SM，导致 GEMM 本身变慢。应比较重叠前后的整个 critical path。

常见误区

“SP 和 CP 都切 Sequence，所以只是不同名字”

SP 主要作用于 TP block 边界的 token-local activations；CP 让整个网络保持 sequence-sharded，并在 Attention 内交换全局上下文。

“AllGather 之后显存还是完整的，所以 SP 没有意义”

SP 的价值在于缩短完整 activation 的驻留区间，并让 LayerNorm、dropout、residual 等区域保持分片。峰值是否下降要结合具体 buffer 生命周期判断。

“CP 会把 Attention 计算复杂度除以设备数，所以一定线性加速”

每 rank 计算量下降，但全局 dense Attention 的 $O(S^2)$ 工作没有消失；还增加了 KV/context 通信、负载均衡和 kernel 变小的成本。

“Local Token Index 就是 Position ID”

Sequence shard 的本地下标必须先映射到真实 global token，再结合 packed segment boundaries 生成位置与 mask。只有连续分片才能简化为加一个 global offset。

“CP Size 越大，能支持的上下文越长，性能也越好”

更大 CP 会缩小本地 block、增加通信参与者并改变拓扑路径。它扩展的是容量上限，不保证时间线性扩展。

小结

Sequence Parallel 与 Context Parallel 的共同点只是沿 token 维分片；它们解决的是不同层次的问题：

  1. SP 复用 TP group，把 row-parallel AllReduce 拆成 ReduceScatter 与延迟的 AllGather；
  2. 两个 collectives 之间的 LayerNorm、dropout、residual activations 保持 sequence-sharded；
  3. 普通 SP 在 Attention 前仍会聚合所需 sequence，不能独立扩展单个 head 的上下文；
  4. CP 让输入和整网 activations 沿 sequence 分片，并在 Attention 内交换全局 K/V；
  5. CP 可使用 KV AllGather、P2P Ring、All-to-All 或分层混合通信；
  6. causal mask、global position、packed segments 与 RNG 必须跟随 logical token；
  7. SP、CP、TP、PP、DP 组成多维 mesh，每个 tensor 和 collective 都要落在正确 group；
  8. 最优方案取决于激活显存、Q/KV heads、网络拓扑、block shape 与 recomputation 成本。

下一篇会进入分布式 Checkpoint：当参数、Optimizer State 和 RNG 已经分散在多维 device mesh 上，怎样保存一份可验证、可异步提交，并能在不同并行度下恢复的训练状态。

参考资料