核心概念：rank / parallel state / micro-batch

阅读 Megatron 源码或调试它的训练任务，必须先把下面这一组"词汇表"内化掉。这一章不教你"怎么做"，而是给你一份字典，后续 5–10 章会反复用到它们。

4.1 rank 一族：global / dp / tp / pp / ep / cp

PyTorch 原生只有一个 rank（全局编号），Megatron 把它拆成了一组：

为什么要拆这么细？ 因为一张物理 GPU 同时属于多个 group：它既是某个 DP group 的一员（要做梯度 all-reduce），又是某个 TP group 的一员（要做参数 all-gather），还是某个 PP group 的一员（要把激活值传给下一个 stage）。每种 group 都有自己独立的通信拓扑，所以 rank 必须分别记。

举个具体例子。8 卡机，配 $\mathrm{TP}=2,\;\mathrm{PP}=2,\;\mathrm{DP}=2$：

• 卡 R0：global rank=0, tp_rank=0, pp_rank=0, dp_rank=0
• 卡 R3：global rank=3, tp_rank=1, pp_rank=0, dp_rank=1
• 卡 R6：global rank=6, tp_rank=0, pp_rank=1, dp_rank=1

掌握"一张卡同时持有多个身份"是后面读所有并行代码的前提。

4.2 parallel state：那个全局对象

megatron/core/parallel_state.py 维护一组 module-level 全局变量，记录上面这些 group。初始化它的入口是：

parallel_state.initialize_model_parallel(
    tensor_model_parallel_size=2,
    pipeline_model_parallel_size=2,
    # virtual_pipeline_model_parallel_size=None,
    # context_parallel_size=1,
    # expert_model_parallel_size=1,
)

这一个调用之后，整个进程后续所有 get_xxx_rank()、get_xxx_world_size()、get_xxx_group() 才有意义。在你的最小例子 examples/run_simple_mcore_train_loop.py 里，initialize_distributed() 就在干这件事。

4.3 micro-batch vs global-batch

Megatron 里"batch"有两层含义，记混就理解不了 PP：

三者关系（设 micro batch 为 $b$，micro batch 数为 $m$，数据并行度为 $d$）：

$$B_\text{global} = b \times m \times d$$

关键变量 $m$（num_micro_batches）：一次 optimizer step 内塞进 pipeline 的小批次数量。它对 pipeline bubble 影响巨大——越大 bubble 占比越小，但单步显存越紧。

$$\mathrm{Bubble\ Ratio} \approx \frac{\mathrm{PP}-1}{m + \mathrm{PP}-1}$$

典型例子（175B 配置）：--global-batch-size 1536，--micro-batch-size 1，$\mathrm{DP}=12$ →

$$m = \frac{1536}{1 \times 12} = 128$$

记忆口诀：micro 是显存维度，global 是收敛维度，二者通过 pipeline depth × DP 联系起来。

4.4 sequence_length / context_parallel / virtual pipeline

这些是稍后才会用到、但概念要早知道的：

第 8、9 章会详细讲，这里只要知道这些词在说什么维度即可。

4.5 model providers 和 forward step：你需要传给框架的两个函数

Megatron 的训练入口本质上是一个泛型循环，它要求你提供两件东西：

1. model_provider(pre_process, post_process)：告诉框架"给定这个 stage 上应不应该有 embedding 和 output layer，请构造对应 module"。pretrain_gpt.py 里这个函数在 gpt_builders.gpt_builder() 里实现。
2. forward_step(data_iterator, model) -> (output, loss_func)：定义"拿一份 micro batch 跑一次 forward 应该做什么、loss 怎么算"。examples/run_simple_mcore_train_loop.py 第 ~140 行附近就是教科书式范本。

框架自身（pretrain_gpt.py 的 main 函数）只负责调度、optimizer step、checkpoint。算法可变的部分（模型 / 数据 / loss）通过这两个 callback 注入。这种"框架 + 用户回调"的关系一旦内化，后面读源码就知道每个文件的角色。

4.6 dataset config：BlendedMegatronDatasetBuilder 与 mock

预训练的 dataset 抽象比微调复杂：

from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
from megatron.core.datasets.gpt_dataset import GPTDatasetConfig, MockGPTDataset

config = GPTDatasetConfig(
    random_seed=0,
    sequence_length=2048,
    reset_position_ids=False,
    reset_attention_mask=False,
    eod_mask_loss=False,
    tokenizer=...,
    ...
)
train_ds, val_ds, test_ds = BlendedMegatronDatasetBuilder(
    GPTDataset,                         # 或 MockGPTDataset 走假数据
    [train_samples, val_samples, test_samples],
    is_built_on_this_rank=lambda: True,
    config=config,
).build()

第 6 章讲得更细，这里先认识它们的名字。

4.7 一张图把这章串起来

后面所有章节都在拆这个流程里的某一格。

4.8 TransformerConfig 字段分类速览

第 5 章会讲 CLI；这里看一下底层 dataclass。megatron/core/transformer/transformer_config.py:52 起的 TransformerConfig 有 208 个字段，分八类：

查字段最快：grep -n "^ [a-z_]*:" megatron/core/transformer/transformer_config.py。CLI flag 对应字段名通常 dash → underscore（--num-layers ↔ num_layers）。

4.9 ParallelState 全表：12 种并行维度

megatron/core/parallel_state.py（1790 行）提供 53 个 get_*_world_size / get_*_rank 函数，覆盖如下并行维度：

初始化入口：parallel_state.initialize_model_parallel(tp, pp, vp, cp, ep, ...)，按笛卡尔积生成所有 group 并存到 module-global 变量里。所有并行的"分卡逻辑"都从这一个函数出发。