数据截至 2026-08-27(上游 commit e79cb4c1bae1)

01 · 五种并行维度与进程组网格

这一章讲什么： Megatron 的全局地图。五种并行各自「切的是哪样东西」、initialize_model_parallel 怎么把 world_size 张成进程组网格、以及训练入口 pretrain() 长什么样。读完你应该能对着一组 --tensor-model-parallel-size 之类的参数，说出每个 rank 在网格里的坐标。

1. 它要解决的小问题

一个 175B 参数的模型，bf16 光参数就 350GB，一张 80GB 的 H100 连参数都装不下，遑论梯度和优化器状态。

于是问题变成：沿哪个维度把训练切开，摊到几千张卡上？ 而且切完之后，跨卡的部分还得拼出和没切时一模一样的数学结果。

Megatron 的历史贡献，是把「切法」系统化成了几个正交维度——每把刀切的方向不同、产生的通信模式不同，可以自由组合。这一章先把五把刀认全。

2. 五种并行，各切什么

一张表里认全五把刀。注意「切什么」这一列——它们切的是不同的东西，这正是「正交」的含义：

维度	全称	切什么	通信原语	代价特征
TP	Tensor Parallelism 张量并行	单层的权重矩阵（按列/按行）	all-reduce（每层两次）	通信量大，只敢开在节点内 NVLink 域（一般 ≤8）
PP	Pipeline Parallelism 流水并行	层（模型的不同段放不同卡）	P2P send/recv（只发生在段边界）	通信量小，但有「气泡」（卡等人）
DP	Data Parallelism 数据并行	batch（每卡一份完整模型、不同数据）	梯度同步（ZeRO 形态：reduce-scatter + all-gather）	通信量与模型大小成正比，但能和计算重叠
CP	Context Parallelism 上下文并行	序列（长序列切成块，attention 跨块交换 KV）	all-gather / P2P 环（attention 内部）	专为长序列；切激活不切权重
EP	Expert Parallelism 专家并行	MoE 的专家（不同卡放不同专家）	all-to-all（按路由结果搬运 token）	只在 MoE 模型里存在

还有一个常被一起提的 SP（Sequence Parallelism 序列并行）：它不是独立维度，而是 TP 的补丁——把 LayerNorm/Dropout 处没被 TP 切开的激活再按序列维切一次，省显存。它复用 TP 的进程组，见 02 章。

一句话直觉

五种并行 = 五种切蛋糕的刀法。 蛋糕有三层（参数、激活、优化器状态），每把刀下刀的方向不同：

TP 横着切参数矩阵；
PP 竖着按层分段；
DP 不切实物，复制整份蛋糕、各喂不同数据；
CP 沿序列维切激活；
EP 把 MoE 专家这盘菜分到不同桌上。

刀法可以叠加：world_size = TP × PP × DP × CP（EP 从 DP 里划出，见 §5）。这就是「5D 并行」说法的出处。

3. 图示：world_size 怎么变成网格

initialize_model_parallel（megatron/core/parallel_state.py:600）做的事，一句话：把 0..world_size-1 这串 rank 编号，按默认 order "tp-cp-ep-dp-pp" 解释成一个多维坐标系，再沿每个轴切出进程组。

以 world_size=16、TP=2、CP=2、DP=2、PP=2 为例（从左到右读：order 越靠左，rank 编号挨得越近）：

rank 编号:  0  1 | 2  3 | 4  5 | 6  7 || 8  9 | ...
            └──┘
            TP 组内（相邻 2 个 rank 同一 TP 组）
            └────────┘
            CP 维度再跨一步
            └────────────────┘
            DP 维度再跨
            └────────────────────────────────┘
            PP 在最外层，跨得最远

切出的组（每个 rank 同时属于 5 个组）:
  TP 组: [0,1], [2,3], [4,5], [6,7], ...
  CP 组: [0,2], [1,3], [4,6], ...
  DP 组: [0,4], [1,5], ...
  PP 组: [0,8], [1,9], ...

怎么读这张图： order 字符串决定「哪个维度变化最快」。tp-cp-ep-dp-pp 意味着相邻编号的 rank 先填 TP 维、再 CP、再 DP、最后 PP。这个排布不是随意的——通信最频繁的维度（TP）被放在 rank 最相邻的位置，而 rank 相邻通常意味着物理上同机、同 NVLink 域。PP 通信最少，被扔到跨节点的远端。

DP 有个细节：它不出现在参数里，而是除出来的——data_parallel_size = world_size // model_size（megatron/core/parallel_state.py:824），其中 model_size = tp * pp * cp * gtp_remat（:814）。剩下的都归 DP。

4. 原理演示：笛卡尔积生成器

Megatron 没有为五种并行各写一套建组逻辑，而是写了一个通用的「按 mask 取正交组」生成器。下面用简化代码演示它的核心数学：

# 示意，非源码
def build_groups(world_size, sizes, mask):
    # sizes  = 各维度的 size，按 order 排列，如 [tp=2, cp=2, dp=2, pp=2]
    # mask   = 要哪几维，如 [True, False, False, False] 表示取 TP 组
    groups = {}
    for rank in range(world_size):
        # 把 rank 按 sizes 分解为多维坐标（像把秒数拆成时分秒）
        coords = decompose(rank, sizes)
        # 组号 = 固定住 mask 维、把其余维的坐标拼成 key
        key = tuple(c for c, m in zip(coords, mask) if not m)
        # 组内位置 = mask 维的坐标
        groups.setdefault(key, []).append(rank)
    return list(groups.values())

重点看：「建组」本质是把 rank 编号当多维数组下标做 reshape/transpose。同一段代码，mask 给 [True,0,0,0] 就出 TP 组，给 [0,0,True,0] 就出 DP 组。

真实实现 generate_masked_orthogonal_rank_groups（megatron/core/parallel_state.py:269）做的就是这件事，docstring 里把数学推导完整写了出来：global_rank = tp_rank + dp_rank * tp_size + pp_rank * tp_size * dp_size，并给了 24 卡 [2,3,4] 的完整算例。包它的 RankGenerator（parallel_state.py:465）则把 order 字符串解析成 ordered_size 列表，get_ranks("tp") 内部就是「按 token 生成 mask → 调生成器」两步（get_mask/get_ranks，parallel_state.py:517、:532）。

5. 真实实现：initialize_model_parallel

函数在 megatron/core/parallel_state.py:600，签名里二十来个参数，前六个就是五把刀的大小（外加虚拟流水）：

def initialize_model_parallel(
    tensor_model_parallel_size: int = 1,
    pipeline_model_parallel_size: int = 1,
    virtual_pipeline_model_parallel_size: Optional[int] = None,
    ...
    context_parallel_size: int = 1,
    ...
    expert_model_parallel_size: int = 1,
    ...
    order: str = "tp-cp-ep-dp-pp",

（parallel_state.py:600-627，:616 是 order 默认值）

函数体分三段：

① 推导 DP。 先算 model_size（TP×PP×CP×GTP_remat），要求 world_size 能被整除（:814-822），然后 data_parallel_size = world_size // model_size（:824）。DP 是唯一不直接给参数的维度——它是「扣掉模型并行后剩下的份数」。

② 造两个 RankGenerator。 注意这里有个容易漏的设计：稠密部分（attention/MLP）和 MoE 专家部分各有一个生成器（decoder_rank_generator 在 :857，expert 的在 :885 起）。原因写在 RankGenerator 的断言里：「EP 和 CP 不能同时 >1 出现在一个生成器里」（parallel_state.py:479）——CP 进默认生成器，EP 进专家生成器。两个生成器还要保证 PP 组一致（:903 的断言）。

③ 逐维度建组。 之后就是一长串 for ranks in generator.get_ranks("xx"): create_group(...)，例如 TP 组在 :1163 附近、PP 组在 :1239。每建一个组，顺手把「我属于哪个组」写进模块级全局变量（_TENSOR_MODEL_PARALLEL_GROUP 等），之后用 get_tensor_model_parallel_group()（parallel_state.py:1703）这类 getter 取。

建完组之后，模型层的构造函数会去查这些全局组来决定自己怎么切——比如 ColumnParallelLinear 默认拿 get_tensor_model_parallel_group() 当自己的通信组。这就是为什么库里建模型时看不到任何「把进程组传给层」的代码：网格是全局单例，层自己认领。最小示例见 examples/run_simple_mcore_train_loop.py:52-55：先 initialize_model_parallel(tp, pp)，然后 model_provider() 里正常 GPTModel(...) 即可。

6. 训练入口：pretrain() 的骨架

库的另一面是参考训练脚本。入口是 pretrain()（megatron/training/training.py:1500），docstring 自己把流程写成了四步：

initialize Megatron（含上面的进程组初始化）；
用 model_provider 建模型、优化器、lr schedule；
用 train_valid_test_dataset_provider 拿数据集；
用 forward_step_func 训练。

调用侧的最小样子是 pretrain_gpt.py：它提供 model_provider（怎么建 GPT）和 forward_step（一个 microbatch 怎么算 loss），其余全交给 pretrain。一个 step 内部是 train_step（training.py:3015）：清梯度 buffer → 调 forward_backward_func（PP 调度器，见 03 章）→ optimizer step（见 04 章）。整条主线的端到端走读在 06 章。

7. 关键细节与坑

order 字符串是排布契约，不是摆设。 它决定 rank→物理卡的映射下哪个维度走 NVLink、哪个跨节点。改了 order 而不改集群拓扑认知，性能可能天差地别。新维度（如 GTP remat）靠 _inject_gtp_remat_axis(order, after="cp")（parallel_state.py:580、:856）插进 order，不用改生成器。
虚拟流水要求 PP>1。 virtual_pipeline_model_parallel_size 非 None 而 PP=1 会直接 RuntimeError（parallel_state.py:826-830）——interleaved 调度没有单段意义。
EP 的 size 从 DP 里扣。 expert 生成器的 dp 是 expert_data_parallel_size，其乘积关系让 EP×expert_DP = DP（inferred：由两个生成器的构造参数关系得出，:885 附近）。所以开 EP 会摊薄每张专家卡上的数据并行份数。
attention 和 MoE 的 PP 组必须一致。 :903 的断言直接比较两个生成器的 get_ranks("pp")——混专家模型不能把层切出两种不同的流水布局。
建组顺序有讲究（SHARP）。 注释写明：用 SHARP 时 dp-cp 组必须最先建，因为硬件只允许最初创建的通信组用 COLLNET 特性（parallel_state.py:958-963 的注释）。这是「不踩过坑写不出来」的那类代码。
全局单例意味着不可重入。 所有 _XXX_GROUP is None 断言（如 :1160）在防二次初始化；测试里要换网格得先 destroy_model_parallel()（示例文件 run_simple_mcore_train_loop.py:41 就这么干）。

8. 代码地图

主题	文件路径	符号名
进程组初始化总入口	`megatron/core/parallel_state.py`	`initialize_model_parallel`
笛卡尔积建组生成器	`megatron/core/parallel_state.py`	`RankGenerator`、`get_mask`、`get_ranks`
正交组数学（docstring 含推导）	`megatron/core/parallel_state.py`	`generate_masked_orthogonal_rank_groups`
order 注入新维度	`megatron/core/parallel_state.py`	`_inject_gtp_remat_axis`
各维度 group getter	`megatron/core/parallel_state.py`	`get_tensor_model_parallel_group`、`get_data_parallel_group`、`get_pipeline_model_parallel_group`、`get_context_parallel_group`、`get_expert_model_parallel_group`
销毁网格（测试用）	`megatron/core/parallel_state.py`	`destroy_model_parallel`
训练主入口	`megatron/training/training.py`	`pretrain`、`train_step`
GPT 预训练脚本	`pretrain_gpt.py`	`model_provider`、`forward_step`
库用法最小示例	`examples/run_simple_mcore_train_loop.py`	`initialize_distributed`、`model_provider`

下一章：02 · 张量并行——五把刀里通信量最大、设计最精巧的那把。

1. 它要解决的小问题​

2. 五种并行，各切什么​

一句话直觉​

3. 图示：world_size 怎么变成网格​

4. 原理演示：笛卡尔积生成器​

5. 真实实现：initialize_model_parallel​

6. 训练入口：pretrain() 的骨架​

7. 关键细节与坑​

8. 代码地图​