Megatron-LM 与 PAI-Megatron：一次训练是怎么装配起来的

读完这节你应该知道：为什么“有 TP/PP 算法”不等于“有可复用的训练系统”。

SC21 论文开篇把约束说得很直接：模型参数已经放不进单卡显存；只靠数据并行又会遇到两个上限，一是每卡 batch 太小、利用率下降，二是可用设备数被 batch size 限制（§1）。所以训练系统必须把 tensor、pipeline、data 三种并行组合起来（论文称 PTD-P，§2）。

但“组合”不是调用一个函数。一次训练至少要回答五个装配问题：模型由哪些层 spec 组成；TP/PP/CP 的进程组怎么建；每个 rank 读哪一段数据；一次 optimizer step 跨多少 micro-batch；checkpoint 按什么 layout 分片、恢复时怎么映射回当前并行组。算法可以复用，这五件事的胶水仍然按模型、按集群、按版本变化。

Megatron-LM 仓库里能直接看到这种“入口按模型分家”：pretrain_gpt.py、pretrain_hybrid.py、pretrain_vlm.py 各有一个入口；examples/t5/pretrain_t5.py 和 examples/bert/pretrain_bert.py 也各自定义 get_batch、forward_step 和 dataset provider。例如 GPT 的 get_batch 在 pretrain_gpt.py:112，forward_step 在 :330，dataset provider 在 :452；T5 的同名函数在 examples/t5/pretrain_t5.py:145、:175、:202。核心组件是共享的，装配脚本不是。

这篇页面讲的是“框架怎么组织一次训练”，不重复 TP/PP 的算法推导。TP 的切分与通信账见 Tensor Parallel（Megatron），PP 的调度与 bubble 见 Pipeline Parallel。

读完这节你应该知道：已有方案各自把哪些问题留给了使用者。

Megatron 的回应是“共享循环 + 可配置模型装配”：训练循环只写一次，模型差异收敛到 layer spec、model provider 和 dataset provider；并行度由配置给出，进程组由 core 统一创建；checkpoint 由 core 的分片格式负责。SC21 把 TP、PP、DP 的组合称为 PTD-P，并进一步讨论 interleaved schedule 与 batch 配置（§2、§5.1）。

读完这节你应该知道：每个行为到底落在哪个文件、哪个函数，以及 PAI 在哪里插进来。

3.1 入口：先把 CLI 参数收成配置容器

pretrain_gpt.py 的 __main__ 只做一条装配链：解析参数 → 构造 GPTModelConfig → 构造 PretrainConfigContainer → 写入全局 run config → 调用 pretrain()（pretrain_gpt.py:496-519）。这一步没有建模型，也没有建进程组；它只是把“这次训练是什么”变成可传递的对象。

配置容器把一次训练拆成若干段：train、validation、model、optimizer、scheduler、data、ddp、dist、rng、logger、checkpoint、profiling、tokenizer、rl、rerun_state_machine、straggler（megatron/training/argument_utils.py:1030-1049）。以后新增一个训练开关，优先落进对应 section，而不是继续往全局 argparse 里塞。

模型结构落在 TransformerConfig：它继承 ModelParallelConfig（megatron/core/transformer/transformer_config.py:61），一边放 num_layers、hidden_size、num_attention_heads、ffn_hidden_size（:72、:170、:173、:211），一边继承 tensor_model_parallel_size、pipeline_model_parallel_size、sequence_parallel、context_parallel_size、expert_model_parallel_size（megatron/core/model_parallel_config.py:59、:112、:123、:129、:186）。core_transformer_config_from_args 按 dataclass 字段名逐项从 argparse 里取值（argument_utils.py:413-417），所以“CLI 参数名 = config 字段名”是这套配置体系能对齐的原因。

GPTModelConfig 再把模型专属参数和 builder 绑在一起：它内嵌 transformer，并用 builder 这个 ClassVar 指向 GPTModelBuilder（megatron/training/models/gpt.py:144-160）。ModelConfig.get_builder_cls() 动态 import builder（megatron/training/models/base.py:128-135），GPTModelBuilder.build_distributed_models() 负责建 stage、加 TP 属性、上 GPU、包 mix-precision 和 DDP（models/gpt.py:351-411、models/dist_utils.py:50-134）。旧的 model_provider / gpt_builder 路径仍然保留：model_provider.py:21-63 把 builder 转发给具体函数，gpt_builders.py:26-56 按 --spec、experimental attention、MoE、heterogeneous、TE/local 选择 layer spec，MTP block 再单独构造，然后构造 GPTModel（:92-110）。两条路径的取舍在 setup_model_and_optimizer 里：有 cfg.model 就走 builder，没有就退回 model_provider_func（training.py:2968-2999）。

3.2 初始化：先建 torch.distributed，再建并行组

pretrain() 的第一件大事是 initialize_megatron(...)（training.py:1853-1866）。它先校验 CUDA（initialize.py:87-89），初始化 runtime services，然后在 _initialize_distributed 里调用 torch.distributed.init_process_group（initialize.py:415-432），最后调用 mpu.initialize_model_parallel(...) 建所有并行组（initialize.py:457-481）。默认 rank 排列是 tp-cp-ep-dp-pp（initialize.py:475）。

源码：megatron/core/parallel_state.py:883-895；GTP 是 generalized tensor parallelism 的权重分片轴，普通训练取 1；以上是未启用 GTP sequence-sharding layout（gtp_remat_num_sequence_shards == 1）时的算法；启用该 layout（gtp_remat_num_sequence_shards != 1）时 model_size 会被 layout.minimum_world_size 覆盖（:889-890）。world_size 不能被 model_size 整除时直接抛 RuntimeError（:892-893）。

并行组不是“按需临时建”的：TP、PP、CP、EP、DP 以及它们的组合组在启动时一次建好，后面的层、DDP、optimizer、checkpoint 都从 ProcessGroupCollection 里取。这样训练步里不需要再判断“我现在在哪个组”，组关系已经编码在 rank layout 里。

3.3 模型装配：spec 决定一层长什么样，builder 决定 stage 怎么摊

layer spec 是“层的配方”。gpt_builder 先看 --spec 是否显式给了模块；否则按 transformer_impl、MoE、heterogeneous、experimental attention、MTP 等分支选择 spec（gpt_builders.py:33-70）。选完 spec 才构造 GPTModel，把 vocab_size、max_sequence_length、pre_process/post_process、并行组等传进去（:92-110）。

PP 和 VPP 的“stage 切片”不在这段代码里：build_distributed_models 会按 virtual_pipeline_model_parallel_size 建多个 stage（models/dist_utils.py:65-125），每个 stage 的 pre_process/post_process 由 VP/PP rank 推导（models/gpt.py:319-324）。这也是为什么同一份模型代码既能跑 PP=1，也能跑 PP=64：模型只描述“一层”，stage 由 builder 组织。

3.4 数据：dataset 只负责样本，入口负责按 TP/CP 切

数据侧先建 dataset，再包 dataloader，最后包 iterator。build_train_valid_test_data_iterators 接受 single、cyclic、external 三种 dataloader 类型（training.py:6309-6348）；GPT 的 provider 最终落到 build_gpt_datasets（megatron/training/datasets/providers.py:76-114），由 BlendedMegatronDatasetBuilder 建 GPTDataset。GPTDataset 包着 IndexedDataset，后者读 .bin/.idx 文件（megatron/core/datasets/gpt_dataset.py:178、:240-264），__getitem__ 返回 tokens、labels、loss_mask、position_ids 等（:302）。

“只让部分 rank 建 dataset”是显式规则：is_dataset_built_on_rank_explicit 要求 TP rank 0，并且只在 first/last PP stage（或 MTP rank）建；sequence packing 例外，会在每个 PP stage 建（providers.py:39-51）。读 batch 时，TP rank 0 先 next(data_iterator)，再广播到 TP 组（pretrain_gpt.py:166-191），最后按 CP 切序列（:210-215）。所以 dataset 不需要知道 TP/CP；并行切分留在入口的 get_batch 里。

3.5 训练步与 checkpoint：循环只认回调，不认具体模型

train_step 的顺序是固定的：清 grad buffer 和 optimizer grad（training.py:3473-3478）→ 调 forward_backward_func（:3573-3585）→ optimizer.step()（:3661-3662）。forward_backward_func 在 train() 入口按 PP/VP 选择一次（training.py:5071-5073）：PP=1 用 no-pipeline，PP>1 且 VP=None 用 1F1B，VP 非空用 interleaved（megatron/core/pipeline_parallel/schedules.py:157-167）。训练循环本身不依赖具体模型的 forward 实现，模型差异通过 forward_step 回调进来（GPT 的实现在 pretrain_gpt.py:330-449）。

checkpoint 是训练状态的一部分，不是训练结束后的导出。save_checkpoint 存模型、优化器、调度器、RNG、迭代数和已消费样本（checkpointing.py:640-880）；load_checkpoint 在模型和优化器建好之后调用（training.py:3155-3188），把分片状态映射回当前并行 layout。torch_dist 的 save/load 策略把 MCore 的 ShardedTensor 转成 PyTorch distributed checkpoint（dist_checkpointing/strategies/torch.py:684-690、:924-945）。一个容易忽略的限制：内置文本 dataloader 的批内位置不存，只有支持该接口的 Energon 多模态 dataloader 会存（checkpointing.py:677-678）；文本训练恢复靠已消费样本数和数据顺序。

读完这节你应该知道：给定并行度和 micro-batch，怎么算出一次 optimizer step 真正吃了多少样本、多少 token。

Megatron 的 batch 账本只有三个角色：micro_batch_size（每张卡一个 micro-batch 里的样本数）、num_microbatches（梯度累积多少个小批）、data_parallel_size（有多少路独立数据副本）。它们和 global batch 的关系是：

源码：megatron/core/num_microbatches_calculator.py:386-397；DP 取 args.data_parallel_size，标准无 GTP 时等于 world_size / (TP × PP × CP)（megatron/training/arguments.py:483-499）。

两个边界条件必须一起看。第一，GBS 必须能被 micro_batch_size × DP 整除，否则 core 直接断言失败（num_microbatches_calculator.py:381-386）；打开 --decrease-batch-size-if-needed 时，它会向下取整到可整除的 GBS（:368-379）。第二，如果没显式给 --global-batch-size，默认值是 micro_batch_size × batch_parallel_size（arguments.py:640-643），也就是只跑一个 micro-batch 的 GBS。

并行组映射：TP、PP、CP 先把 world 切掉，剩下才是 DP

第二个交互把配置项直接映到并行组。给定 world_size、TP、PP、CP，先算一个模型副本占多少 rank：model_size = TP × PP × CP；再用 DP = world_size / model_size 得到数据副本数。core 只接受整除的配置，否则 initialize_model_parallel 会报错（parallel_state.py:892-893）。

用 SC21 的 1T 行做一次交叉核对：Table 1 给出 TP=8、PP=64、3072 张 GPU；Table 1 没有列 CP，按 CP=1 代入（推断，依据是表中未出现 CP 列），model_size = 8 × 64 = 512，DP = 3072 / 512 = 6。也就是说这一行主要靠 PP 把模型摊到 64 个 stage，数据副本只有 6 路。这个结果与表中 batch=3072 并排看即可，不需要反推 micro-batch。

读完这节你应该知道：一次训练从 launcher 到第一个 optimizer step，中间到底按什么顺序经过哪些对象。

启动阶段的顺序是单向的：入口解析参数并构造配置容器 → initialize_megatron 建进程组与并行组 → setup_model_and_optimizer 建 stage、包 DDP、建 optimizer，并在函数末尾（若给了 --load）调用 load_checkpoint → 建 dataset/dataloader/iterator → 进入 train()（training.py:1853-1866、:2135-2147，内部 :3155-3188、:2160-2200、:2283-2301）。checkpoint 在数据 iterator 之前加载，是因为分片状态要按当前的 TP/PP/DP layout 映射，恢复的已消费样本数会决定 iterator 起点。

训练步是闭环：train_step 先清梯度，再让 forward_backward_func 处理 num_microbatches 个 micro-batch；每个 schedule 在末尾调用 finalize_model_grads 做 DP/CP 梯度归约与按 token 数缩放（training.py:5005-5006、schedules.py:857-865、finalize_model_grads.py:723-735），最后统一 optimizer.step()（含 grad clip 与参数更新，optimizer.py:580-591）。模型侧只提供 forward_step：它从 iterator 取一个 batch、调用 model(...)、返回输出和 loss 函数（pretrain_gpt.py:330-449）。训练循环不认识 GPT 的 loss，也不认识 Qwen 的 tokenizer；它只认这个回调契约。

读完这节你应该知道：这套装配在论文里跑出了什么数字，以及这些数字的条件是什么。

实验跑在 Selene 集群：每节点 8 张 80GB A100，每节点 8 个 200Gbps HDR InfiniBand HCA，单卡 16 位峰值 312 TFLOP/s（§5）。所有行使用 vocabulary 51200、序列长度 2048，并开启 interleaved pipeline schedule 与 scatter/gather 优化（§5.1）。

Table 1 · 弱扩展吞吐（转录自 SC21 Table 1）。每一行是不同规模的 GPT 模型，不是同一个模型换卡数；因此不能把相邻行读成加速比。

把表读成三件事。第一，论文把整张表概括为到 3072 张 A100 的 super-linear scaling：模型变大后矩阵乘更大，GPU 利用率提高，而通信时间相对计算没有明显增长（§5.1）。逐行看并非单调——18.4B 行是 135 TFLOP/s、43%，最大行是 163 TFLOP/s、52%；峰值占比的范围是 43% 到 52%。第二，1T 行的并行配置是 TP=8、PP=64、3072 张卡；按 core 的公式（CP 按 1 计，推断，因为表中没有 CP 列）得到 model_size=512、DP=6，说明这一行主要靠 64 级 PP 摊开模型，数据副本只有 6 路。第三，§5.1 写 “a trillion parameters”，Abstract 写 “1 trillion parameters”。

论文正文还给两个端到端估算：175B 的 GPT-3 在 1024 张 A100 上、batch 1536、每卡 140 TFLOP/s，训练 300B token 约 34 天；1T 模型在 3072 张 A100 上、每卡 163 TFLOP/s，训练 450B token 约 84 天（§5.1）。这些吞吐是端到端数字，包含数据加载、optimizer step 与通信，不是单个 GEMM 的峰值（§5.1）。

读完这节你应该知道：哪些事该找 Megatron-Core，哪些事该找 PAI patch，以及升级时会在哪里断。

所以“PAI 不改 Megatron 核心”要加限定词。它的 mcore 主路径是外挂：不把新模块写进 core，只通过 PYTHONPATH 把 megatron_patch 和固定的 Megatron 快照拼在一起（run_mcore_qwen3.sh:6），Qwen3 的 pretrain_qwen.py 最终调用 Megatron-LM 的 pretrain（megatron.training）。但仓库里同时存在两份“旧世界”：megatron_patch/training.py 是旧版训练循环的副本，megatron_patch/initialize.py 是旧版初始化的副本；它们服务的是仍按旧 API 写的 legacy 示例。把 PAI 理解成“纯 patch、绝不含训练循环”会漏掉这条路径；而 fixes/ 下的 patch 又会直接改 core，强度见下。

版本耦合：两处硬绑定 + 一处可选 patch

第一处是启动脚本的 PYTHONPATH：run_mcore_qwen3.sh:6 指向 backends/megatron/Megatron-LM-250624；这个子模块在 PAI 仓库里的 gitlink 是 6a6cd478839d90cf09a837adf8c79cbc844bc920。同一仓库虽然还放了 250908 快照，但 Qwen3 示例没有用它。

第二处是代码接口。PAI 的 pretrain_qwen.py:139-144 用位置参数调用 pretrain(provider, model_provider, ModelType, forward_step, extra_args_provider=...)；本页读到的 Megatron-LM 70f7a26 里，pretrain 的第一个参数已经是 cfg_container（training.py:1746）。把 PAI 的入口直接套到新 core 上，先撞的是签名，不是模型。

第三处是可选 fixes。megatron_patch/fixes/yarn_args/README.md:5-7 明确写“仅适用于 Megatron-LM-250328”，并给出 commit 6ba97dd37150a6bfba03d31808674211cf2a4d0d；同属 fixes/ 的 optimizer_offloading 也限定同一版本。这类 patch 会直接改 core 文件，和“外挂 adapter”不是同一种耦合强度。

模型适配成本：不是改一个配置项

接一个新模型，至少要准备四类东西：模型侧的 layer spec / model provider（pretrain_qwen.py:39-131 是一个可读的样板）；tokenizer 类型与词表 padding（megatron_patch/arguments.py:156 注册 --patch-tokenizer-type，megatron_patch/tokenizer/__init__.py:46-54 按类型构造）；数据 provider（megatron_patch/data/__init__.py:45-55）；以及 HF↔MCore 的分布式转换脚本（toolkits/distributed_checkpoints_convertor/README.md:3-5）。这些都不是 TransformerConfig 里的一个布尔开关。

权重转换还有自己的约束：转换器基于 torch_dist，其 README 要求训练时把 --ckpt-format torch 显式改成 torch_dist；转换脚本不会复制 tokenizer，需要手动拷配置；当前只测了 GPU 路径（toolkits/distributed_checkpoints_convertor/README.md:10-17）。这里又有一层版本差：本页读到的 core 70f7a26 里 CheckpointConfig.ckpt_format 默认已经是 torch_dist（megatron/training/config/training_config.py:603），而 PAI 转换器 README 仍按旧默认 torch 写。转换完不等于可以直接训练，还要把 tokenizer、配置和启动脚本补齐。

框架不解决的问题

• 不替你选并行度。 core 的入口接收调用者给的 TP/PP/CP/EP 尺寸，只校验整除（parallel_state.py:892-893）；模型多大、切几路、PP 多少 stage，仍要人来定。
• 不保证每个模型都支持。 GPT 有 GPTModelConfig 和 gpt_builder，T5/BERT/Mamba/VLM 各有自己的入口（pretrain_gpt.py、examples/t5/pretrain_t5.py、examples/bert/pretrain_bert.py、pretrain_mamba.py、pretrain_vlm.py）；新结构仍要写 spec。
• 不保存内置文本 dataloader 的批内位置。 只有支持该接口的 Energon 多模态 dataloader 会存；文本训练靠已消费样本数和数据顺序恢复（checkpointing.py:677-678）。
• PAI 示例不是完整断点续训。 Qwen3 脚本显式传 --no-load-optim、--no-load-rng、--no-save-optim（run_mcore_qwen3.sh:474-475、:486），它适合从已有权重继续训练，不是从 optimizer 状态原地恢复。
• 数据质量、学习率、集群调度不在框架内。 这些环节由数据管线、超参搜索和调度器负责；训练框架只负责把给定的配置跑成可复现的 step 与 checkpoint。

什么时候别用这套：只需要单卡或纯 DDP 的小模型，Megatron 的 CUDA 断言、并行组和分片 checkpoint 只会增加约束（initialize.py:87-89）；需要比 PAI 固定快照更新的 Megatron 特性，又不想 port patch，直接读 core 源码比改 PAI 入口更省事；模型不在 PAI 模型库里，适配成本要先按上面四类文件估一遍。