01
什么是 Rank
Rank 就是每个 GPU 进程的全局编号。启动分布式训练时，每个 GPU 运行一个独立进程，torch.distributed 给每个进程分配唯一整数 ID，从 0 到 world_size - 1。
但一个 GPU 同时参与多种并行（TP、PP、DP 等），在每种并行维度的通信组内，它又有自己的 local rank。
| 概念 | 含义 | 例子 |
|---|---|---|
| Global Rank | 全局唯一编号，0 ~ world_size-1 | 16 卡中第 7 号 → rank=7 |
| Local Rank | 在某个通信组内的编号 | TP 组 [g4, g5] 中，g5 的 tp_rank=1 |
同一个 GPU（如 global rank=5）可以同时有 tp_rank=1、pp_rank=0、dp_rank=1。这些通信组内的 GPU 通过集合通信操作（AllReduce、AllGather 等）交换数据——具体操作详见 集合通信操作详解。
02
Rank 映射公式
Megatron 默认排列顺序为 tp-cp-ep-dp-pp，含义是 TP 变化最快（最内层），PP 变化最慢（最外层）。完整公式包含 5 个维度：
parallel_state.py — Rank 映射公式（完整 5 维）
L282-286
# order = "tp-cp-ep-dp-pp"（默认）
global_rank = tp_rank
+ cp_rank * tp
+ ep_rank * tp * cp
+ dp_rank * tp * cp * ep
+ pp_rank * tp * cp * ep * dp
# 各维度取值范围：
# tp_rank ∈ [0, tp_size) Tensor Parallelism
# cp_rank ∈ [0, cp_size) Context Parallelism
# ep_rank ∈ [0, ep_size) Expert Parallelism
# dp_rank ∈ [0, dp_size) Data Parallelism
# pp_rank ∈ [0, pp_size) Pipeline Parallelism
#
# 其中 world_size = tp × cp × ep × dp × pp这本质上是一个多维数组的线性索引：把 (tp, cp, ep, dp, pp) 五维坐标展开为一维整数。谁在最内层（步长=1），谁变化就最快、编号就最相邻。
在下面的示例中，我们先用 TP/DP/PP 三维建立直觉（cp=1, ep=1 时公式退化为三维），然后在第 05 节扩展到 CP 和 EP。
03
具体示例：16 GPU, TP=2, PP=4, DP=2
| PP stage 0 | PP stage 1 | PP stage 2 | PP stage 3 | |
|---|---|---|---|---|
| DP=0 | g0 g1← TP → | g4 g5← TP → | g8 g9← TP → | g12 g13← TP → |
| DP=1 | g2 g3← TP → | g6 g7← TP → | g10 g11← TP → | g14 g15← TP → |
由此生成三种通信组：
TP 组（切分同一层的参数，通信量最大）
[g0,g1] [g2,g3] [g4,g5] [g6,g7] [g8,g9] [g10,g11] [g12,g13] [g14,g15]
→ 相邻 rank 在同一 TP 组（步长=1）
PP 组（组成流水线，通信量小）
[g0,g4,g8,g12] [g1,g5,g9,g13] [g2,g6,g10,g14] [g3,g7,g11,g15]
→ 步长 = tp_size × dp_size = 4
DP 组（相同模型处理不同数据，梯度同步）
[g0,g2] [g1,g3] [g4,g6] [g5,g7] [g8,g10] [g9,g11] [g12,g14] [g13,g15]
→ 步长 = tp_size = 2
看懂分组的前提：每个 GPU 到底持有什么？
以一个 24 层 GPT 模型（hidden_size=4096）为例：
- PP=4：24 层切成 4 段（stage），每段 6 层
- TP=2：每层的权重矩阵 [4096×4096] 按行切成上下两半 [2048×4096]
- DP=2：整个切好的模型复制 2 份，各自处理不同的 mini-batch
| PP stage 0 层 0-5 | PP stage 1 层 6-11 | PP stage 2 层 12-17 | PP stage 3 层 18-23 | |
|---|---|---|---|---|
| 数据批次 A (DP=0) |
g0 上半权重 g1 下半权重 |
g4 上半权重 g5 下半权重 |
g8 上半权重 g9 下半权重 |
g12 上半权重 g13 下半权重 |
| 数据批次 B (DP=1) |
g2 上半权重 g3 下半权重 |
g6 上半权重 g7 下半权重 |
g10 上半权重 g11 下半权重 |
g14 上半权重 g15 下半权重 |
现在回头看三种通信组，每个组里的 GPU 之间为什么需要通信就很清楚了：
三种通信组的具体数据流 展开看详细示例
TP 组 [g0, g1]：同一层的上下半合并（以 RowParallelLinear 为例）
完整计算：\(Y = X \times W\)，其中 \(X=[1,2,3,4]\)，\(W\) 是 4×2 矩阵，完整结果 \(Y = [30, 70]\)。
TP=2 时 W 按行切成上下两半，X 也对应切开，每个 GPU 算一项，AllReduce 求和还原完整结果：
| g0 上半权重 | g1 下半权重 | |
|---|---|---|
| X_上 = [1, 2] W_上 = [[1,5],[2,6]] |
← TP → | X_下 = [3, 4] W_下 = [[3,7],[4,8]] |
| Y_0 = [5, 17] | AllReduce SUM ↓ |
Y_1 = [25, 53] |
Y = Y_0 + Y_1 = [30, 70] ✓
数学原理
$$Y = X \cdot W = \begin{bmatrix} X_\text{上} & X_\text{下} \end{bmatrix} \begin{bmatrix} W_\text{上} \\ W_\text{下} \end{bmatrix} = X_\text{上} W_\text{上} + X_\text{下} W_\text{下}$$
每个 GPU 算一项，AllReduce 求和就是完整结果。不通信的后果：g0 只有 [5,17]，g1 只有 [25,53]，都不是正确的 [30,70]。
PP 组 [g0, g4, g8, g12]：流水线的相邻 stage 间传递激活值
flowchart LR
subgraph G0["g0 — 层 0-5"]
i0["输入 X
算层 0-5
输出 a5"] end subgraph G4["g4 — 层 6-11"] i4["收到 a5
算层 6-11
输出 a11"] end subgraph G8["g8 — 层 12-17"] i8["收到 a11
算层 12-17
输出 a17"] end subgraph G12["g12 — 层 18-23"] i12["收到 a17
算层 18-23
输出 Loss"] end i0 -- "P2P Send
激活值 a5" --> i4 i4 -- "P2P Send
激活值 a11" --> i8 i8 -- "P2P Send
激活值 a17" --> i12
算层 0-5
输出 a5"] end subgraph G4["g4 — 层 6-11"] i4["收到 a5
算层 6-11
输出 a11"] end subgraph G8["g8 — 层 12-17"] i8["收到 a11
算层 12-17
输出 a17"] end subgraph G12["g12 — 层 18-23"] i12["收到 a17
算层 18-23
输出 Loss"] end i0 -- "P2P Send
激活值 a5" --> i4 i4 -- "P2P Send
激活值 a11" --> i8 i8 -- "P2P Send
激活值 a17" --> i12
只需 P2P（点对点），不需要 AllReduce。每次只传一个 micro-batch 的激活值，数据量 = batch × seq_len × hidden，远小于参数量。
DP 组 [g0, g2]：相同模型+不同数据 → 梯度求平均
| g0 数据批次 A | g2 数据批次 B | |
|---|---|---|
| 持有层 0-5 上半权重 用批次 A 做前向+反向 |
完全相同 的参数 |
持有层 0-5 上半权重 用批次 B 做前向+反向 |
| ∇W_A | ReduceScatter ↓ |
∇W_B |
平均梯度 = (∇W_A + ∇W_B) / 2 → 参数始终保持一致
以 g0 的视角串联：一次前向 + 反向经历了哪些通信？
| 阶段 | 通信操作 | 通信组 | 通信对象 | 目的 |
|---|---|---|---|---|
| 前向：算层 0 | AllReduce | TP [g0,g1] | g1 | 合并上下半的部分输出 |
| 前向：算层 1-5 | AllReduce ×5 | TP [g0,g1] | g1 | 每层都要合并 |
| 前向：层 5 算完 | P2P Send | PP [g0,g4,g8,g12] | g4 | 把激活值传给下一 stage |
| 反向：收到梯度 | P2P Recv | PP [g0,g4,g8,g12] | g4 | 从下一 stage 收梯度 |
| 反向：算层 5-0 | AllReduce ×6 | TP [g0,g1] | g1 | 每层的梯度也要合并 |
| 反向：梯度算完 | ReduceScatter | DP [g0,g2] | g2 | 跨数据副本聚合梯度 |
04
为什么 TP 在最内层
代码注释明确写了：adjacent ranks are on the same DGX box。
物理拓扑与并行维度的对应
机器 A (8 卡 NVLink): g0 g1 g2 g3 g4 g5 g6 g7
机器 B (8 卡 NVLink): g8 g9 g10 g11 g12 g13 g14 g15
TP 组 [g0,g1] → 同机器，NVLink（900 GB/s） ✓ 通信量大，需要快
DP 组 [g0,g2] → 同机器，NVLink（900 GB/s） ✓ 梯度同步可以接受
PP 组 [g0,g4,g8,g12] → g0→g4 同机，g4→g8 跨机 IB（400 Gb/s）
✓ 只传激活值，量小
设计原则
通信量大的并行维度 → 放内层 → 物理上相邻 → 走高速互联。TP 每层前向/反向都要 all-reduce，通信最频繁；PP 每个 micro-batch 只传一次激活值，通信最少。
05
Context Parallelism 与 Expert Parallelism
前面 3 节用 TP/DP/PP 三维建立了直觉。实际上 Megatron 的默认排列是 5 维：tp-cp-ep-dp-pp。CP 和 EP 夹在 TP 和 DP 之间，各自解决不同的问题。
| 维度 | 切什么 | 解决什么问题 | 通信原语 |
|---|---|---|---|
| TP | 权重矩阵（行/列切分） | 单层参数太大，一张卡放不下 | AllReduce / AllGather |
| CP | 序列长度（sequence 维度） | 长序列的 Attention 显存爆炸 | AllGather KV（前向）、ReduceScatter（反向） |
| EP | MoE 专家（每卡持有部分 expert） | Expert 数量太多，一张卡放不下 | AllToAll（token 分发/收集） |
| DP | 数据（不同 mini-batch） | 加速吞吐量 | ReduceScatter / AllGather（梯度） |
| PP | 模型层（stage 划分） | 层数太多，一张卡放不下 | P2P Send/Recv（激活值） |
5.1 Context Parallelism (CP)：切分序列长度
问题：训练长序列（如 128K tokens）时，Attention 的 KV cache 显存随序列长度平方增长。一张卡的显存不够存完整 KV。
方案：将序列切成 cp_size 段，每个 CP rank 只处理一段。Attention 计算时，Q 留在本地，K/V 通过 AllGather 从所有 CP rank 收集完整副本。
CP=4 时，一条 8192 token 的序列被切成 4 段
原始序列: [t0 t1 t2 ... t2047 | t2048 ... t4095 | t4096 ... t6143 | t6144 ... t8191]
CP rank 0 CP rank 1 CP rank 2 CP rank 3
Attention 计算时：
Q: 每个 rank 只有自己那段的 Q → 不需要通信
K: 需要完整序列的 K → AllGather 从所有 CP rank 收集
V: 需要完整序列的 V → AllGather 从所有 CP rank 收集
反向传播时：
dK, dV: 每个 rank 算出完整的 dK/dV → ReduceScatter 只保留自己那段的梯度
CP 不影响权重
CP 只切分序列，不切分权重矩阵。这意味着同一个 CP 组内的 rank 持有完全相同的权重副本。所以权重梯度需要在 CP rank 间做 AllReduce——Megatron 的做法是把 CP rank 合并到 DP 组中（dp-cp 联合组），一起做梯度同步。
CP 通信组示例：8 GPU, TP=2, CP=2, DP=2, PP=1
rank 排列（order = tp-cp-ep-dp-pp）：
TP=0 TP=1 TP=0 TP=1
CP=0: g0 g1 g4 g5 ← DP=0, DP=1
CP=1: g2 g3 g6 g7 ← DP=0, DP=1
CP 组（切分序列的 rank，需要 AllGather KV）：
[g0, g2] [g1, g3] [g4, g6] [g5, g7]
→ 步长 = tp_size = 2
DP-CP 联合组（权重梯度同步）：
[g0, g2, g4, g6] [g1, g3, g5, g7]
→ 同时跨 CP 和 DP 两个维度做 AllReduce
5.2 Expert Parallelism (EP)：分发 MoE 专家
问题：MoE（Mixture of Experts）模型有大量专家（如 64 个），每个专家都是一个完整的 FFN。所有专家放在一张卡上显存不够。
方案：将专家分配到 ep_size 个 rank 上，每个 rank 只持有 num_experts / ep_size 个专家。Router 决定每个 token 去哪个专家后，通过 AllToAll 把 token 发送到对应专家所在的 rank。
EP=4, 8 个专家（每个 rank 持有 2 个专家）
EP rank 0 EP rank 1 EP rank 2 EP rank 3
Expert 0, 1 Expert 2, 3 Expert 4, 5 Expert 6, 7
Step 1: Router 对每个 token 打分，决定去哪个 expert
token A → Expert 3 (在 rank 1)
token B → Expert 7 (在 rank 3)
token C → Expert 0 (在 rank 0，本地！)
Step 2: AllToAll — 每个 rank 把 token 发给目标 expert 所在的 rank
┌──────────────────────────────────────────────────┐
│ rank 0 ──→ rank 1: token A │
│ rank 0 ──→ rank 3: token B │
│ rank 1 ──→ rank 0: token D (某个分配到 Expert 0 的)│
│ ... │
└──────────────────────────────────────────────────┘
Step 3: 每个 rank 对收到的 token 执行本地 expert 计算
Step 4: AllToAll（反向） — 把计算结果发回原始 rank
EP 会"吃掉"DP
EP 的 rank 是从 DP rank 中分出来的。如果原始 DP=8，开启 EP=4 后，Expert DP 只剩 8/4=2。这意味着 expert 参数的数据并行度降低了，每个 expert 只有 2 份副本参与梯度同步，而非 8 份。公式：
expert_dp_size = world_size / (tp × ep × pp)
EP 通信组示例：8 GPU, TP=2, EP=2, DP=2, PP=1
rank 排列（order = tp-cp-ep-dp-pp, cp=1）：
TP=0 TP=1 TP=0 TP=1
EP=0: g0 g1 g4 g5 ← DP=0, DP=1
EP=1: g2 g3 g6 g7 ← DP=0, DP=1
EP 组（持有不同 expert 的 rank，需要 AllToAll 交换 token）：
[g0, g2] [g1, g3] [g4, g6] [g5, g7]
→ 步长 = tp_size = 2
Expert DP 组（持有相同 expert 的 rank，梯度同步）：
[g0, g4] [g1, g5] [g2, g6] [g3, g7]
→ Expert DP size = 2（比完整 DP=4 小了 EP=2 倍）
TP-EP 联合组（AllToAll 通信的实际范围）：
[g0, g1, g2, g3] [g4, g5, g6, g7]
→ MoE 的 AllToAll 在 TP-EP 范围内执行
5.3 CP 与 EP 在同一 RankGenerator 中互斥
Megatron 的 RankGenerator.__init__ 中有一个断言：assert ep == 1 or cp == 1。这不是数学上无法同时大于 1，而是因为 Transformer 中没有任何组件同时需要 CP 和 EP——它们作用于完全不同的计算。
一个 MoE Transformer Layer 中，CP 和 EP 作用于不同组件
Self-Attention ← decoder_rank_generator (cp=2, ep=1)
├── Q·Kᵀ 是跨位置运算：pos 0 要看 pos 1,2,...,N 的 K,V
├── 序列越长 KV cache 越大 → 需要 CP 切分序列
└── 没有 expert → EP 无意义
│
↓
MoE FFN (Expert 层) ← expert_decoder_rank_generator (cp=1, ep=2)
├── Router 逐 token 打分，token 独立通过各自的 expert
├── token 之间无交互 → CP 切分序列无任何收益
└── 64 个 expert 放不下一张卡 → 需要 EP 分配
为什么 MoE FFN 不需要 CP？
CP 的价值在于切分 Attention 的序列维度——因为 Q·Kᵀ 中每个位置要与所有其他位置计算注意力，KV cache 随序列长度增长。而 FFN/Expert 是逐 token 独立计算的：token A 选了 Expert 3，token B 选了 Expert 7，它们各自独立通过自己的 expert。token 之间没有交互，切分序列不会减少任何通信或显存开销。
如果强行在一个 RankGenerator 中同时设 CP>1 和 EP>1 会怎样？
回顾 rank 公式：world_size = tp × cp × ep × dp × pp，dp 是被其余维度"挤"出来的。同时开启 CP 和 EP 会多占一个维度，白白压缩 DP：
| 方案（8 GPU, TP=2, PP=1） | CP | EP | DP | 效果 |
|---|---|---|---|---|
decoder_rank_generator |
2 | 1 | 2 | Attention 用 CP 切序列，保留 DP=2 |
expert_decoder_rank_generator |
1 | 2 | 2 | MoE 用 EP 分 expert，保留 DP=2 |
| 假设合并到一个 RankGenerator | 2 | 2 | 1 | DP 被压到 1，丧失数据并行 |
合并的代价是双重浪费：EP 组在 Attention 计算时完全用不上，CP 组在 MoE FFN 计算时完全用不上——它们消耗了 DP 维度却没有提供任何收益。分成两套 RankGenerator，同一组 GPU 在不同计算阶段使用不同的通信组，两个阶段都保留了完整的数据并行度。
5.4 五维并行的通信量对比
| 维度 | 通信操作 | 频率 | 数据量 | 理想互联 |
|---|---|---|---|---|
| TP | AllReduce | 每层前向+反向各 1 次 | batch × seq × hidden | NVLink |
| CP | AllGather + ReduceScatter | 每个 Attention 层 1 次 | batch × seq × head_dim × num_kv_heads | NVLink / NVSwitch |
| EP | AllToAll | 每个 MoE 层前向+反向各 1 次 | 取决于 token 分配，通常 < TP | NVLink / IB |
| DP | ReduceScatter + AllGather | 每步结束 1 次 | 全部参数梯度 | IB / RoCE |
| PP | P2P Send/Recv | 每个 microbatch 1 次 | batch × seq × hidden（激活值） | IB（量小，延迟更重要） |
06
RankGenerator：如何生成通信组
parallel_state.py — RankGenerator 核心接口
L456-531
class RankGenerator(object):
def __init__(self, tp, ep, dp, pp, cp, order, rank_offset=0):
self.world_size = tp * dp * pp * cp * ep
# 按 order 字符串确定各维度的排列顺序
self.ordered_size = [self.name_to_size[token] for token in order.split("-")]
def get_ranks(self, token):
"""通过 token 获取对应的通信组列表。
token 可以是 'tp', 'pp', 'dp', 也可以是 'tp-dp' 这样的组合。
"""
mask = self.get_mask(self.order, token)
return generate_masked_orthogonal_rank_groups(
self.world_size, self.ordered_size, mask
)工作原理：用 mask 指定哪些维度"在组内变化"，哪些"在组间变化"。
parallel_state.py — 初始化时创建各种通信组
L776-784, L1018-1050
# ---- Decoder RankGenerator：用于 Attention + 非 MoE 层 ----
decoder_rank_generator = RankGenerator(
tp=tensor_model_parallel_size,
ep=1, # Attention 不需要 EP
dp=data_parallel_size,
pp=pipeline_model_parallel_size,
cp=context_parallel_size, # CP 在这里生效
order=order, # 默认 "tp-cp-ep-dp-pp"
)
# 用 get_ranks() 生成各种通信组
for ranks in decoder_rank_generator.get_ranks('tp'): # TP 组
group = create_group(ranks, ...)
for ranks in decoder_rank_generator.get_ranks('pp'): # PP 组
group = create_group(ranks, ...)
for ranks in decoder_rank_generator.get_ranks('dp'): # DP 组
group = create_group(ranks, ...)
for ranks in decoder_rank_generator.get_ranks('dp-cp'): # DP-CP 联合组
group = create_group(ranks, ...)
for ranks in decoder_rank_generator.get_ranks('cp'): # CP 组
group = create_group(ranks, ...)
# ---- Expert RankGenerator：用于 MoE 层 ----
expert_decoder_rank_generator = RankGenerator(
tp=expert_tensor_parallel_size, # 可以与 TP 不同
ep=expert_model_parallel_size, # EP 在这里生效
dp=expert_data_parallel_size, # = world / (tp_e × ep × pp)
pp=pipeline_model_parallel_size,
cp=1, # MoE 不需要 CP
order=order,
)
for ranks in expert_decoder_rank_generator.get_ranks('ep'): # EP 组
group = create_group(ranks, ...)
for ranks in expert_decoder_rank_generator.get_ranks('dp'): # Expert DP 组
group = create_group(ranks, ...)
for ranks in expert_decoder_rank_generator.get_ranks('tp-ep'): # TP-EP 联合组
group = create_group(ranks, ...)组合组的用途
get_ranks('dp-cp')— 权重梯度的 AllReduce。因为 CP 不切分权重，CP rank 和 DP rank 都需要同步梯度。get_ranks('tp-ep')— MoE 的 AllToAll 通信范围。token 在 TP-EP 组内分发到目标 expert。get_ranks('tp-dp')— Embedding 权重的梯度 AllReduce，需要在所有持有相同 embedding 副本的 rank 间进行。
07
总结
flowchart TD
WS["world_size = tp × cp × ep × dp × pp"]
WS --> RANK["Global Rank = 多维坐标的线性展开"]
RANK --> TP["TP 通信组
步长=1, 最内层
NVLink
AllReduce 每层"] RANK --> CP_G["CP 通信组
步长=tp
NVLink
AllGather KV"] RANK --> EP_G["EP 通信组
步长=tp×cp
NVLink/IB
AllToAll token"] RANK --> DP_G["DP 通信组
步长=tp×cp×ep
IB
梯度 AllReduce"] RANK --> PP_G["PP 通信组
步长=tp×cp×ep×dp
IB
激活值 P2P"]
步长=1, 最内层
NVLink
AllReduce 每层"] RANK --> CP_G["CP 通信组
步长=tp
NVLink
AllGather KV"] RANK --> EP_G["EP 通信组
步长=tp×cp
NVLink/IB
AllToAll token"] RANK --> DP_G["DP 通信组
步长=tp×cp×ep
IB
梯度 AllReduce"] RANK --> PP_G["PP 通信组
步长=tp×cp×ep×dp
IB
激活值 P2P"]
下一步阅读
- 集合通信操作详解 — AllReduce / AllGather / ReduceScatter / AllToAll / P2P 等通信原语在 Megatron 中的具体用法
- pretrain_gpt.py 源码精读 — 理解数据如何按 PP/TP/CP 维度切分
- training.py 初始化篇 —
initialize_megatron()中如何调用mpu.initialize_model_parallel()创建这些通信组