14. 序列并行（Sequence Parallelism）解决了什么问题？和 TP 是什么关系？

整理 Sequence Parallelism 的动机、实现机制与组合策略。

简单回答

Sequence Parallelism（SP）是把序列长度维度切分到多张 GPU 上，让每张卡只处理 sequence 的一部分。它解决的问题是 长序列训练时 activation 显存爆炸——TP 切了权重和部分 activation 但 LayerNorm、Dropout、residual 这些 element-wise 操作的 activation 仍然按完整序列长度存。SP 把这些剩余的 activation 也切分掉。SP 通常和 TP 组合使用——Megatron-LM 的 SP 是和 TP 共用 GPU，TP 切权重的同时 SP 切 element-wise 操作的 activation。还有更激进的 Ring Attention / DeepSpeed Ulysses 把 attention 计算也按序列切分，能训 100K-1M 上下文。

详细解释

为什么 TP 不够

回忆 TP 的切分方式（第 04 篇）。TP 把 attention 和 FFN 的权重矩阵按某个维度切分到多张卡上：

  • Attention 切 head 维度（每张卡负责部分 heads）
  • FFN 切 hidden dim（第一层按列切、第二层按行切）

TP 在 矩阵乘法 操作上 activation 是切分的——每张卡只算它负责的那部分。但在 element-wise 操作上 activation 是完整的：

  • LayerNorm / RMSNorm
  • Dropout
  • 残差连接
  • GELU/SiLU 激活函数

这些操作不能简单地按维度切，因为每个位置的输出依赖完整的输入向量。所以 TP 实现里这些操作是在每张卡上重复跑一份完整的 activation。

短序列时这无所谓——activation 不是瓶颈。但长序列（8K、32K、128K）时这部分 activation 显存占比急剧上升。LLaMA-70B 在 32K 序列上 activation 显存可能比模型权重还大。这时不切 sequence 就训不动。

Megatron 的 Sequence Parallelism

Megatron-LM（NVIDIA）2022 年提出的 SP 思路是：

  • 把 element-wise 操作（LayerNorm、Dropout、残差）按序列维度切分到多张卡
  • 每张卡只负责 sequence 的一部分位置（比如序列 8K，TP=8，每张卡 1K 个位置的 LayerNorm）
  • TP 的矩阵乘法仍然按权重切分

进入和退出 SP 区域要做 reshape：

[seq_len, hidden] (按 hidden 切，TP 模式)
       ↓ AllGather (按 hidden 维度还原)
[seq_len, hidden_full]
       ↓ ReduceScatter (按 seq_len 维度切分)
[seq_len_local, hidden_full] (按 seq 切，SP 模式)
       ↓ LayerNorm/Dropout (本地计算)
       ↓ AllGather (按 seq 还原)
[seq_len, hidden_full]
       ↓ ...继续 TP 矩阵乘法

在 TP 和 SP 之间切换需要额外的 AllGather 和 ReduceScatter 通信。但这部分通信总量和 TP 原本的通信差不多（数据量一样只是分布方式变了），所以总通信量没增加多少。

收益是显存。激活显存可以降到原来的 1/TP。32K 序列 70B 模型上，SP 让 activation 显存从 ~80GB 降到 ~10GB。配 gradient checkpointing 后训练 32K-128K 上下文成为可能。

Megatron SP 复用 TP 的 GPU 组——同一组 GPU 在矩阵乘时是 TP，在 element-wise 时是 SP，没有额外硬件开销。

Ring Attention / Context Parallelism

Megatron SP 只切了 element-wise 操作，attention 矩阵乘法本身还是每张卡完整跑（attention 复杂度 O(n²) 还是按完整序列算）。当上下文极长时（100K+），attention 自己就爆显存。

Ring Attention（来自 Berkeley，2023）把 attention 计算本身也按序列切分。思路是：

  • 序列切成 N 段分到 N 张 GPU 上
  • 每张卡的 Q 是它本地那段
  • K/V 在 GPU 之间环形传递——每张卡看到所有的 K/V，但是分批次传过来
  • 用 online softmax 增量算 attention，不需要存完整 attention 矩阵

通信和计算可以重叠——一段 K/V 传来时算上一段的 attention，传完算完同时进行。

理论上 Ring Attention 能让 attention 的显存按 1/N 切分，且通信开销不显著（被计算覆盖掉）。实际能训 1M-10M 上下文。Gemini 1.5 Pro 的 1M 上下文据信用了类似技术。

DeepSpeed Ulysses

Microsoft 的 DeepSpeed Ulysses 是另一种 sequence 维度并行方案，思路和 Ring Attention 不同：

  • 输入按序列维度切到 N 张 GPU（每张卡处理 1/N 序列）
  • 在 attention 计算前用 AllToAll 把数据从 seq_local, head_full 转成 seq_full, head_local——按 head 维度切
  • 在 head 切分下做 attention（类似 TP 切 head）
  • attention 后再 AllToAll 转回 sequence 切分

用 AllToAll 通信替代 Ring Attention 的环形传递，对网络拓扑要求不同。在 NVLink 域内 AllToAll 很快，跨机时可能不如 Ring。

组合策略

实际生产里 SP 和其他并行通常组合：

SP + TP + DP：Megatron 的标配。SP 和 TP 共用 GPU 组，DP 在外层。70B-540B 模型 + 32K-64K 序列的标准方案。

Ring Attention + TP + PP + DP：要训 100K 以上上下文时加上 Ring Attention 这一层。Ring Attention 维度独立于 TP/PP/DP，形成 4D 并行。

Context Parallelism + Tensor Parallelism：NVIDIA 的 Megatron 在 2024 年加入 CP，类似 Ring Attention 的思路集成进 Megatron 框架。

显存收益的实际量级

举个具体数字：70B 模型，TP=8，DP=8，micro-batch=1，BF16 训练。

  • 不开 SP：32K 序列 activation ≈ 80 GB / GPU（单卡 80GB H100 几乎用完）
  • 开 SP：activation ≈ 10 GB / GPU，留出大量空间给更长序列或更大 micro-batch

128K 序列上：

  • 不开 SP：activation 可能 320 GB，根本训不了
  • 开 SP：40 GB，训得了
  • 开 SP + Ring Attention：< 10 GB，训得很轻松

显存优势是定量级的，不是边际优化。

工程上的注意点

Dropout 的随机数同步。SP 下 dropout 在不同 GPU 上独立计算，但要保证训练可复现就需要每张卡看到的 dropout mask 在 sequence 维度上拼起来等价于完整 dropout。要专门同步随机种子。

Position Embedding 处理。每张卡只看到部分位置，position embedding 要按本地位置范围生成。RoPE 这类相对位置编码自然兼容，绝对位置编码要小心 offset。

Causal Mask 在切分下的处理。Decoder-only 模型的因果掩码在序列切分后还要保证每段只能看到自己之前的位置。Ring Attention 里 K/V 环形传递时要按"环位置"决定要不要 mask。

Loss 的归约。Loss 是序列上 token 的平均，SP 下每张卡算自己那段的 loss 再 AllReduce 平均。要小心权重——按 token 数加权而不是按 GPU 数。

适用边界

SP 不是万能的。短序列（< 4K）用 SP 收益很小但通信开销固定，反而拖慢。所以一般有阈值——序列长度超过某个值才开 SP。

序列长但 batch size 也大的场景，SP 的相对收益小（activation 总量已经被分到 batch 维度上）。SP 主要价值在 长序列 + 小 batch 场景，这恰好是大模型预训练的典型工况。

面试时可以这样答

Sequence Parallelism 解决长序列训练时 activation 显存爆炸的问题。TP 切了权重和矩阵乘法的 activation，但 element-wise 操作 LayerNorm、Dropout、残差 这些每张卡仍然存完整序列的 activation。长序列时这部分占比很大。

Megatron 的 SP 把 element-wise 操作按序列维度切到 TP 同一组 GPU 上。TP 和 SP 之间切换通过 AllGather 和 ReduceScatter 重新分布 activation。通信总量基本不变，但激活显存能降到 1/TP。70B 模型 32K 序列 activation 从 80GB 降到 10GB。

更激进的是 Ring Attention 和 DeepSpeed Ulysses，把 attention 计算本身也按序列切分。Ring Attention 用环形传递 K/V 配合 online softmax，能训 1M 上下文。Gemini 1.5 的长上下文据信用了类似技术。Ulysses 用 AllToAll 把 sequence 切分转成 head 切分做 attention，对拓扑要求不同。

实际配置上 SP + TP + DP 是 70B-540B 模型 + 32K-64K 序列的标准。要训 100K+ 再加 Ring/Context Parallelism 形成 4D 并行。

工程要点包括 Dropout 随机数同步、Position embedding 按本地位置生成、Causal mask 在切分下要保证因果性、Loss 按 token 加权 AllReduce。

适用边界是长序列小 batch——短序列开 SP 收益小但通信开销固定反而拖慢，序列长 batch 也大的场景 SP 相对收益小。SP 主要价值在大模型预训练这种长序列小 batch 的工况。

常见追问

  1. Ring Attention 的环形通信和计算怎么重叠？理论上能完全 overlap 吗？
  2. SP 配 FlashAttention 时 attention kernel 怎么改造？
  3. SP 在多模态训练（图像 token + 文本 token）下的特殊处理是什么？
13. 3D 并行（DP + TP + PP）怎么组合？为什么大模型训练要把三种并行叠在一起？
15. MoE 的专家并行（Expert Parallelism）怎么做？为什么 MoE 训练特别难？