项目背景与定位

1.1 一句话定位

Megatron-LM 是 NVIDIA 出品的、专门为千亿参数级别 Transformer 在 GPU 集群上高效预训练而设计的代码库。它不是一个开箱即用的"训练 SaaS"，而更像 NVIDIA 自己用来跑论文、跑客户案例、再不断把"我们当时这样写最快"沉淀回 GitHub 的工程参考实现。

如果你看过 GPT-3、Megatron-Turing NLG 530B、Llama 2、Mixtral、Nemotron、Grok-1 等模型的技术报告，里面频繁出现的"我们采用 TP=8、PP=16、SP、selective activation recomputation……"这些术语，绝大多数都是 Megatron-LM 这一行行代码里冒出来的。它的影响力体现在两个层面：一是直接被使用，二是被其他框架（DeepSpeed-Megatron、Megatron-DeepSpeed、Megatron-Core、NeMo、ColossalAI 的部分模块）当作底座或抄作业的目标。

1.2 为什么要存在这样一个项目

要回答"为什么要专门写一个 Megatron-LM"，得回到 2019 年的时间点。当时业界的常态是：

• PyTorch DDP（数据并行）已经成熟，可以把同一个模型复制到很多张卡上，每张卡看不同的数据；
• 但 GPT-2 1.5B 那一代模型，单卡显存已经塞不下完整模型了，更别说优化器状态和激活值；
• 直接用 ZeRO 还很初期，PyTorch 自带的并行能力非常有限。

Megatron 当时给出的核心答案是"沿着 Transformer 自身的结构去切"：

1. Attention 的 QKV 投影和 MLP 的两层 Linear 可以按列 / 按行切到不同 GPU，靠 AllReduce 拼起来 —— 这就是张量并行 (Tensor Parallelism, TP)；
2. 不同 Transformer Layer 可以拆到不同 GPU 上，让流水线推着 micro-batch 走 —— 这就是流水线并行 (Pipeline Parallelism, PP)；
3. 再叠加传统的数据并行 (DP)，构成"3D 并行"。

所谓 3D 并行的总卡数 $W$ 满足：

$$W = \mathrm{TP} \times \mathrm{PP} \times \mathrm{DP}$$

例如 GPT-3 175B 的标准配置（见 examples/gpt3/train_gpt3_175b_distributed.sh）就是 $\mathrm{TP}=8,\;\mathrm{PP}=16,\;\mathrm{DP}=12$，即 $W = 8\times16\times12 = 1536$ 张卡。

之后几年里，Megatron-LM 又陆续把 Sequence Parallelism（SP）、Context Parallelism（CP）、Expert Parallelism（EP）、Selective Activation Recomputation、Distributed Optimizer 等新点子加了进来。可以说，当下大模型工业训练里能听到的所有"并行词汇"，几乎都能在它的代码里找到对应实现。

1.3 它和这些项目是什么关系

初学者最容易混淆的就是几个"长得很像"的名字：

1.4 设计哲学：硬件友好优先

通读 Megatron 源码会感觉到非常浓的"NVIDIA 工程师味"：

• 不掩盖硬件细节。NCCL、CUDA stream、tensor 形状、通信拓扑全都写在代码里能看见的位置，没有过度的中间抽象层。
• 优先 H100/A100，其次才是其他。fused RMSNorm、fused softmax、FlashAttention-3、TransformerEngine FP8 这些算子在 H100 上一路开到顶。
• 追求可复现的极限性能。仓库里到处是"用 X 个 GPU 训这个尺寸的模型每秒能算多少 tokens"的硬基准，PR 合并前都要跑性能 regression。
• 不追求"封装漂亮"。函数签名经常 20+ 个参数；同一个概念在不同地方有不同名字；CLI flag 数量惊人。它接受这一切，换来的是"读到代码就知道在哪条 stream 上算什么"。

这意味着两件事：你想用它，必须接受陡峭的上手曲线；但只要你愿意读，它愿意把所有事情告诉你。这恰恰是写一本中文 tutorial 的价值所在 —— 帮你跨过那条陡峭的曲线。

1.5 它适合用来做什么、不适合做什么

1.6 本书将怎么带你走完这趟旅程

接下来的章节会按这样的顺序展开：

• 第 2–3 章先把环境搭起来、跑通一个 345M 参数的小 GPT 预训练，让你有"这东西在我电脑上动起来了"的实感；
• 第 4 章把整本书后面会反复出现的关键词（rank、parallel group、micro-batch、pipeline stage）讲清楚，是后续所有章节的术语字典；
• 第 5–9 章是"硬核区"：配置系统、数据流水线、三种并行的实现细节；
• 第 10 章讲那些让你"训得起也训得快"的优化技巧；
• 第 11–12 章是工程师最关心的"出事了怎么办"以及"模型练好了怎么放出去用"；
• 第 13 章给你一张代码地图，下次自己读源码不再迷路。

读完这本书的目标，不是你能"从零写一个 Megatron-LM"，而是当 Megatron-LM 出现在你的项目里时，你敢动它、读得懂它、改得了它。这就够了。

1.7 2024–2026 重要更新时间线

Megatron-LM 不是 2020 那个 "GPT-3 论文配套代码" 了；最近两年加了一堆新东西。下表按时间倒序，每行附本书后续讨论位置：

1.8 Megatron-LM / Megatron-Core / NeMo：谁是谁

这三个名字经常混用，搞清楚不绕弯：

本书主线是 Megatron-LM 仓库本体 + 其 core/ 目录。后面提到 "MCore" 时指的就是 megatron/core/。