Sequence Parallel（Megatron SP）怎么做？
为什么通信量不变？省多少 activation memory？
先弄清楚：每个概念是什么
- TP 没切到的区域
- TP 只切 attention 和 MLP 里的 GEMM。夹在中间的 LayerNorm、dropout、residual add 每卡都算一遍、存一整份 [\(s, b, h\)]。
- Sequence parallel (SP)
- 这些运算对每个 token 各算各的，所以沿 \(s\) 切开，每卡只管 \(s/t\) 个 token。例：\(s = 2048\)、\(t = 8\) → 每卡只存 256 个 token 的 LayerNorm 输入。
- Reduce-scatter 与 all-gather
- Reduce-scatter 先求和，再让每卡只拿结果的 \(1/t\)。All-gather 把各卡那块拼回完整的。先后做这两步等于一次 all-reduce（图 2），详见 all-reduce。
- Activation memory
- forward 时存下来给 backward 用的 tensor。下面按 bytes 算（16-bit，dropout mask 1 byte），不做并行时每层 \(sbh(34 + 5as/h)\)。
看图：TP 区和 SP 区怎么交替，通信为什么不变
30 秒回答
Plain TP only splits the GEMMs, so LayerNorm, dropout and the residual add are replicated on every rank with full s×b×h activations. Those ops are independent per token, so sequence parallelism splits them along the sequence dimension. The all-reduce after the row-parallel linear becomes a reduce-scatter that leaves each rank s/t tokens, and an all-gather rebuilds the full sequence before the next column-parallel linear; backward swaps them. A ring all-reduce is exactly reduce-scatter plus all-gather, so communication volume is unchanged. Every activation is now divided by t: per layer, sbh(10 + 24/t + 5as/(ht)) becomes sbh(34 + 5as/h)/t.
推导 / 关键公式
MLP 这一段怎么接（上标 \(s\) = 按 sequence 切）
\[ [Y_1^s, Y_2^s] = \mathrm{LayerNorm}([X_1^s, X_2^s]) \] \[ Y = g(Y_1^s, Y_2^s) \quad \text{all-gather} \] \[ W_i = \mathrm{GeLU}(Y A_i)\, B_i \quad \text{column → row} \] \[ [W_1^s, W_2^s] = \bar g(W_1, W_2) \quad \text{reduce-scatter} \]\(g\) 在 forward 做 all-gather、backward 做 reduce-scatter，\(\bar g\) 相反。每层 forward + backward 共 4 次 all-gather + 4 次 reduce-scatter，替代原来的 4 次 all-reduce。
通信量为什么不变
\[ \underbrace{\frac{2(t-1)}{t}S}_{\text{all-reduce}} = \underbrace{\frac{t-1}{t}S}_{\text{reduce-scatter}} + \underbrace{\frac{t-1}{t}S}_{\text{all-gather}} \]Activation memory（bytes / 层）
\[ \text{TP：}\ sbh\Big(10 + \frac{24}{t} + \frac{5as}{ht}\Big) \] \[ \text{TP + SP：}\ \frac{sbh}{t}\Big(34 + \frac{5as}{h}\Big) \]那个不随 \(t\) 降的 10 = 两个 LayerNorm 的输入（4）+ 两个 dropout mask（2）+ 两个 column linear 的输入（4）。SP 之后连 column linear 的输入也只存本卡那段 \(Y_i^s\)。算 weight gradient 要完整的 \(Y\)，所以 backward 再 all-gather 一次，它和算 \(\partial L/\partial Y\) 的 GEMM overlap，几乎不增加时间。
手算例子：GPT-3 175B
\(a = 96\)、\(s = 2048\)、\(h = 12288\)、\(L = 96\)、\(t = 8\)、\(b = 1\)：\(sbh\) = 24 MiB，\(5as/h = 80\)。1F1B 下第一个 pipeline stage 要存 \(L\) 层的量：
| 配置 | 每层系数 × \(sbh\) | 每层 | × 96 层 |
|---|---|---|---|
| TP | 10 + 3 + 10 = 23 | 552 MiB | 51.8 GiB |
| TP + SP | (34 + 80) / 8 = 14.25 | 342 MiB | 32.1 GiB |
| TP + SP + selective recompute | 34 / 8 = 4.25 | 102 MiB | 9.6 GiB |
SP 单独省 38%。再加 selective recompute（只重算 attention 里那个 \(5as/h\) 项，多 2.7% FLOPs，见 activation recomputation），总共降到 TP 的约 1/5。
必答点
- 为什么：TP 只切 GEMM，LayerNorm、dropout、residual 的 activation 每卡一整份（\(10sbh\)）
- 怎么做：这些区域沿 \(s\) 切；row 之后的 all-reduce 换成 reduce-scatter，column 之前加 all-gather
- 通信量不变：all-reduce = reduce-scatter + all-gather
- 效果：activation 全部 /\(t\)；加 selective recompute 后约 5× 节省
追问
SP 和 context parallel / Ulysses 有什么区别？
SP 的 attention 里每卡仍看完整序列。Context parallel 连 attention 也按 \(s\) 切，要交换 K、V，见 Context Parallel。
通信次数翻倍了，会不会更慢？
字节数一样。拆开后反而更好和 GEMM 切块 overlap，Megatron Core 的 tp_comm_overlap 要求先开 SP。
来源Korthikanti et al. 2022 Abstract（p.1，5×、recompute overhead 降 90%+）、§4.1（p.4，eq. 1）、§4.2.1（p.5，eq. 2）、§4.2.2（pp.6–7，Fig. 5–6，eq. 3–4）、§4.2.3（p.8，eq. 5）、§5（pp.8–9，Table 2，GPT-3 省 70%、多 2.7% FLOPs）、Table 3（p.10）· NCCL Collective Operations（ReduceScatter + AllGather = AllReduce）· nccl-tests PERFORMANCE.md · Megatron-LM arguments.py（tp_comm_overlap 需要 SP）