Chapter 13

源码导读：从 pretrain_gpt.py 顺藤摸瓜

📌 所有路径基于 NVIDIA/Megatron-LM · commit 432d76b2b8bf （v0.15.0）。目标：读完之后你给 Megatron 提 PR 时知道改哪个文件。

13.1 仓库地图

Megatron-LM/
├── pretrain_gpt.py             ★ GPT 训练主入口（用户写法）
├── pretrain_mamba.py           Mamba 主入口
├── pretrain_hybrid.py          Hybrid 主入口
├── pretrain_vlm.py             视觉语言模型主入口
├── train_rl.py                  RL 训练主入口
├── gpt_builders.py             ★ GPT model_provider
├── mamba_builders.py           Mamba builder
├── hybrid_builders.py          Hybrid builder
├── model_provider.py           通用 model provider 入口
├── megatron/
│   ├── training/                ★ 训练循环 + arguments
│   │   ├── arguments.py         ★ CLI 入口（第 5 章）
│   │   ├── training.py          ★ 主循环 setup_model_and_optimizer / pretrain
│   │   ├── checkpointing.py     ★ checkpoint save/load
│   │   ├── tokenizer/           tokenizer 抽象
│   │   ├── initialize.py        分布式初始化
│   │   └── utils.py
│   ├── core/                    ★ Megatron-Core 库（独立 pip 包）
│   │   ├── tensor_parallel/     第 7 章
│   │   ├── pipeline_parallel/   第 8 章
│   │   ├── transformer/         核心 transformer 实现
│   │   ├── models/              模型架构（GPT / Mamba / T5 / Bert）
│   │   ├── datasets/            第 6 章
│   │   ├── distributed/         distributed optimizer 等
│   │   ├── fusions/             fused kernel
│   │   ├── extensions/          TransformerEngine 集成
│   │   ├── dist_checkpointing/  第 12 章
│   │   ├── inference/           推理引擎
│   │   ├── export/              导出工具
│   │   ├── moe/                  MoE 实现
│   │   └── ...
│   ├── post_training/           RL / SFT 后训练
│   ├── elastification/          弹性扩缩容
│   └── inference/                推理服务
├── tools/                       命令行工具
│   ├── preprocess_data.py       第 6 章
│   ├── checkpoint/              第 12 章
│   ├── report_theoretical_memory.py  第 11 章
│   └── ...
├── examples/                    各种训练脚本
├── tests/                       单元测试
├── docker/                      Dockerfile
├── docs/                        文档源
└── scripts/                     辅助脚本

13.2 主调用栈：pretrain_gpt.py 启动后发生了什么

flowchart TB CLI["torchrun ... pretrain_gpt.py <一长串 args>"] CLI --> Main["pretrain_gpt.py · main()"] Main --> Init["megatron/training/initialize.py
initialize_megatron()"] Init --> A["torch.distributed.init_process_group"] Init --> B["parallel_state.initialize_model_parallel(TP, PP, CP, EP)"] Init --> C["set seeds / RNG tracker"] Init --> D["build tokenizer"] Main --> Pretrain["megatron/training/training.py
pretrain(...)"] Pretrain --> MP["setup_model_and_optimizer()"] MP --> ModelProv["gpt_builders.py · gpt_builder()
构 GPTModel"] MP --> Optim["build optimizer (Adam/Lamb/Muon)"] MP --> Sched["build LR scheduler"] Pretrain --> DL["build_train_valid_test_data_iterators()"] Pretrain --> Loop[train(...) 主循环] Loop --> Step["train_step(...)"] Step --> FB["pipeline_parallel/schedules.py
get_forward_backward_func()()"] FB --> Fwd[forward_step_func] FB --> Bwd[backward via autograd] Step --> Opt[optimizer.step] Step --> Save["checkpointing.py · save_checkpoint() 每 N 步"]

13.3 关键文件逐个看

pretrain_gpt.py（用户脚本侧）

从 ch04 学到的"用户要提供两个 callback"在这里具体化：

# pretrain_gpt.py（简化）
from megatron.training import pretrain
from gpt_builders import gpt_builder

def forward_step_func(data_iterator, model):
    batch = next(data_iterator)
    output = model(batch["tokens"], position_ids=..., attention_mask=...)
    return output, partial(loss_func, batch["loss_mask"])

if __name__ == "__main__":
    pretrain(
        train_valid_test_dataset_provider=...,
        model_provider=gpt_builder,
        model_type=ModelType.encoder_or_decoder,
        forward_step_func=forward_step_func,
        ...
    )

megatron/training/training.py

这是 Megatron 的"框架本体"，约 2000 行。核心函数：

函数	做什么
`pretrain(...)`	顶层入口；初始化 + 调 `train`
`setup_model_and_optimizer(...)`	用 model_provider 造模型，挂上 distributed wrapper
`train(...)`	主循环：iter → train_step → log → save
`train_step(...)`	单步：调 forward_backward_func + optimizer.step
`evaluate_and_print_results(...)`	验证

megatron/core/transformer/

核心模型实现，分层组织：

文件	说明
`transformer_config.py`	`TransformerConfig` dataclass，所有结构超参
`transformer_block.py`	多层 transformer block 堆叠
`transformer_layer.py`	单层 transformer（attention + MLP + norm）
`attention.py`	self-attention 实现
`dot_product_attention.py`	不同 backend 的 attention
`mlp.py`	MLP / SwiGLU
`moe/`	MoE 实现（router、experts、token dispatch）
`spec_utils.py`	ModuleSpec 抽象（让 attention/MLP 可替换）

13.4 ModuleSpec 抽象：Megatron 的"积木"

从 Megatron-Core 起，attention / MLP / norm 等都通过 ModuleSpec 注入：

from megatron.core.transformer.spec_utils import ModuleSpec
from megatron.core.tensor_parallel import ColumnParallelLinear

mlp_spec = ModuleSpec(
    module=MLP,
    submodules=MLPSubmodules(
        linear_fc1=ModuleSpec(module=ColumnParallelLinear),
        linear_fc2=ModuleSpec(module=RowParallelLinear),
    ),
)

好处：

切 backend（TE 替换 PyTorch）只要换 ModuleSpec；
SwiGLU / GeLU / Mamba 块都能注入；
不用为每种模型重写整套 transformer 代码。

看 megatron/core/models/gpt/gpt_layer_specs.py 是理解 ModuleSpec 怎么"组装模型"的最好入口。

13.5 阅读路线总结

"读完 Megatron"是个工程量惊人的任务（~50 万行 Python + C++）。建议这样推进：

跑通 ch03 最小例子：建立感觉；
读 pretrain_gpt.py + gpt_builders.py：理解用户侧 callback；
读 megatron/training/training.py · pretrain → train → train_step：理解主循环；
读 megatron/training/initialize.py：理解分布式初始化；
读 megatron/core/tensor_parallel/mappings.py + layers.py（ch07）：理解 TP；
读 megatron/core/pipeline_parallel/schedules.py（ch08）：理解 1F1B；
读 megatron/core/transformer/transformer_block.py：理解 model forward；
读 megatron/core/datasets/blended_megatron_dataset_builder.py（ch06）：理解数据；
挑一个 PR 提交（哪怕修 typo）：跑通 contribution 流程。

13.6 想提 PR？从这里入手

改动类型	对应文件
新模型架构	`megatron/core/models/<new>/` + 一个 layer_specs.py
新 attention backend	`megatron/core/transformer/dot_product_attention.py`
新 fused kernel	`megatron/core/fusions/<new>.py`
新优化器	`megatron/core/optimizer/<new>_optimizer.py`
新 schedule	`megatron/core/pipeline_parallel/schedules.py`
新 dataset 类型	`megatron/core/datasets/<new>_dataset.py`
新 tokenizer	`megatron/training/tokenizer/tokenizer.py`
新 checkpoint format	`megatron/core/dist_checkpointing/` + `tools/checkpoint/`
新 CLI flag	`megatron/training/arguments.py` 加 `add_argument`
新 MoE 路由策略	`megatron/core/transformer/moe/router.py`

提 PR 前看 CONTRIBUTING.md 和 tests/，每个新 feature 都要带单元测试 + 性能 regression test。NVIDIA 对合入要求严格但反馈很快。

13.7 与外部生态的对应关系

外部框架	与 Megatron 的依赖
Megatron-Core (pip)	= `megatron/core/` 部分，独立发版
NeMo	用 Megatron-Core 当底座
Megatron-DeepSpeed	fork 仓库，DS ZeRO + Megatron 并行
verl	训练后端之一是 Megatron-Core
ColossalAI	borrowing TP/PP 思想，独立实现
TorchTitan	PyTorch DTensor 路线，独立实现
Megatron-Bridge	NV 出，做 HF ↔ Megatron 互转

13.8 megatron/core/ 真实子目录树（v2026）

13.1 那份是简化版。下面是 ls megatron/core/ 真实输出加用途批注，配合源码读时必备：

megatron/core/
├── parallel_state.py            ★ 53 个 get_*_world_size/get_*_rank
├── transformer/                  Transformer 算子层
│   ├── transformer_config.py    ★ TransformerConfig (208 字段)
│   ├── transformer_layer.py:291 ★ TransformerLayer 主类
│   ├── multi_latent_attention.py:1437  MLA
│   ├── multi_token_prediction.py        MTP
│   ├── moe/                              MoE 全部组件
│   │   ├── router.py:138                 TopKRouter
│   │   ├── token_dispatcher.py:53/354/1395  AllGather/AllToAll/Flex 三 dispatcher
│   │   └── experts.py:1080               SequentialMLP
│   ├── attention.py / spec_utils.py / module.py
│   └── identity_op.py / mlp.py / dot_product_attention.py
│
├── tensor_parallel/                ★ TP 实现
│   ├── layers.py:770/1134/197    ColumnParallel / RowParallel / VocabParallel
│   ├── mappings.py                4 个核心原语
│   ├── random.py                  CUDA RNG tracker
│   └── cross_entropy.py           VocabParallel CE
│
├── pipeline_parallel/              ★ PP 调度
│   ├── schedules.py:48/592/896/2039  dispatch + 三 schedule
│   ├── hybrid_cp_schedule.py             PP+CP 联合
│   └── p2p_communication.py
│
├── distributed/                    ★ DDP / param overlap
│   ├── distributed_data_parallel.py:22   DDP wrapper
│   ├── ddp_config.py
│   ├── finalize_model_grads.py
│   ├── param_and_grad_buffer.py
│   └── fsdp/                              MCore 自家 FSDP
│
├── optimizer/                      ★ 12 个文件
│   ├── __init__.py:962             get_megatron_optimizer 工厂
│   ├── distrib_optimizer.py:102    DistributedOptimizer（ZeRO-1 风格）
│   ├── clip_grads.py / grad_scaler.py
│   └── optimizer_param_scheduler.py
│
├── dist_checkpointing/             ★ 分布式 ckpt
│   ├── core.py                     save() / load() 顶层 API
│   ├── serialization.py / mapping.py
│   └── strategies/
│       ├── async_utils.py          ★ async I/O
│       ├── filesystem_async.py     ★ 非阻塞磁盘写
│       ├── fully_parallel.py / torch.py / nvrx.py / two_stage.py
│
├── datasets/                       ★ 21 个文件
│   ├── blended_megatron_dataset_builder.py
│   ├── gpt_dataset.py
│   ├── indexed_dataset.py
│   ├── packed_seq_params.py        ★ 变长打包
│   └── megatron_tokenizer.py
│
├── fusions/                        ★ 14 个 fused 算子
│   ├── fused_mla_yarn_rope_apply.py  ★ MLA + YaRN RoPE
│   ├── fused_bias_geglu.py / fused_bias_swiglu.py
│   ├── fused_softmax.py / fused_layer_norm.py
│   └── fused_indices_converter.py
│
├── extensions/                     ★ TE / Triton wrapper
│   ├── transformer_engine.py        TE 集成入口
│   └── transformer_engine_spec_provider.py
│
├── models/                         ★ 顶层模型
│   ├── gpt/                          GPT / Llama / Mistral 系
│   ├── t5/                            T5 encoder-decoder
│   ├── mamba/                         Mamba SSM
│   ├── multimodal/                    VLM
│   ├── bert/
│   └── retro/                         RETRO 检索增强
│
├── inference/                      ★ 推理引擎
│   ├── engines/
│   ├── text_generation/
│   ├── sampling/
│   └── moe/                           MoE 推理 dispatch
│
└── ssm/                             ★ Mamba SSM 子模块
    ├── mamba.py
    └── mamba_context_parallel.py

13.9 megatron/ 顶层非 core 子目录

目录	用途
`megatron/training/`	训练 loop / launcher / arguments.py / 数据加载
`megatron/inference/`	推理时的工具函数（生成 / 解码）；底层调 `core/inference`
`megatron/legacy/`	★ 老 API（2024 前），新代码请用 `core/`
`megatron/elastification/`	★ 弹性训练（动态加减节点）
`megatron/post_training/`	★ SFT / DPO 后训练
`megatron/rl/`	★ 强化学习（reward 模型 / PPO）

13.10 examples/ 模型清单

NV 官方 recipe 都在这里，从这里复制比读 docs 快：

目录	模型
`examples/gpt3/`	GPT-3 175B baseline
`examples/llama/`	Llama 2 / 3（RoPE + GQA）
`examples/mixtral/`	Mixtral 8x7B（稀疏 MoE）
`examples/mamba/`	Mamba SSM（S6 层）
`examples/gptoss/`	GPT-OSS（YaRN 长上下文 + attention sink）
`examples/t5/`	T5 encoder-decoder
`examples/bert/`	BERT MLM
`examples/multimodal/`	VLM（CLIP / LLaVA 风格）
`examples/megatron_fsdp/`	MCore 自家 FSDP recipe

13.11 三条推荐阅读路径

路线 A · 工程师"我要跑大模型"：pretrain_gpt.py → arguments.py（5 章已列）→ training.py:train_step → schedules.py:48；
路线 B · 研究员"我要改 attention"：transformer_layer.py:291 → 子模块 spec 系统 → multi_latent_attention.py:1437 抄一份；测试在 tests/unit_tests/transformer/；
路线 C · 性能工程师"我要榨吞吐"：fusions/ 全过一遍 → extensions/transformer_engine.py → distrib_optimizer.py:102 → schedules.py 的 overlap 逻辑。

13.12 这章你需要带走什么

知道 pretrain_gpt.py → training.py → train_step → forward_backward_func 这条主调用栈；
知道 megatron/core/ 是可独立 pip 装的库；
知道 ModuleSpec 是 Megatron-Core 的"可替换积木"机制；
知道想动哪类功能就改哪个目录；
提 PR 必须配测试，NVIDIA 风格严谨。

← 上一章

Checkpoint：Megatron ↔ HuggingFace 互转

论文与延伸阅读