Megatron 3D Parallel Rank Layout Example

本文聚焦于：在 张量并行(Tensor Parallelism, TP)、数据并行(Data Parallelism, DP) 和 流水线并行(Pipeline Parallelism, PP) 同时启用时，Megatron 将如何分配 GPU。此处以一个 16 张卡的集群为例。

具体来说：

• 总共 16 张 GPU，记为 g0 到 g15
• g0 到 g7 在第一个 node，g8 到 g15 在第二个 node
• tensor_model_parallel_size = 2
• pipeline_model_parallel_size = 4
• data_parallel_size = 2
• 这里只讨论 TP / DP / PP 三维；等价于不额外展开 CP / EP
• 下面的编号对应官方文档给出的 16 卡示例；如果修改 order，group 划分会变化

一眼看懂的 rank 布局

Megatron Core 文档在 3D 并行的 rank 生成里使用下面这个关系：

$$\color{red}\text{global\_rank} = \text{tp\_rank} + \text{dp\_rank} \times tp + \text{pp\_rank} \times tp \times dp$$

一个反直觉的结论是：DP rank 相同时，GPU id 并非是连续的；反而是 PP rank 相同时，会按顺序分配 GPU。

在本例里 tp = 2、dp = 2、pp = 4，所以 rank 可以排成下面这张表：

pp_rank \ (dp_rank, tp_rank)	(0, 0)	(0, 1)	(1, 0)	(1, 1)
0	g0	g1	g2	g3
1	g4	g5	g6	g7
2	g8	g9	g10	g11
3	g12	g13	g14	g15

看这张表时可以这样记：

• TP group：固定 pp_rank 和 dp_rank，只让 tp_rank 变化
• DP group：固定 pp_rank 和 tp_rank，只让 dp_rank 变化
• PP group：固定 dp_rank 和 tp_rank，只让 pp_rank 变化
• Model Parallel group：固定 dp_rank，同时让 tp_rank 和 pp_rank 变化

实用结论

• 看到 连续两个 rank，大概率是在同一个 TP group
• 看到 隔 2 个 rank 的一对，通常是在同一个 DP group
• 看到 隔 4 个 rank 串起来的一列，通常是在同一个 PP group
• 看到 两列 TP pair 沿 PP 方向拼起来，通常就是一个 model parallel group

注意事项

• _EMBEDDING_GROUP 只有在默认 embedding rank 规则下才是首尾 stage 成对出现；这一点在使用自定义 embedding rank 或 encoder / decoder split 时需要重新核对。
• 最容易忽略的前提是 order：官方也支持其他 rank 初始化顺序；一旦顺序变了，尤其是 DP group，编号就不再是这里这一版。

References

• 官方 API 文档直接给出了这个 16 GPU、TP=2、PP=4 时的 TP / PP / DP group 示例，并说明相邻 rank 最好落在同一台 DGX 上： Megatron Core core.parallel_state API
• 官方源码解释了 rank 生成公式、order 对 group 的影响，以及 dp_group 的构造方式： NVIDIA/Megatron-LM parallel_state.py