数据截至 (上游 commit e79cb4c1bae1)
Megatron-LM — 架构与原理
30 秒导读： Megatron-LM 是 NVIDIA 的大规模 Transformer 预训练框架，也是「张量并行、流水并行、序列并行」这些技术的策源地（2019 年论文的代码实体）。它回答的核心问题是：一个几百亿、几千亿参数的模型放不进一张 GPU，怎么把它沿不同维度切开摊到几千张卡上，还让通信开销尽量小。读懂它，就读懂了今天几乎所有大模型训练框架（包括 PyTorch 原生并行）的设计语汇。
1. 这是什么（零基础也能懂）
一句话定义
Megatron-LM 是一个给超大 Transformer 模型做预训练的框架：它把「模型的一层怎么切到多张卡上算」做成库（Megatron Core，megatron/core/），再配好一整套训练脚本（pretrain_gpt.py 等）告诉你怎么在几千张 GPU 上把 GPT 类模型从头训出来。
它要解决谁的什么问题
假设你有一块 80GB 显存的 H100，和一个 175B 参数的模型：
- 光参数（bf16）就要 350GB——一张卡连模型都装不下。
- 加上梯度（bf16 又是 350GB）和 Adam 优化器状态（fp32 主参数 + 两个动量，约 1400GB），单机想都别想。
- 就算装下了，单卡算一个 batch 要算到天荒地老。
于是必须回答：沿哪个维度切？ 切参数？切层？切数据？切序列？Megatron 的历史贡献就是把「切法」系统化成几个正交维度，每个维度给出通信量上最优的切分点。今天是 所有「3D 并行」「5D 并行」说法的源头。
它能做什么
| 能力 | 具体支持 | 在哪 |
|---|---|---|
| 张量并行（TP） | 单层权重矩阵按列/行切到多卡 | megatron/core/tensor_parallel/layers.py |
| 流水并行（PP） | 层按段分到多卡，interleaved 1F1B 调度 | megatron/core/pipeline_parallel/schedules.py |
| 序列并行（SP） | LayerNorm/Dropout 处按序列维切激活 | megatron/core/tensor_parallel/mappings.py |
| 上下文并行（CP） | 长序列的 attention 按序列分片计算 | megatron/core/context_parallel/ |
| 专家并行（EP） | MoE 专家摊到多卡，All-to-All 调度 token | megatron/core/transformer/moe/ |
| 分布式优化器 | ZeRO 式优化器状态分片 | megatron/core/optimizer/distrib_optimizer.py |
| 混合精度 | FP16/BF16/FP8/FP4 | megatron/core/fp8_utils.py 等 |
| 模型 | GPT、T5、BERT、Mamba、混合架构、多模态 | megatron/core/models/ |
用起来什么样
最小形态是一个普通 torchrun 启动的脚本。核心心智模型从命令行参数就能读出来——五种并行各一个 size：
# 示意（基于 examples/gpt3/ 与 README 的真实用法精简）
torchrun --nproc_per_node 8 --nnodes 64 pretrain_gpt.py \
--tensor-model-parallel-size 8 \ # 一层切 8 份（节点内）
--pipeline-model-parallel-size 8 \ # 模型切 8 段（跨节点）
--num-layers 64 --hidden-size 8192 \
--num-experts 128 --moe-router-topk 8 # MoE 可选
库用法（不跑整套脚本）则是 examples/run_simple_mcore_train_loop.py:36 演示的样子：先 torch.distributed.init_process_group，再 parallel_state.initialize_model_parallel(tp_size, pp_size) 建立进程组网格，之后模型层的构造函数自动按这个网格切自己。
一句话直觉
把 Megatron 想成「一套切蛋糕的标准刀法」。 蛋糕是「参数 + 激活 + 优化器状态」三样东西，刀法有五种（TP/PP/DP/CP/EP），每把刀切的方向不同、产生的通信模式不同。框架的全部内容就是：每把刀怎么下刀（数学）、切完各块怎么拼回去（通信原语）、以及怎么保证多块拼起来的结果和没切时一模一样（正确性）。
2. 顶层全景（它大概怎么转）
2.1 代码的两个世界
仓库分两层，定位完全不同：
megatron/core/— 可组合的库：并行原语、Transformer 积木、优化器。框架开发者拿它搭自己的训练系统（NVIDIA 自家的 NeMo、外面无数框架都建立其上）。megatron/training/+pretrain_gpt.py— 参考训练脚本：用 core 组装出完整预训练流程，是「标准用法」的活教材。
本系列六章几乎全在 core 里打转——那里是并行机制本体。
2.2 顶层部件图
一个训练 step 里各部件的关系（从左到右是数据流方向）：
┌─────────────────────────────────────────────┐
pretrain_gpt.py │ pretrain() 训练主循环 │
(入口脚本) ───►│ megatron/training/training.py:1500 │
│ │
│ train_step (training.py:3010) │
│ │ │
│ ▼ │
│ forward_backward_func ←── 三选一调度器 │
│ (pipeline_parallel/schedules.py:53) │
│ │ │
└────┼────────────────────────────────────────┘
▼
┌───────────────────────────────┐
│ GPTModel 前向 (一个 PP 段) │
│ embedding → N×TransformerLayer│
│ → output_layer → CE loss │
└───────┬───────────────┬───────┘
│ │
┌─────────▼────┐ ┌──────▼─────────┐
│ TP 原语 │ │ MoE 三件套 │
│ Column/Row │ │ Router→Dispatch│
│ ParallelLinear│ │ →Experts→Combine│
└───────┬───────┘ └──────┬─────────┘
│ │
▼ ▼
┌─────────────────────────────┐
│ 进程组网格 parallel_state │
│ (谁在哪个 TP/PP/DP/EP 组) │
└─────────────────────────────┘
▲
┌─────────────┴───────────────┐
│ DistributedOptimizer │
│ 梯度 reduce-scatter / │
│ 参数 all-gather / 分片 step │
└─────────────────────────────┘
怎么读这张图： 上半部是「调度」（哪个 microbatch 在哪个阶段算前向还是反向），中间是「模型」（每一层内部怎么被 TP/MoE 切开），底部是「底座」（所有跨卡行为最终都落到 parallel_state 建的进程组上；训练收尾时分布式优化器再沿 DP 组做一次参数同步）。
2.3 部件一句话职责
| 部件 | 干什么 | 在哪个文件 |
|---|---|---|
parallel_state | 建 TP/PP/DP/CP/EP 全部进程组，回答「我是谁、我的队友是谁」 | megatron/core/parallel_state.py:600（initialize_model_parallel） |
| TP 通信原语 | copy/reduce/scatter/gather 四个 autograd 函数，前向反向互为对偶 | megatron/core/tensor_parallel/mappings.py:199-260 |
ColumnParallelLinear / RowParallelLinear | 权重矩阵按列/行切分的线性层 | megatron/core/tensor_parallel/layers.py:986 / :1382 |
get_forward_backward_func | 按 PP/VP 配置选三种调度器之一 | megatron/core/pipeline_parallel/schedules.py:53 |
| interleaved 调度 | 虚拟流水下的 warmup→1F1B→cooldown 编排 | megatron/core/pipeline_parallel/schedules.py:1019 |
P2PCommunicator | 相邻流水段之间收发激活/梯度 | megatron/core/pipeline_parallel/p2p_communication.py:145 |
DistributedOptimizer | 优化器状态按 DP rank 分片（ZeRO） | megatron/core/optimizer/distrib_optimizer.py:113 |
_ParamAndGradBuffer | 连续梯度/参数 buffer，分桶做异步通信重叠 | megatron/core/distributed/param_and_grad_buffer.py:1051 |
MoELayer 三件套 | Router 选专家 → Dispatcher 搬 token → GroupedMLP 算 | megatron/core/transformer/moe/moe_layer.py:215 |
GPTModel | embedding + TransformerBlock + 并行输出层的组装 | megatron/core/models/gpt/gpt_model.py:52 |
pretrain / train_step | 训练主循环与一个 step 的完整编排 | megatron/training/training.py:1500 / :3010 |
2.4 主线走一遍（一个训练 step）
下面这条线是 train_step（megatron/training/training.py:3015）的骨架，五种并行各在其位：
① 清零梯度
model_chunk.zero_grad_buffer() —— 梯度 buffer（不是 param.grad）清零
│
▼
② 前向 + 反向
forward_backward_func(...)（training.py:3105 调用）
→ 调度器按 microbatch 逐个跑：每卡只算自己 PP 段、每段内部
TP 各卡算自己那 1/tp 的权重、MoE 层 token 被 All-to-All 送走又收回
→ 反向时梯度落进连续 grad buffer，bucket 满了就异步 reduce-scatter
│
▼
③ 优化器 step
DistributedOptimizer.step()
→ 每卡只更新自己拥有的那一段参数（fp32 主副本）
→ 更新完把新参数写回 buffer，all-gather 让所有 DP rank 拿到全量新参数
│
▼
④ 下一个 step
这条线最值得记住的一点： 数据并行在 Megatron 里没有独立的 all-reduce 步骤。梯度同步被拆成了「反向过程中分桶 reduce-scatter」（每卡只留 1/DP 的梯度）+「step 后 all-gather 参数」，这就是 ZeRO-1 的工程形态，藏进了 buffer 的 hook 里。
3. 阅读地图（建议顺序）
六章由浅入深。时间有限的话，读 01 → 02 → 03 就能抓住 Megatron 区别于其他训练框架的全部 要害。
| 顺序 | 章节 | 讲什么 | 适合谁 |
|---|---|---|---|
| 1 | 01-parallel-dimensions.md | 五种并行各切什么、进程组网格怎么建、训练入口 | 所有人必读，这是全局地图 |
| 2 | 02-tensor-parallel.md | Column/Row 切分数学、四个通信原语、MLP/Attention 怎么配对 | 想懂「Megatron 论文」的人 |
| 3 | 03-pipeline-parallel.md | 1F1B、interleaved 调度、P2P 通信 | 关心气泡率与吞吐的人 |
| 4 | 04-distributed-optimizer.md | ZeRO 式分片、grad buffer、reduce-scatter/all-gather | 关心显存与 DP 通信的人 |
| 5 | 05-moe.md | Router、All-to-All token 调度、GroupedMLP、负载均衡 | 做 MoE 模型的人 |
| 6 | 06-gpt-model.md | GPTModel 组装、一个 step 端到端、输出层与词表并行 | 想把前面五章串起来的人 |
4. 巧妙之处（可借鉴的技术）
每条先白话点出妙在哪，细节在各章展开。
-
通信原语的「前向/反向对偶」。 张量并行只需要四个 autograd 函数就封死所有通信：copy（前向恒等、反向 all-reduce）与 reduce（前向 all-reduce、反向恒等）互为对偶，scatter 与 gather 同理（
megatron/core/tensor_parallel/mappings.py:199-260）。写层的人几乎不用想反向通信——autograd 自动配好。见 02 章。 -
Column→Row 配对让 MLP 只做一次通信。 第一个矩阵按列切、第二个按行切，中间激活不用同步，整个 MLP 块前向只有一次 all-reduce（在 RowParallelLinear 出口）。这是 2019 年论文的核心观察，落在
MLP.__init__里gather_output=False+input_is_parallel=True两个开关上（megatron/core/transformer/mlp.py:219、:240）。 -
进程组用「笛卡尔积生成器」统一描述。 不分叉写五种并行的建组逻辑，而是一个
RankGenerator按 order 字符串（"tp-cp-ep-dp-pp"）生成各维度的 rank 列表（megatron/core/parallel_state.py:465、:858）。加新维度（如 GTP remat）只是给 order 加个 token。 -
interleaved 调度把气泡再砍一刀。 每张卡持有多个不连续的「模型块」（virtual pipeline），warmup 期更短，气泡从 (p-1)/m 降到 (p-1)/(v·m)（
megatron/core/pipeline_parallel/schedules.py:929的 warmup 公式）。见 03 章。 -
梯度 buffer 的「按段归属」让 ZeRO 无需切参数边界。 DP rank 拥有的是连续 grad buffer 的一段区间，不是「某些参数」——一个参数可以被两个 rank 各管一半。
Range类 + 四重区间映射把这个做 得干净利落（megatron/core/optimizer/distrib_optimizer.py:78、:139）。见 04 章。 -
MoE 的 dispatch/combine 是教科书式四拍。 route → dispatch（permute + All-to-All）→ expert 算 → combine（All-to-All + unpermute），四步接口稳定在
MoELayer.forward（megatron/core/transformer/moe/moe_layer.py:637），调度器实现可换（AllGather / AllToAll / DeepEP / 混合 EP）。见 05 章。
5. 边界与局限
诚实地说，Megatron 不是万能钥匙。
- 单机没法玩。 整套设计假设多节点 GPU 集群；源码卡片也明说「nothing here is exercisable on one machine」（
aiRef/sources/megatron-lm.md）。没有 GPU、没有 NCCL，连单元测试大多都跑不了。 - 深度绑定 NVIDIA 栈。 NCCL、Transformer Engine、CUDA graph、FP8/FP4 都是 NVIDIA 专用；AMD/TPU 上它不是答案。
- TP 切分有理论上限。 张量并行的通信量在 hidden size 除以 tp 后仍随激活大小线性增长，因此 TP 一般只在节点内（NVLink 域）开到 8；再大就得靠流水/数据并行。
- 流水并行要求层数可整除、形状统一。 interleaved 调度要求
pipeline_model_parallel_size > 1（megatron/core/parallel_state.py:823-826附近的检查），且 P2P 收发双方张量形状必须事先谈拢（get_tensor_shapes，megatron/core/pipeline_parallel/schedules.py:2115）——变长序列要靠variable_seq_lengths特殊处理。 - 正确性责任在用户。 框架不阻止你配出荒谬的组合（比如 TP>8 跨节点）；它给你绳子和性能计数器，不给护栏。
GPTModel已标记 DEPRECATED。 本 commit 中GPTModel.__init__里就挂着弃用警告，指向HybridModel迁移指南（megatron/core/models/gpt/gpt_model.py:130-135）。它是存量主线，不是未来方向——本系列仍以它为讲解对象，因为它仍是理解 Megatron 模型的最短路径。
6. 横向对比
同书架上解决「大模型怎么训」的兄弟项目，取舍各不相同：
| 项目 | 与 Megatron 的关系 | 核心取舍 |
|---|---|---|
| DeepSpeed | 同一问题的微软答案 | ZeRO 论文出处；DP 分片做得更早更彻底，TP/PP 直接借用 Megatron 的实现思路 |
| megablocks | MoE 专家计算的另一种实现 | 用 block-sparse GEMM 表达专家计算；Megatron 选择 grouped GEMM（TE GroupedLinear），两者数学同构、工程路线不同 |
| Accelerate | 定位不同层 | Accelerate 是「帮你写分布式代码」的薄封装；Megatron 是「并行机制本身」的实现 |
| verl | 下游使用者 | RL 后训练框架，把 Megatron 当训练后端之一来调度（verl/workers/engine/）——读 Megatron 是读 verl 训练侧的地基 |
一句话：Megatron 是策源地——后来者的 TP/PP/SP 概念词汇大多从这里流出；它比 DeepSpeed 更「模型并行原生」，比 megablocks 更「全栈」，比 Accelerate 深得多。
7. 代码地图（导航索引）
给要跳进源码的人/agent：按主题列关键文件与符号。各章末尾有更细的本章地图。
| 主题 | 文件路径 | 符号名 |
|---|---|---|
| 进程组网格（一切并行的底座） | megatron/core/parallel_state.py | initialize_model_parallel、RankGenerator、get_tensor_model_parallel_group |
| TP 通信原语 | megatron/core/tensor_parallel/mappings.py | copy_to_tensor_model_parallel_region、reduce_from_tensor_model_parallel_region、scatter_to_tensor_model_parallel_region、gather_from_tensor_model_parallel_region |
| TP 线性层 | megatron/core/tensor_parallel/layers.py | ColumnParallelLinear、RowParallelLinear、VocabParallelEmbedding |
| 词表并行交叉熵 | megatron/core/tensor_parallel/cross_entropy.py | VocabParallelCrossEntropy、vocab_parallel_cross_entropy |
| 调度器选择 | megatron/core/pipeline_parallel/schedules.py | get_forward_backward_func |
| interleaved 1F1B | megatron/core/pipeline_parallel/schedules.py | forward_backward_pipelining_with_interleaving、get_pp_rank_microbatches |
| 流水 P2P | megatron/core/pipeline_parallel/p2p_communication.py | P2PCommunicator、_batched_p2p_ops |
| 分布式优化器 | megatron/core/optimizer/distrib_optimizer.py | DistributedOptimizer、Range、_build_model_gbuf_param_range_map |
| 梯度/参数 buffer | megatron/core/distributed/param_and_grad_buffer.py | _ParamAndGradBuffer、start_grad_sync、start_param_sync |
| MoE 层 | megatron/core/transformer/moe/moe_layer.py | MoELayer |
| MoE 路由 | megatron/core/transformer/moe/router.py | TopKRouter |
| token 调度 | megatron/core/transformer/moe/token_dispatcher.py | MoEAlltoAllTokenDispatcher、MoEAllGatherTokenDispatcher |
| 专家计算 | megatron/core/transformer/moe/experts.py | TEGroupedMLP、SequentialMLP |
| MLP 积木 | megatron/core/transformer/mlp.py | MLP |
| Transformer 层/块 | megatron/core/transformer/transformer_layer.py / transformer_block.py | TransformerLayer、TransformerBlock |
| GPT 模型 | megatron/core/models/gpt/gpt_model.py | GPTModel |
| 训练主循环 | 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 |