序列并行 SP / 上下文并行 CP

9.1 SP 与 CP 听起来都像"切序列"，但完全不是一回事

SP 是 TP 的"补丁"，CP 是独立维度。两者可以同时开，但用途完全不同。

9.2 SP：TP 的"补丁"

问题

第 7 章讲过 TP 切 attention / FFN 的 weight，但 LayerNorm 之前的激活、dropout、loss 这些位置仍然是完整复制的。设 hidden $H$、seq $S$、batch $B$：

$$ M_\text{norm act} = B \cdot S \cdot H \cdot 2\,\text{bytes} \quad (\text{每卡完整}) $$

Llama-3 70B 70 层、$B{=}4$、$S{=}8192$、$H{=}8192$ → 每层 norm 输入约 0.5 GB，70 层就 35 GB，加上反向梯度 70 GB，压根装不下。

SP 的招

把这些"夹在 TP 区域之间"的激活也按 seq 维切到 $T$ 张卡，需要时再 gather 回来：

关键变换：原来 TP 区域结尾的 all-reduce 变成 reduce-scatter；原来 TP 区域开头的"复制"变成 all-gather。通信总量没变，但激活内存降到 $1/T$。

SP 显存收益

$$ M_\text{act, full TP} = O(B \cdot S \cdot H) \;\longrightarrow\; M_\text{act, TP+SP} = O(B \cdot S \cdot H / T) $$

对长序列 + 大 TP 的训练，SP 几乎是必开。

9.3 SP 怎么用

# SP 必须跟 TP > 1 一起开
--tensor-model-parallel-size 4 \
--sequence-parallel

不开 TP 的情况下 --sequence-parallel 没意义（毫无切分对象）。

SP 与 dropout

SP 区域内的 dropout 在 seq 维度被切，每卡只看自己那段，dropout mask 互相独立，这是正确的（不像 TP 区域里同步 RNG 那样）。

9.4 CP：让长上下文成为可能

动机

attention 内存增长是 $O(S^2)$：

$$ M_\text{attn} = B \cdot N_\text{heads} \cdot S^2 \cdot \text{dtype bytes} $$

$S{=}32$K 时 attention 矩阵 $\approx 32^2 = 1024$ 倍于 $S{=}1$K 的情况。FlashAttention 把内存降到 $O(S)$，但仍然单卡装不下千 K 序列。

CP 的招

把一条序列沿时间维切到 $C$ 张卡，每卡处理 $S/C$ 个 token。问题是 attention 是"每个 token 看所有历史 token"，跨卡怎么办？

答案：ring attention —— 让序列在卡之间循环传递，每张卡轮流接收别人的 KV 做局部 attention，最后把所有 partial output 聚合起来。

关键点

• 因果 mask（causal）下，前面 token 不能看后面，只需通信一半左右；
• 每张卡显存只装 $1/C$ 序列对应的 KV cache 和激活；
• 通信量随序列长度线性增长，但用 NVLink/IB 还能扛。

9.5 CP 怎么用

# 4 卡 CP，sequence length 32K
--seq-length 32768 \
--max-position-embeddings 32768 \
--context-parallel-size 4 \
--rotary-base 1000000              # 长上下文常配长 RoPE base
# CP 通常与 TP / SP 一起开
--tensor-model-parallel-size 4 \
--sequence-parallel

9.6 SP 与 CP 的关系

9.7 全维度并行总和

Megatron 在最大化并行后，单 step world size：

$$ W = \mathrm{TP} \times \mathrm{PP} \times \mathrm{DP} \times \mathrm{CP} \times \mathrm{EP}_\text{eff} $$

其中 $\mathrm{EP}_\text{eff}$ 是 MoE 的有效专家并行度（与 TP 复用）。Llama-3 405B 训练 PP=16 × TP=8 × CP=8 × DP=8 = 8192 卡这种规模。

9.8 长上下文配方：SP + CP + Flash-Attn

9.9 关键源码

9.10 常见错误

9.11 Mamba SSM + Context Parallel

2025-Q3 加入，专门解决 Mamba 这种 SSM 长序列下激活也会 OOM 的问题（递归式状态而非 attention，但激活仍随 seq 线性增长）：

启用：--use-mamba --mamba-num-layers 24 --context-parallel-size 4。注意：Mamba 与 Transformer 层混合的"hybrid 模型"在 pretrain_hybrid.py。

9.12 Heterogeneous Models（异构层）

2025-Q4 加入（arguments.py:3194 _add_heterogeneous_args）：

• 允许同一个模型里不同 transformer layer 用不同配置（hidden size / 是否 MoE / 是否 Mamba）；
• 主要给"前后 Dense + 中间 MoE"、"前后 Transformer + 中间 Mamba" 这类混合架构用；
• CLI：--heterogeneous-layers-config-path layers.json，每层一条 spec；
• PP 切分时 Megatron 会自动考虑各层负载差异。

9.13 Packed Sequences（变长打包）

SFT 阶段常见的"短序列浪费 padding"问题。megatron/core/datasets/packed_seq_params.py + arguments.py:3244 _add_experimental_args 里的相关 flag：

--use-packed-sequences \
--packed-seq-cu-seqlens-style flash_attn   # 与 flash-attn 的 cu_seqlens 对齐

原理：把短样本拼到一个 seq 里，用 cu_seqlens 告诉 attention "这一段到哪结束"。flash-attn-2/3 原生支持。SFT 提速 2-3 倍。

9.14 这章你需要带走什么

• SP 是 TP 的补丁，切 norm/dropout 激活，必须配 TP；
• CP 是独立维度，沿 seq 切 attention，长上下文必须；
• SP 通信量不增，激活降 1/T；
• CP 通信量随 seq 长，但让单卡能装下；
• 长上下文配方：SP + CP + flash-attn + recompute 一起上。