深度解读：Megatron-LM在千卡级集群中的超高效能

4.1 Megatron-LM：千卡级集群预训练的“硬核”框架

Megatron-LM 是一个基于 PyTorch 的分布式训练框架，专门面向 Transformer 大语言模型 （尤其是 GPT/LLaMA 类 decoder-only 架构）的超大规模训练。它在工程上追求“极致吞吐”，在千卡级甚至万卡级训练中，往往能取得非常强的性能与可扩展性。

但需要明确的是：Megatron-LM 的优势来自于它对训练路径的深度侵入式优化，这也带来了“框架与模型强耦合、抽象弱、可维护性差”等一系列工程代价。整体上，它像一台强力但不够精致的赛车： 性能很猛，但改装和维护成本高 。

1. 为什么需要 Megatron-LM：大模型训练的三大硬约束

训练大模型最先遇到的不是算法，而是三个硬约束： 显存、计算时间、并行策略的通信结构 。

1.1 显存约束：参数、梯度、优化器状态的“3 倍法则”

以 Adam/AdamW 为例，训练时显存主要来自三部分：

• 参数： $P$
• 梯度： $P$
• 优化器状态（以 Adam 为例有一阶矩和二阶矩）：约 $2P$

因此，单精度（FP32）下训练时的总量近似是：

$$\text{Memory}\approx 4P$$

如果用 bf16/fp16 存参数和梯度，但优化器状态常用 FP32 保存，估算会更复杂，但数量级仍然非常大。

案例：175B 参数模型（GPT-3 量级）

假设参数用 FP16/bf16（2 bytes）存储，优化器状态 FP32（4 bytes）：

• 参数： $175B\times 2\approx 350\text{GB}$
• 梯度： $175B\times 2\approx 350\text{GB}$
• Adam 一阶矩： $175B\times 4\approx 700\text{GB}$
• Adam 二阶矩： $175B\times 4\approx 700\text{GB}$

总计约：

$$350+350+700+700=2100\text{GB}\approx 2.1\text{TB}$$

这还没包含激活值、通信 buffer、临时张量等。现实训练中，整体需求常更高。

1.2 计算时间约束：单卡训练几乎不可行

即使把模型“塞进”单卡（例如通过 CPU offload），计算量仍会导致训练时间爆炸。

一个简化的经验量级估算：Transformer 的训练 FLOPs 大致与 参数量、序列长度、训练 token 数 成正相关。对于百亿到千亿参数模型，单卡训练会是“年”为单位的工程。

因此大模型训练必须依赖多卡并行，核心问题就变成：
怎么切模型、怎么切数据、怎么切序列，才能让 GPU 忙起来且通信不拖后腿？

1.3 并行策略约束：通信模式决定可扩展上限

不同并行策略对应不同的通信模式（AllReduce / AllGather / ReduceScatter / P2P 等），而通信结构通常决定了训练在大规模集群上能否线性扩展。

2. 数据并行（DP）：最简单、最强扩展，但受 batch 限制

2.1 DP 的基本方式

假设有 $N$ 张 GPU，每张 GPU 保存完整模型副本。每个 step 的 batch 被分成 $N$ 份，每张卡计算本地梯度，然后通过 AllReduce 得到全局平均梯度：

$$g=\frac{1}{N}\sum _{i=1}^N{g}_i$$

更新参数：

$$\theta \leftarrow \theta -\eta g$$

2.2 AllReduce 是什么

AllReduce 的目标是：每张卡输入一段张量 $x_i$ ，输出相同的归约结果（例如求和）：

$$y=\sum _{i=1}^N{x}_i$$

然后每张卡都拿到同一个 $y$ （或 $y/N$ ）。

2.3 DP 的两个主要瓶颈

1） batch 变小导致 GPU 利用率下降
当 GPU 数增大、全局 batch 不变时，单卡 micro-batch 过小，kernel 启动开销、通信占比上升，吞吐下滑。

2） 设备数受 batch 可扩展性限制
DP 往往需要足够大的 batch 才能把更多 GPU 吃满。否则进入“加卡不加速”的区域。

3. 激活检查点（Activation Checkpointing）：用算力换显存

DP 解决不了“模型塞不进显存”的问题时，常用激活检查点降低激活显存。

3.1 核心思想

前向传播不保存所有中间激活，只保存少量检查点。反向传播时，对缺失的激活进行重算。

假设一段网络 $f=f_k\circ \cdots \circ f_1$ ，标准反传要保存每层激活 $h_j$ 。检查点策略只保存部分层的激活，反向时重算缺失部分。

3.2 代价

• 优点：显存显著下降
• 缺点：额外的前向重算成本，训练时间上升

4. 模型并行（MP）：把模型切开，解决“单卡放不下”的根本问题

模型并行的两大主流形态：

• 张量并行（TP）： 层内切分（Intra-layer）
• 流水线并行（PP）： 层间切分（Inter-layer）

Megatron-LM 的核心贡献之一就是把 TP+PP 做得非常工程化，并在大规模训练中表现出色。

5. 张量并行（TP）：把大矩阵拆到多卡上算

张量并行的典型场景：线性层的权重矩阵非常大，单卡放不下，或单卡算太慢。

以一个线性层为例：

$$Y=XA$$

其中 $X\in {\mathbb{R}}^{bs\times h}$ （把 batch 和 seq 合并）， $A\in {\mathbb{R}}^{h\times h^{\prime }}$ 。

TP 的关键是：把 $A$ 沿某个维度切分，让不同 GPU 分担计算。

5.1 行并行（Row Parallelism）

把 $A$ 按“行”切成 $A_1,A_2$ ：

$$A=\left[\begin{array}{c}{A}_1\\ A_2\end{array}\right]$$

同时把 $X$ 按“列”切成 $[X_1{X}_2]$ ，则：

$$XA=[X_1{X}_2]\left[\begin{array}{c}{A}_1\\ A_2\end{array}\right]=X_1{A}_1+X_2{A}_2$$

两个 GPU 各自算一部分后，需要把结果做加和（通常通过 AllReduce）得到完整 $Y$ 。

5.2 列并行（Column Parallelism）

把 $A$ 按“列”切成 $A_1,A_2$ ：

$$A=[A_1{A}_2]$$

则输出自然也被切成两段：

$$Y=XA=[X{A}_{1}X{A}_2]=[Y_1{Y}_2]$$

此时每张卡可以独立计算自己的输出分块，最后只需要在需要“完整向量”的地方再做拼接或通信。

6. Transformer 中的 TP：Attention 和 MLP 都能切

Transformer block 主要是两大计算密集模块：

1）Self-Attention
2）MLP/FFN

Megatron 的经典做法是：对两者都做 TP 切分，并插入必要的同步通信点。

6.1 先看 MLP：为什么常用“第一层列切，第二层行切”

标准 FFN（以隐藏维度 $h$ ，中间维度 $4h$ 为例）：

$$\text{FFN}(x)=W_2\cdot \sigma (W_{1}x)$$

其中 $W_1\in {\mathbb{R}}^{h\times 4h}$ ， $W_2\in {\mathbb{R}}^{4h\times h}$ ， $\sigma$ 常是 GeLU。

关键事实：非线性不可分

如果把第一层做行切，需要在 GeLU 前把各卡的部分结果合并，否则：

$$\sigma (u_1+u_2)\ne \sigma (u_1)+\sigma (u_2)$$

这会引入额外同步点，通信更频繁。

因此常见设计：

• 第一层 $W_1$ ： 列并行 （各卡独立算出自己那部分 $4h/N$ 的中间激活，并各自做 GeLU）
• 第二层 $W_2$ ： 行并行 （把 $4h$ 维拆开乘回 $h$ ，最后 AllReduce 汇总得到完整输出）

这能最大化 GeLU 前后的独立计算，减少不必要的通信。

6.2 Self-Attention 的 TP：按 head 维度天然可并行

多头注意力可以写成（忽略 mask 细节）：

$$\text{Attn}(X)=\text{Concat}({\text{head}}_1,\dots ,{\text{head}}_H)W_O$$

每个 head：

$${\text{head}}_i=\text{Softmax}\left(\frac{Q_i{K}_i^{\top }}{\sqrt{d}}\right)V_i$$

天然并行点：不同 head 之间互不依赖。
因此可以把 head 分配到不同 GPU 上，使每个 GPU 负责一部分 head 的 $Q,K,V$ 投影和注意力计算。最后在输出投影或后续层需要完整向量时做合并。

一个常见要求是：head 数尽量能被 TP 的 GPU 数整除，以平衡负载。

7. 输出层 softmax / cross-entropy 的并行：词表太大时必须优化通信

输出层 logits：

$$z=h{W}_{\text{vocab}}$$

如果词表大小为 $v$ ，logits 张量规模约为 $bsv$ 。当 $v$ 很大时，直接 AllGather 拼完整 logits 再做 softmax，会产生巨大通信开销。

softmax：

$$p_j=\frac{e^{z_j}}{{\sum }_{k=1}^v{e}^{z_k}}$$

关键在于分母的归一化项：

$$Z=\sum _{k=1}^v{e}^{z_k}$$

如果词表被分片到不同 GPU 上，每个 GPU 先算自己分片的部分和：

$$Z_i=\sum _{k\in {\mathcal{V}}_i}{e}^{z_k}$$

再 AllReduce 得到全局：

$$Z=\sum _{i=1}^N{Z}_i$$

这样就不需要 AllGather $bsv$ 级别的大张量，而只需要 AllReduce $bs$ 规模的归一化量，通信从“随词表线性增长”降到“随 batch*seq 增长”。

8. 流水线并行（PP）：按层切分，降低单卡显存，但会引入空泡

8.1 PP 的基本方式

把网络层按顺序切成多个 stage，每张卡存一段层。例如 8 层、2 卡：

• GPU0： $L_0$ - $L_3$
• GPU1： $L_4$ - $L_7$

前向时激活从 GPU0 传到 GPU1，反向时梯度再传回来。

8.2 空泡（Bubble）问题：流水线会出现“有人在等人”

如果一个 batch 不切 micro-batch，那么 GPU0 前向结束后要等 GPU1 前向结束才能进入反向，GPU1 反向时 GPU0 又可能空闲。

解决方式：把 batch 切成多个 micro-batch，流水线交错执行。

9. 主流 PP 调度：GPipe、PipeDream、Virtual Pipeline

9.1 GPipe：micro-batch 填充流水线，但空泡明显

GPipe 用 micro-batch 把流水线填满，前向全部跑完再反向全部跑完。空泡率较高。

9.2 PipeDream：1F1B 调度减少激活占用

PipeDream 在稳定阶段形成 1F1B（一个前向、一个反向交替），显存更省（激活保存压力更小），空泡改善但仍受 stage 数影响。

9.3 Virtual Pipeline：用更多 stage 细分减少空泡

Virtual Pipeline 的核心思想：在 device 数不变时，把 stage 切得更细（增加虚拟 stage 数），让设备之间更频繁通信以减少等待，从而降低空泡率、缩短 step 时间。

案例（示意）
16 层，4 张卡：

• 普通 PP：每卡 4 层
• Virtual Pipeline（虚拟 stage=2）：每卡分两段，每段 2 层，层分配变得交错，从而在调度上更灵活、空泡更小，但 P2P 通信次数增加。

10. 3D 并行：DP + TP + PP 的混合范式

实际训练通常不会只用一种并行方式，而是混合使用：

• DP：扩展到更多 GPU，吞吐好
• TP：解决单层矩阵太大问题
• PP：解决层数/整体参数太大问题

10.1 三者的典型对比

• 显存效率：模型并行（TP/PP）优于 DP
• 通信效率：PP 的通信量通常更低（P2P），TP 通信频繁且容易成为瓶颈
• 计算效率：DP 最容易线性扩展；TP 通信频繁导致效率可能较差；PP 受空泡影响

11. 4D 并行：在 3D 基础上加入 Context Parallel（CP）

当上下文长度极长（长序列训练、超长上下文 SFT/RL 等），单卡的序列维度计算与激活显存会非常大。此时除了 DP/TP/PP，还会引入 CP（Context Parallel）在序列维度上进一步切分计算与激活。

直觉上可以理解为：

• TP 切“特征维/矩阵维”
• PP 切“层维”
• CP 切“序列维”
• DP 切“数据维”

因此被称为 4D 并行 。

12. Megatron-LM 的工程优劣势：为什么“又爱又恨”

12.1 为什么它在大规模预训练中依然强势

• 吞吐很强 ：在千卡以上集群中，Megatron 在 kernel 融合、通信调度、并行组合方面经过大量实战打磨，常能比通用框架更快。
• 并行形态齐全 ：DP/TP/PP（以及更进一步的 EP、CP）组合灵活，能覆盖从几十亿到千亿参数的训练需求。
• 大规模可扩展性更稳 ：当规模上去后，很多“通用抽象框架”的额外开销会放大，而 Megatron 的“贴地飞行”反而占优。

因此在千卡以上预训练场景里，它经常是最优选择之一。

12.2 典型缺点：性能背后的代价

1） 框架抽象弱，模型与框架强耦合
缺少清晰的分层与模块化抽象，很多实现细节直接侵入模型结构。
结果是：模型很难与框架解耦，需要手动切割并适配，且更偏向 GPT/LLaMA 类。

2） 通信精度问题（尤其某些 ring AllReduce 路径）
在混合精度下，某些通信路径可能引入明显数值误差，训练稳定性受影响，尤其在对数值敏感的阶段（例如 RLHF 的 reward/advantage 波动更大）。

3） 推理侧稳定性问题可能反噬训练（例如 KV cache 相关 bug）
在 RLHF/RLVR 训练中，推理链路异常往往会直接导致 rollout 质量崩溃或训练失败。KV cache 的错误会放大这种风险。

4） 混合精度实现差异：bf16 没有 master weight 的设计取舍
某些实现路径下不保留 FP32 master weights，可能降低数值稳定性或影响某些优化器行为。

5） 显存申请策略“粗犷”，叠加 PyTorch allocator 容易 OOM
碎片化、临时 buffer 峰值、通信 buffer 与激活峰值重叠等问题，常导致“理论上能跑、实际上 OOM”。

13. 规模选择建议：千卡以上 Megatron，千卡以内选择更灵活

• 千卡以上 ：Megatron-LM 往往是极具竞争力甚至最优的工程选择之一（吞吐、并行组合、实战经验积累明显）。
• 千卡以内 ：如果目标不是极限吞吐，FSDP 或 Torchtitan 这类更模块化、更易维护的方案通常更省心。

14. 一个贯穿式案例：为什么需要 3D/4D 并行，而不是只堆 DP

假设训练一个中等偏大的模型，单卡放不下完整权重与优化器状态：

• 只用 DP：每卡都要存完整模型，显存不够，直接失败。
• 加 TP：把关键线性层权重切开，单卡权重压力下降，但通信变多。
• 再加 PP：把层切开，单卡存的层数减少，权重/优化器进一步下降，但会出现空泡，需要 micro-batch 调度。
• 训练长上下文：激活与注意力计算沿序列维度爆炸，引入 CP 把序列维度切开，降低单卡序列压力。

最终形成 DP+TP+PP(+CP) 的组合，才能把训练跑起来且吞吐可接受。

15. 小结：Megatron-LM 的定位

Megatron-LM 的核心价值可以用一句话概括：

它不是最优雅的框架，但在千卡级预训练里，它往往是最快的那台机器。

• 追求吞吐与规模扩展：Megatron 很强
• 追求框架抽象、模型解耦与可维护性：Megatron 代价很高
• 预训练（尤其千卡以上）：Megatron 依然是主力选项
• 中小规模或需要快速迭代：FSDP/Torchtitan 等方案更合适