Device Mesh：把 3D/4D 并行的 rank 算术收进一张网格

先建立共同认知：每个 rank 同时属于多个 group，错一个坐标就会把不该通信的卡配到一起。

SC21 把大模型训练拆成三个方向：data parallelism（DP）复制模型、切分 batch；pipeline model parallelism（PP）把层切成阶段；tensor model parallelism（TP）把层内矩阵切到多张卡。论文把它们组合成 PTD-P，并要求 p · t · d = n（§3.1）。这三项已经需要三套 group；长上下文训练再加 context parallelism（CP），把序列切成块，就变成四维。

四维不是四个独立开关。一个 rank 在 DP 组里要和同模型副本交换梯度，在 TP 组里要对层内激活做 all-reduce，在 PP 组里要把激活发给上下游，在 CP 组里要和序列块邻居交换 KV。四套 group 的成员集合不同，但都从同一个 global rank 解出来；只要 rank 顺序或维度顺序有一处不一致，通信就会安静地配错对象。

换个说法：k 个维度就有 k 套 group 家族；每套又有 n / size 个子组。rank 列表的数量随维度线性增长，也随 world_size 线性增长，而每张卡都要在每个家族里被分对一次。

前三项来自 SC21 §3.1 的 p · t · d = n；CP 是后加维度，Megatron-Core 把它乘进 model_size（megatron/core/parallel_state.py:883-895）。

所以问题不是「能不能建 group」——new_group 当然能建。问题是：当维度从 2 变成 4，谁来保证每个 rank 在每个维度上都被分到正确的组？谁来给消费者一个「只看 DP」或「只看 TP」的子网格？这就是 02 要说的旧方案痛点，也是 03 的 mesh 抽象要接住的东西。

已有方案能建出 group，但每个消费者都要自己重算 rank；维度一多，重复和错配就来了。

三者的共同缺口不是建组能力，而是表示：一个 group 属于哪一维、当前 rank 的坐标是什么、怎样只取某一维的子网格。没有这层表示，每个消费者都要维护一份 rank 算术；新增 CP 或 EP 时，所有消费者都得跟着改。DeviceMesh 的做法是把这份算术收进一个对象；Megatron 的做法是把它收进一个 order 字符串和生成器。两者都不是免费午餐，但都把「维度」变成了显式参数。

本页把 DeviceMesh 当作 PyTorch 侧的抽象来读；Megatron-Core 70f7a26 的 parallel_state.py 建组路径没有调用 DeviceMesh，它用 RankGenerator + create_group 达到同样的正交分组；其 FSDP/DTensor 适配层则使用 DeviceMesh.from_group（megatron/core/distributed/torch_fully_sharded_data_parallel.py:81）。后面的对比会一直区分这两条建组路径。

读完这节你应该能从一个 global rank 反推出坐标，并说出它在每个维度上和谁同组。

DeviceMesh 的 docstring 说得很直接：mesh 是一个 n 维数组，数组里的每个值是默认 process group 的 global rank；每个维度上的通信可以分开进行（torch/distributed/device_mesh.py:188-196）。它遵循 SPMD：所有 rank 必须拿到同一张 mesh，不一致的 mesh 会静默 hang（同文件 203-207）。

PyTorch 文档里的 2×4 例子（torch/distributed/device_mesh.py:228-242）：第 0 维跨主机，第 1 维机内。沿第 0 维 reduce 会配对 (0,4), (1,5), ...；沿第 1 维 reduce 会配对 (0,1,2,3) 与 (4,5,6,7)。

从 shape 到对象：init_device_mesh

init_device_mesh(device_type, mesh_shape, mesh_dim_names=...) 用 mesh_shape 的乘积生成长度为该乘积的 rank map：torch.arange(layout.numel())（torch/distributed/device_mesh.py:1720-1725）。这就是行主序的单位映射：rank 0 在坐标 (0,0,...,0)，最后一个 rank 在最后一格。维名不得重复，且长度与 shape 一致（同文件 1691-1702）。

一旦有了对象，常用接口是：size(dim) 取某一维大小（1289-1292）；get_local_rank(dim) 取当前 rank 在该维 group 内的序号（1308-1326）；get_coordinate() 取相对所有维的坐标（1371-1376）；mesh["tp"] 切出 1D 子网格（770-789）；get_group("tp") 取当前 rank 在该维的 process group（846-899）。当 ndim > 1 时，get_group() 必须给维度，否则报错（861-867）。

group 是怎么生成的：固定其余维，只变一维

_init_process_groups 对每个维度调用一次 _init_one_process_group（665-694）。后者把该维的子布局重映射成 rank 张量 pg_ranks_by_dim，再对每个子组调用 new_group；每个 rank 都会为所有子组进入调用，但只保存包含自己的那个 group（547-663）。这就是为什么创建 group 是集合操作，不能只让组内 rank 参与。

3D/4D 布局：坐标是 rank 的分解

以行主序的 DP-TP-PP-CP 为例，shape 为 (dp, tp, pp, cp) 时：

这是页面交互采用的简化布局；步长依次是 TP·PP·CP、PP·CP、CP、1。Megatron 的真实 rank 顺序由 order 字符串决定，函数签名默认是 tp-cp-ep-dp-pp（megatron/core/parallel_state.py:645），docstring 里的示例公式又是 tp-dp-pp 顺序（同文件 293-309）。换顺序就换步长，不能把这里的公式直接套到 Megatron 上。

3D 是 CP=1 的特例：rank = dp·(TP·PP) + tp·PP + pp。SC21 的 1T 配置是 TP=8、PP=64，所以 DP 每跨一格，rank 跳 8×64；这不是 Megatron 的默认顺序，只是把 3072 个 rank 摊成一张三维表的其中一种方式（Table 1）。

Megatron 的配置默认值也在同一层：TransformerConfig 继承 ModelParallelConfig（transformer_config.py:61），tensor_model_parallel_size、pipeline_model_parallel_size、context_parallel_size 默认都是 1（model_parallel_config.py:59, 112, 129）。没有 data_parallel_size 字段；DP 由 world_size / (TP×PP×CP×gtp_remat) 推出，默认 gtp_remat=1（megatron/core/parallel_state.py:883-895）。

4D 之后，为什么还要 mesh 这层抽象

所以「4D 需要 mesh 抽象」的准确说法是：4D 需要一个显式的维度表示。它可以是 PyTorch 的 DeviceMesh 对象，也可以是 Megatron 的 RankGenerator + order 字符串；关键是消费者不再各自维护一套隐式 rank 算术。抽象不减少通信，只减少「把 group 配错」和「加一维时改所有调用点」的成本。

同一张网格上，DP/TP/PP/CP 只是固定另外三维、变化自己这一维。

把 rank 写成坐标 (dp, tp, pp, cp) 后，某一维的 group 就是「其余坐标不变，只让这一维取遍 0..size-1」得到的 rank 集合。组大小等于该维大小，组数等于 world_size / size。下面这张表是页面整理的四个维度；右侧锚点给出 Megatron-Core 里对应的建组循环。

§04 · 四维 group 家族（页面整理；源码锚点见各行）

表中 parallel_state.py 均指 megatron/core/parallel_state.py。EP 不在上面那张 dense 4D 网格里：Megatron 用单独的 expert RankGenerator（megatron/core/parallel_state.py:950-976），并且同一个 RankGenerator 不允许 EP 与 CP 同时大于 1（同文件 508-509）。所以 MoE 的“5D”不是一张均匀 mesh，而是 dense grid 与 expert grid 两套分组推断。

CP 那一行有一个容易忽略的细节：CP rank 持有不同的序列块，但权重是复制的；因此权重梯度也要在 CP 之间求和。Megatron 的做法是把 CP 并进 DP 的权重梯度 all-reduce 组，而不是单独再做一次（megatron/core/parallel_state.py:717-722）。TP 的组内通信则是层内激活的 all-reduce：一个 Transformer 层前向 2 次、反向 2 次（TP 论文 §3）；SC21 §3.2 把每层每卡的 TP 通信量记为 8bsh((t−1)/t)。

顺序不是流水账；每一步的通信落在不同的 group 家族上。

一个训练步可以按「数据 → 层间 → 层内 → 序列」读。注意这只是一种展开顺序：实际执行中 PP 的 microbatch、TP 的 all-reduce 和 CP 的 KV 交换会重叠，不会严格串行。

反向时顺序反过来：TP 的 all-reduce 在反向再做两次（TP §3）；PP 反向发梯度；DP 在 batch 末 all-reduce 梯度；CP 的权重梯度并入 DP 组（megatron/core/parallel_state.py:717-722）。

§05 · 训练步通信账（页面整理；parallel_state.py 指 megatron/core/parallel_state.py）

PP 的气泡：p 是阶段数，m 是每 pipeline 的 microbatch 数（SC21 §2.2.1；GPipe §2.3 写作 O((K−1)/(M+K−1))）。气泡是 PP 的固定成本，和 mesh 形状无关。

先看配置如何落在网格上，再看吞吐；这是弱扩展实验，不是同模型强扩展。

以下数字都来自 SC21 论文在 Selene 上的实测，不是 PyTorch 或 Megatron-Core 的基准。每个节点 8 张 80GB A100，节点内 NVLink/NVSwitch，节点间 8 条 200Gbps HDR InfiniBand，三层的 leaf/spine/core fat-tree；单卡 16 位峰值 312 TFLOP/s（§5）。Table 1 是弱扩展：模型从 1.7B 长到 1008B，GPU 数从 32 长到 3072。

Table 1（SC21 v5）· Selene 上的弱扩展端到端吞吐；DP 列为页面派生 = GPU/(TP×PP)

Table 1 全表转录；DP 列是页面派生，不是原表列。弱扩展下每一行的模型和 GPU 数都不同，不能把相邻行当成强扩展曲线。

先看配置怎么变：TP 在所有行都不超过 8；PP 从 1 长到 64；DP 从 32 缩到 6。论文的 Takeaway #1 解释了这个选择：TP 用到单机 GPU 数 g，再用 PP 跨机扩展（§3.2）。1008.0B 行（论文称 1T）的 TP=8、PP=64、DP=6 正好落在 3072 张 A100 上，每 GPU 163 TFLOP/s，达到峰值 52%，聚合 502 PFLOP/s（Table 1）。

论文把 1T 行更高的利用率归因于更大的矩阵乘：模型越大，GEMM 越大，通信时间相对计算时间没有明显上升（§5.1）。这不是说 mesh 本身提高了效率；mesh 只是把 3072 个 rank 分成了正确的通信组。

Table 2（SC21 v5）· PTD 与 ZeRO-3（不带模型并行）的配对比较；按模型、batch、GPU 数配对读

* 530B 的 ZeRO-3 在 560 GPU、microbatch 4 下放不下，论文改用 640 GPU、batch 2560 给出估算（Table 2 注）。

配对看最清楚：174.6B、1536 GPU、batch 1536、microbatch 1 时，PTD 141 vs ZeRO-3 44 TFLOP/s/GPU，300B token 的训练时间 23 vs 74 天；529.6B、2240 GPU、batch 2240、microbatch 1 时，159 vs 48，42 vs 140 天（Table 2）。这些是端到端数字，包含通信、数据加载和优化器；不是单纯的 GEMM 吞吐。

论文还报告：GPU 翻倍、batch 不变时，PTD 相对 ZeRO-3 的吞吐提升 70%，175B 与 530B 两个模型都是这个幅度（§5.2）。70% 是相对提升，条件是 batch 不变、只加倍 GPU 数；不是绝对百分点。

形状、rank 顺序、group 创建、异构链路和框架衔接，五处会限制扩展效率。

1. 形状：乘积必须对得上 world_size

init_device_mesh 用 mesh_shape 的乘积生成长度为该乘积的 rank map（torch/distributed/device_mesh.py:1720-1725）；代码里只检查 rank map 覆盖 layout 的大小（同文件 305-309），没有显式的 world_size 相等检查。页面把 world_size = DP×TP×PP×CP 当作使用前提。Megatron 则显式检查 world_size % model_size == 0，其中 model_size = TP×PP×CP×gtp_remat（默认 gtp_remat=1），DP 由除法得到（megatron/core/parallel_state.py:883-895）。

DeviceMesh 还要求 layout 正交、不重叠（torch/distributed/device_mesh.py:297-300），维名不重复且长度与 shape 一致（同文件 1691-1702）。这些约束不是形式主义：mesh 是矩形数组，每个 rank 必须恰好落在一格；如果 shape 或维名对不上，初始化会报错；如果 rank 顺序对不上，组可能建在错误的卡上。

2. 集合约束：group 创建不是本地操作

new_group 要求主组所有进程进入，并按相同顺序建组（torch/distributed/distributed_c10d.py:7282-7285）。DeviceMesh 对每一维遍历所有子组，每个 rank 都为所有子组调用 new_group（torch/distributed/device_mesh.py:639-649, 665-694）。因此 group 创建是集合操作：rank 顺序不一致会 hang；组数随 world_size 线性增长。

按当前 Megatron-Core 的四个主循环推算 1T 配置（TP=8、PP=64、DP=6、CP=1）：DP 512 组、TP 384 组、PP 48 组、CP 3072 个 singleton 组，四个循环合计 4016 个 NCCL group；默认 create_gloo_process_groups=True 时，DP 循环还会为每个 DP 组再建一个 Gloo group（megatron/core/parallel_state.py:648, 1217-1220），create_group 调用约 4528 次（megatron/core/parallel_state.py:1210-1408；不含 GTP weight-shard、embedding、model-parallel、DP-with-CP 等额外循环）。这是推算，不是 SC21 的实测；SC21 论文没有 CP 轴。它说明一件事：把 CP 设为 1 仍保留一维，就会为每个 rank 建一个 singleton group；不需要这一维时，最好别把它放进 shape。DeviceMesh 同样会对 size-1 维循环生成 singleton group（665-694）。

3. 拓扑：逻辑相邻不等于物理相邻

DeviceMesh 的 rank map 是逻辑数组；默认 preserve_rank_order=False，子组 rank 会按升序排序（torch/distributed/device_mesh.py:213-220），new_group 的 sort_ranks 默认也是 True（torch/distributed/distributed_c10d.py:7332-7336）。如果 rank 编号与物理拓扑不一致，排序后同组 rank 可能不再相邻；需要保序时要显式打开 preserve_rank_order。

Megatron 的 order 字符串就是干这个的：注释说最左边的 token 步长最小、rank 最相邻（megatron/core/parallel_state.py:609-617）；docstring 也提醒调用方让相邻 rank 落在同一台 DGX（817-820）。SC21 §5.9 在 1T/3072 GPU 上测得 PP 阶段间 p2p 的有效二分带宽 892 GB/s，而 DP 副本间 all-reduce 为 12.9 TB/s；同一张逻辑 mesh 上，不同维度的物理链路并不等价。SC21 §3.2 还说，TP 度数超过单机 GPU 数时，all-reduce 要走跨节点链路，开销可能不实用；Takeaway #1 建议 TP 用到单机 GPU 数 g，再用 PP 跨机。

4. 框架衔接：FSDP 要的是 DP 子网格

fully_shard 的 mesh 是 data parallel mesh；1D 时全分片，2D 时第 0 维复制、第 1 维分片，即 HSDP 的 (Replicate(), Shard(0))（torch/distributed/fsdp/_fully_shard/_fully_shard.py:184-190）。_validate_mesh 默认只接受 1D/2D mesh，2D 还要求有 mesh_dim_names（torch/distributed/fsdp/_fully_shard/_fsdp_init.py:50-76）；要传更大的 SPMD mesh，得用 dp_mesh_dims 指定 shard/replicate 维（torch/distributed/fsdp/_fully_shard/_fully_shard.py:245-249）。

因此，不带 dp_mesh_dims 时，4D mesh 不能直接塞给 FSDP；要么先切出 mesh["dp"] 子网格，要么用 dp_mesh_dims 指定 DP 维（它内部会切子网格，torch/distributed/fsdp/_fully_shard/_fsdp_init.py:108-140）。本页讨论的 Megatron dense 建组路径不走 FSDP：它用 RankGenerator + create_group 建自己的 DP/TP/PP/CP 组（megatron/core/parallel_state.py:494-577, 1210-1464），DDP/optimizer 直接消费这些组；Megatron-Core 另有一套 FSDP 适配层（megatron/core/distributed/fsdp/），会在 DP 组之上构造 DeviceMesh。两套实现不能混着读：DeviceMesh 是 PyTorch 的组合式并行接口，RankGenerator 是 Megatron 的建组器。

5. 什么时候别用 4D mesh

如果模型能放进单机、batch 足够大，DP 或 DP+TP 往往已经够用；再加 PP/CP 只会多一套 group 和一类通信，除非它们确实解决了显存或序列长度约束。如果物理拓扑不规则，先把 rank 顺序和子网格映射对清楚，再谈 4D；mesh 只表达逻辑分组，不表达链路带宽。

一句话边界：mesh 能把「谁和谁通信」表达清楚，但不能把慢链路变快，也不能消除 all-reduce、p2p、KV 交换的字节数。扩展效率最终由通信量、链路带宽、pipeline 气泡、CP 的块大小和 group 启动成本共同决定；维数不是越多越好。