从一维到多维:图解 Megatron 的 SP / CP / EP 与 Transformer 唯一通信处

本文从一个反直觉的观察出发:一个 Transformer 层里,99% 的计算模块根本不需要跨 GPU 通信,唯一的"全员大会"只有 Attention。理解了这一点,Megatron-LM 那套看似复杂的 6 维并行(TP / SP / PP / DP / CP / EP)就不再是黑魔法,而是一张围绕 Attention 设计的"交通调度图"。

关联阅读:本文讲训练并行维度的拆解;VLA 场景下为何选 FSDP 而非 Megatron 见 为什么 VLA 训练选 FSDP 而不是 Megatron。

一、先看清楚:Transformer 层里到底有什么

把一个标准 Transformer 层展开,数据流是这样的(Pre-LN 结构):

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
    层输入 x [batch, seq, d]
         │
         ├────────────────────────┐ (残差)
         ▼                        │
①  LayerNorm                      │
         ▼                        │
②  Linear (QKV 投影)              │
    xW_q  xW_k  xW_v → Q, K, V   │
         ▼                        │
③  Attention   ★ 唯一的红灯区 ★    │
    score[t,j] = Q[t]·K[j]        │
    softmax → Σ w_j · V[j]        │
         ▼                        │
④  Linear (输出投影 W_o)           │
         ▼                        │
       ( + ) ←────────────────────┘
         │
         ├────────────────────────┐ (残差)
         ▼                        │
⑤  LayerNorm                      │
         ▼                        │
⑥  MLP = Linear(d→4d)             │
        → GELU                    │
        → Linear(4d→d)            │
         ▼                        │
       ( + ) ←────────────────────┘
         │
    层输出 → 重复 × N 层

现在问一个关键问题:如果把这 8 个 token 切成两半分给两张 GPU,哪些模块还能各算各的?

token-local vs token-global:一张表看懂

模块 每个词的计算 依赖其他词? 能否按序列切分
LayerNorm 对 [The] 自己的 hidden 维度归一化 否 ✅
Linear(QKV 投影等) hidden × W,纯矩阵乘 否 ✅
MLP(FFN) Linear + GELU + Linear 否 ✅
Attention [here] 的 Q 要和 [The][cat][sat] 的 K 点积 是 ❌ 必须补通信

前三个模块的共同点:计算一个 token 的输出,只需要这个 token 自己的输入向量 + 共享权重。它们像"各自批改各自的学生作业",随便分发,零通信。

Attention 则是"全班讨论":第 t 个 token 的注意力权重 score[t,j] = Q[t]·K[j] 要遍历所有 j 位置。把序列切开,GPU0 上的 [here] 就看不到 GPU1 上的 [The] 的 K/V——信息被截断了。

用一张色块图总结(绿 = token-local 可切,红 = token-global 需通信,黄 = TP 权重切分区):

1
2
3
4
5
6
7
8
┌─────────────────────────────────────────────┐
│ ① LayerNorm        ┃ 绿 ┃ SP/CP 直接切      │
│ ② Linear (QKV)     ┃ 黄 ┃ TP 按权重列切     │
│ ③ Attention        ┃ 红 ┃ CP Ring Attention │
│ ④ Linear (W_o)     ┃ 黄 ┃ TP 行切+AllReduce │
│ ⑤ LayerNorm        ┃ 绿 ┃ 直接切            │
│ ⑥ MLP              ┃ 黄 ┃ TP 切 / EP 切专家 │
└─────────────────────────────────────────────┘

整个 Transformer 层只有 ③ 一处红灯。 后面所有的并行设计,本质上都是在回答同一个问题:怎么让绿灯区随便切、把红灯区的通信代价降到最低、以及绕开或者补齐它。

二、让真实数字流过整个网络

光看色块图还不够直观,我们让一个真实的数字流跑一遍。下面所有数字都来自一个真实运行的小模型(d_model=4, seq=4, batch=2, 1 个注意力头;真实大模型是 d=4096+、seq=8K+,但逻辑一模一样)。

2.1 batch / seq / d 是什么

张量形状 [batch, seq, d],三个维度分别是:

  • batch = 2 一次喂几句话(互不相关)
  • seq = 4 每句话几个 token
  • d = 4 每个 token 用多长的向量表示
1
2
3
4
[
  [ [t0的4个数], [t1的4个数], [t2的4个数], [t3的4个数] ],   ← 句子1 "The cat sat here"
  [ [t0的4个数], [t1的4个数], [t2的4个数], [t3的4个数] ],   ← 句子2 "A dog ran fast"
]

关键直觉:矩阵的"行"是 token,"列"是特征。 后面所有讨论都围绕"哪些计算沿行独立、哪些计算要跨行"。

2.2 让 “The cat sat here” 流过整个网络

第 0 步:tokenize + embedding。 “The cat sat here” → [0, 1, 2, 3] → 查嵌入表 → x[2, 4, 4]。句子 1 的 x(每行 = 一个 token 的身份向量):

1
2
3
4
5
        d1      d2      d3      d4
The  [-1.126, -1.152, -0.251, -0.434]
cat  [ 0.849,  0.692, -0.316, -2.115]
sat  [ 0.322, -1.263,  0.350,  0.308]
here [ 0.120,  1.238,  1.117, -0.247]

从这一刻起,"cat"这个词就变成了向量 [0.849, 0.692, -0.316, -2.115],后续所有模块处理的都是这种数字,不再有"词"的概念。

① LayerNorm:每行自己归一化。 对每一行(每个 token 自己的 4 个数)求均值方差,变成均值 0、方差 1:

1
2
3
4
5
        d1      d2      d3      d4
The  [-0.954, -1.020,  1.214,  0.760]
cat  [ 0.908,  0.775, -0.079, -1.603]
sat  [ 0.571, -1.732,  0.611,  0.550]
here [-0.688,  1.072,  0.882, -1.266]

验证:每行均值≈0(1e-7 量级),每行方差≈1。注意看:算 The 那一行的归一化,用到的数字全部来自 The 自己那 4 个数——cat 行变成什么、在不在同一张 GPU 上,完全无所谓。这就是 token-local。此时如果把"The cat"放 GPU0、"sat here"放 GPU1,各算各的,结果一模一样,零通信。

② Linear(QKV 投影):整行乘同一个权重。 Q = x @ W_q,每个 token 的向量独立地乘权重矩阵:

1
2
3
4
5
6
Q(每行 = 该token的"查询"向量):
       q1      q2      q3      q4
The  [-0.236,  0.295, -0.016,  1.213]
cat  [ 0.306, -0.338,  0.567, -1.023]
sat  [-0.257, -0.641,  0.254,  0.594]
here [ 0.261,  0.551,  0.373, -0.121]

K、V 同理各有一份。The 的 Q 只由 The 的输入向量算出,依然 token-local。TP 切的是这里的 W_q(列切)——即把 4 列权重分给不同 GPU,各算 2 列,而不是切 token。

③ Attention:唯一跨行的地方。 现在算 score[t,j] = Q[t]·K[j]——注意下标:Q 的行 t 点乘 K 的行 j。第 1 行 The 的注意力分数要跟所有 4 行 K 做点积:

1
2
3
4
5
6
score 矩阵(4×4,行=查询token,列=被看的token):
          K(The)  K(cat)  K(sat)  K(here)
Q(The)   [-0.530,  0.307, -0.233, -0.120]
Q(cat)   [ 0.417, -0.282,  0.318, -0.064]
Q(sat)   [-0.341, -0.013, -0.108, -0.341]
Q(here)  [ 0.092,  0.085,  0.148,  0.077]

softmax 后(每行加权和=1):

1
2
3
4
Q(The)   [0.162, 0.375, 0.218, 0.245]   ← "The" 分给 cat 37.5% 的注意力
Q(cat)   [0.331, 0.165, 0.300, 0.205]
Q(sat)   [0.215, 0.299, 0.272, 0.215]
Q(here)  [0.248, 0.246, 0.262, 0.244]

输出 = 权重 @ V("The"的新表示 = 16.2%·V(The) + 37.5%·V(cat) + …)。看 Q(The) 那一行:它引用了 K(cat)、K(sat)、K(here)——跨行交互发生了。如果 cat/sat/here 的 K/V 在别的 GPU 上,The 的输出就是错的。这就是 CP 必须用 Ring Attention 把 K/V 沿环传一圈的原因:让每一行的 Q 都"见过"全部 K/V。

对比一下 ①② 和 ③:

  • ①② 的计算是"行内"的(每行独立);
  • ③ 的 score 矩阵是"行×行"的(Q 的行 × K 的行)——这个 4×4 矩阵就是通信需求的来源,序列越长它越大(seq=128K 时是 128K×128K,这就是 O(N²))。

④ 输出投影 + 残差。 Wo(attn_out) + h,又回到行内计算,加上残差跳线(把 ① 之前的输入直接加回来):

1
[ 0.540,  0.755,  0.163, -1.819]   ← cat 行

⑤ 第二个 LayerNorm。 同 ①,行内归一化。

⑥ MLP:升维 → 激活 → 降维。

1
2
3
4
Linear W_up: d=4 → d_ff=8   ("cat"那行变成 8 个数)
[ 0.632,  0.729,  0.258, -0.508,  0.987, -0.241, -0.315, -0.160]
       ↓ GELU(逐元素,还是行内)
       ↓ Linear W_down: 8 → 4

加残差 → 本层最终输出 [2, 4, 4]:

1
2
3
4
5
        d1      d2      d3      d4
The  [-1.069, -1.101,  1.231,  0.584]
cat  [ 0.081,  0.936, -0.100, -1.433]
sat  [ 0.433, -1.860,  0.802,  0.348]
here [-1.333,  1.452,  0.889, -1.015]

→ 送入下一层,重复 N 次(如 Llama-70B 是 80 层)。

整个过程一句话:每行(token)带着自己的 4 个数出发,在 ③ 处短暂地"看"了一眼所有其他行,其余时间都在自己行内做变换。

2.3 三个维度正好对应三种并行

现在回头看 [batch, seq, d],你会发现并行的切分刀口就落在这三个轴上:

1
2
3
4
5
6
7
8
              batch 轴(2句)
                ↓ DP 在这里切:句子1给一组GPU,句子2给另一组
┌───────────────────────────────┐
│ [The][cat][sat][here]          │ ← seq 轴:SP/CP 在这里切
│  ↓    ↓    ↓    ↓              │    "The cat"给GPU0, "sat here"给GPU1
│  各token独立的4维向量           │
│    [d1 d2 d3 d4]               │ ← d 轴:TP 在这里切
└───────────────────────────────┘    W_q 的4列分给不同GPU
维度 切哪里 用上面的例子说
DP 切 batch 句子1给 GPU 组 A,句子2给 GPU 组 B,最后同步梯度 两句话本来互不相干,怎么分都对
SP / CP 切 seq token [The][cat] 给 GPU0,[sat][here] 给 GPU1 ①②⑤⑥ 行内计算照常;只有 ③ 的 score 矩阵被切成 4 块,GPU0 的 Q(The) 需要 K(sat)K(here) → Ring Attention 环形传递
TP 切 d 每个 token 的 4 维向量(和权重矩阵的 4 列)分给多卡 GPU0 算每行的前 2 个特征,GPU1 算后 2 个,All-Reduce 合并
EP(MoE) ⑥ 的 MLP 换成 8 个专家,每卡持 2 个 [The] 被路由到专家 1(GPU0),[cat] 被路由到专家 3(GPU1),All-to-All 流动

最妙的一点:③ 的 score 矩阵形状恰好是 [seq, seq]——seq 维切分(CP)直接切这个矩阵的行和列,d 维切分(TP)则完全不碰它(切的是产生 Q/K/V 的权重)。这就是为什么 CP 管"序列长"、TP 管"模型宽",两者天然正交、可以叠加。

三、Megatron 到底有"几维"并行?

先澄清一个常见困惑:网上有人说 Megatron 是 5D,有人说是 6D,到底几个?

NVIDIA 官方口径是 5D-parallelism:TP + PP + DP + CP + EP(EP 被明确定位为"第五维度")。如果把与 TP 绑定的 SP 也算独立一维,社区习惯上称 6D。MoE 场景下还有 ETP(专家张量并行,即 DeepSeek 论文里的 ESPT)和 EDP(专家数据并行)两个扩展维度。

各维度的定位一览:

维度 切分对象 解决什么瓶颈 核心通信原语
DP batch 维度的数据 吞吐扩展 梯度 All-Reduce
TP 权重矩阵(行/列) 单层权重单卡放不下 All-Reduce
SP TP 区域外的激活(序列维) 激活显存冗余 All-Gather + Reduce-Scatter
PP 模型层(按 stage) 层数太多跨机切 stage 间点对点
CP 整条序列(含 Attention) 长上下文 O(N²) KV 环形 P2P(Ring Attention)
EP MoE 的专家集合 专家数量多 Dispatch/Combine All-to-All

初学者最容易混淆的就是 SP / CP / EP 这三个名字里都带"并行"的东西。下面逐个拆开。

四、SP:TP 的"伴生机制",切的是 TP 管不到的激活

配置:TP=2,SP=2,共 4 卡,序列 8 个 token。

SP 的出发点很朴素:TP 切权重,但 LayerNorm 和 Dropout 没有 TP 的 All-Reduce 保护,它们的激活在每张卡上都是全量冗余的。SP 把这部分激活也按序列切开:

  • TP 区域之外(LayerNorm、Dropout):激活按序列切两半,GPU0/1 持 token 1–4,GPU2/3 持 token 5–8——绿区,零通信;
  • 进入 TP 区域(QKV 投影、MLP):先 All-Gather 把序列拼完整,再按 TP 切权重计算;
  • 退出 TP 区域:原来的 All-Reduce 替换为 Reduce-Scatter,输出回到序列切分状态。

注意一个微妙的事实:SP 并没有减少总通信量——它只是把 TP 的 2 次 All-Reduce 换成等量的 All-Gather + Reduce-Scatter(数学上 All-Reduce = All-Gather + Reduce-Scatter)。真正省下的是绿区的激活显存,从"每卡全量"降到 1/N。

两个关键定位:

  • SP 与 TP 共用同一个进程组,--sequence-parallel 不是独立维度,是 TP 的开关;
  • SP 全程绕开 Attention 的跨 token 需求——因为进 TP 区域前序列已经拼回完整,Attention 看到的永远是全序列。

适用场景:TP 已经开到 8(单机 NVLink 域内),激活显存仍然紧张时打开。这就是为什么 Megatron 的大模型配方里 SP 几乎总是伴随 TP 出现。

五、CP:唯一敢"切开 Attention"的方案

配置:CP=2,序列 8K,GPU0 处理 token 1–4K,GPU1 处理 4K–8K。

  • LayerNorm / Linear / MLP:各卡独立计算,零通信(绿区属性);
  • Attention:GPU0 的 Q 需要看到 token 4K–8K 的 K/V——问题就在这里;
  • 反向传播:梯度通过 Reduce-Scatter 回到序列分片状态。

这就是 Ring Attention 登场的时刻。

Ring Attention:让 KV 沿环传一圈

核心思想:每张卡保存自己那段 K/V 分块,让 KV 块沿"环"逐站传递,传一圈后每张卡的 Q 都"见过"所有 K/V。

设 4 卡、序列 8K、每卡持 2K token 的 Q/K/V:

轮次 GPU0(Q₀) GPU1(Q₁) GPU2(Q₂) GPU3(Q₃)
初始 持 K₀V₀ 持 K₁V₁ 持 K₂V₂ 持 K₃V₃
第 1 轮 用 K₀V₀ 算部分注意力 K₁V₁ K₂V₂ K₃V₃
传递 ↓ 收 K₃V₃ 收 K₀V₀ 收 K₁V₁ 收 K₂V₂
第 2 轮 用 K₃V₃ 累积 K₀V₀ K₁V₁ K₂V₂
… (共 4 轮)
结束 Q₀ 已见过全部 K/V ✓ ✓ ✓ ✓

三要素缺一不可:

  • 分块计算:softmax 分子和分母按块累积(在线 softmax 技巧),不用一次性算出完整 N×N 注意力矩阵;
  • 环形 P2P 通信:KV 沿环单方向流动,避免全局 All-Gather 的显存峰值;
  • 计算通信重叠:GPU 在算当前 KV 块的注意力时,网络已经在预取下一块。

Megatron 的 CP 是这套算法的工业级实现:底层用环形 P2P 改写 All-Gather/Reduce-Scatter,配合 FlashAttention 内核、causal mask 的负载均衡优化。

SP 和 CP 的本质区别(面试高频)

NVIDIA 官方原文说得很清楚:“Unlike prior SP which only splits the sequence of Dropout and LayerNorm activations, CP partitions the network inputs and all activations along sequence dimension.”

SP CP
切分范围 仅 TP 区域外的激活(LayerNorm/Dropout) 整条前向序列,包括 Attention
是否独立维度 否,绑定 TP 进程组 是,独立进程组
如何处理 Attention 进 TP 区域前拼回完整序列,绕开 用 Ring Attention 穿过去
解决的瓶颈 激活显存冗余 长序列 O(N²)
典型配置 TP=8 时伴随开启 32K/128K 长上下文训练

一句话:SP 是"绕着 Attention 走",CP 是"穿过 Attention 走"。

六、EP:切专家,与序列维度完全正交

配置:EP=4,模型含 8 个专家,每卡持有 2 个专家的完整权重。

EP 面向 MoE 模型(DeepSeek-V3、Mixtral 等)。MoE 层用 Router + 多个专家 MLP 替换了上面结构图里的 ⑥——注意被替换的正是 token-local 的绿区,这也是 EP 能干净利落切分的前提。

设一个 batch 有 8 个 token,Router 做 Top-2 路由:

Token 路由到 所在 GPU
t1 专家 1, 5 GPU0(持专家 1–2)、GPU2(持专家 5–6)
t2 专家 3, 8 GPU1、GPU3
t3 专家 2, 7 GPU0、GPU3

四步流程:

  1. Router 计算:每个 token 经门控网络得到 Top-K 专家编号;
  2. Dispatch(All-to-All):GPU0 发现 t1 要去 GPU2 的专家 5,同时 GPU2 也有 token 要来——双向 All-to-All 把 hidden state 发往目标卡;
  3. 专家计算:各卡收齐发往自己专家的 token,批量执行 MLP;
  4. Combine(All-to-All 回传):结果沿原路发回来源卡,按路由权重加权合并。

EP 与 CP/SP 在维度上完全正交:前两者切序列位置,EP 切专家集合;通信原语一个是 All-Gather/环形 P2P,一个是 All-to-All。这也是为什么 DeepSeek-V3 能同时叠加 CP + EP + TP——各自负责不同模块的不同瓶颈,互不冲突(新版本 Megatron-Core 支持 MLA + CP + EP 组合,需 TransformerEngine ≥ 2.5)。

补充:当单个专家本身过大时,还可在专家内部再切一刀——即 ETP(DeepSeek 称为 ESPT,Expert-Sliced Tensor Parallelism)。MoE 层的完整并行折叠就是 ETP × EP × EDP × PP,与注意力层的 TP × CP × DP × PP 两套布局在 PP 上保持对齐。

1
2
3
注意力层布局                 MoE 层布局
TP × CP × DP × PP   ←对齐→   ETP × EP × EDP × PP
     └──── PP 共享 ────┘

七、一套直观类比

  • LayerNorm / Linear / MLP 像"各自批改各自的作业"——随便分发,零通信;
  • Attention 像"全班讨论"——必须让每个人听到所有人发言,Ring Attention 就是环形传话筒;
  • SP 是"教室内部重新分工"——只切 TP 旁边的绿区,省显存不省通信;
  • CP 是"把一本书拆给 4 个人读"——段落级切分,读到需要引用别处时通信一次(KV 环形传递);
  • EP 是"每个老师专教自己擅长的科目"——token 按需被路由到对应专家,学完回来加权汇总。

八、选型速查

症状 药方
模型大到 TP 已开满(单机 8 卡 NVLink 域),激活显存仍紧张 TP + SP
上下文长到单卡放不下注意力(32K/128K 训练、长文本 prefill) 加 CP(Ring Attention)
模型是 MoE,专家需要分布到多卡 加 EP(All-to-All dispatch/combine)
层数多,跨机切分 加 PP
都不够,吞吐还要拉满 以上叠加 + DP,如 DeepSeek-V3 同时用 CP + EP + TP(ESPT)

九、总结

回到开头那句话:一个 Transformer 层 = “绿(随便切)→ 黄(TP 切权重)→ 红(唯一要通信的注意力)→ 黄 → 绿 → 残差” 的循环。

所有的多维并行策略,本质上都是在为这三个色区做分工:

  • SP 负责绿区的序列维切分(省激活显存,绕开 Attention);
  • CP 负责红区的序列维切分(Ring Attention 穿过 Attention,解长上下文);
  • EP 负责 MoE 化的绿区(专家维切分 + token 流动);
  • TP 负责全部黄区(权重切分);
  • PP 和 DP 在层间和 batch 维度上完成最后的拼图。

理解了"Attention 是唯一需要跨 token 通信的地方"这个支点,6 维并行就不再是玄学——它们只是围绕这个支点搭起来的、各司其职的脚手架。

本文结构图与流程图均为示意,实际的 Megatron-Core 实现细节(RankGenerator、Parallel Folding 的进程组映射、causal mask 负载均衡等)可参考 NVIDIA 技术报告与开源代码。写于 2026 年 8 月,基于 Megatron-Core 新版本(支持 DeepSeek-V3 架构 MLA + CP + EP)。

参考