4.1 Megatron-LM：千卡级集群预训练的"硬核"框架
Megatron-LM 是一个基于 PyTorch 的分布式训练框架，专门面向 Transformer 大语言模型（尤其是 GPT/LLaMA 类 decoder-only 架构）的超大规模训练。它在工程上追求"极致吞吐"，在千卡级甚至万卡级训练中，往往能取得非常强的性能与可扩展性。
但需要明确的是：Megatron-LM 的优势来自于它对训练路径的深度侵入式优化，这也带来了"框架与模型强耦合、抽象弱、可维护性差"等一系列工程代价。整体上，它像一台强力但不够精致的赛车：性能很猛，但改装和维护成本高。
1. 为什么需要 Megatron-LM：大模型训练的三大硬约束
训练大模型最先遇到的不是算法，而是三个硬约束：显存、计算时间、并行策略的通信结构。
1.1 显存约束：参数、梯度、优化器状态的"3 倍法则"
以 Adam/AdamW 为例，训练时显存主要来自三部分：
- 参数：PPP
- 梯度：PPP
- 优化器状态（以 Adam 为例有一阶矩和二阶矩）：约 2P2P2P
因此，单精度（FP32）下训练时的总量近似是：
Memory≈4P \text{Memory} \approx 4P Memory≈4P
如果用 bf16/fp16 存参数和梯度，但优化器状态常用 FP32 保存，估算会更复杂，但数量级仍然非常大。
案例：175B 参数模型（GPT-3 量级）
假设参数用 FP16/bf16（2 bytes）存储，优化器状态 FP32（4 bytes）：
- 参数：175B×2≈350GB175B \times 2 \approx 350\text{GB}175B×2≈350GB
- 梯度：175B×2≈350GB175B \times 2 \approx 350\text{GB}175B×2≈350GB
- Adam 一阶矩：175B×4≈700GB175B \times 4 \approx 700\text{GB}175B×4≈700GB
- Adam 二阶矩：175B×4≈700GB175B \times 4 \approx 700\text{GB}175B×4≈700GB
总计约：
350+350+700+700=2100GB≈2.1TB 350 + 350 + 700 + 700 = 2100\text{GB} \approx 2.1\text{TB} 350+350+700+700=2100GB≈2.1TB
这还没包含激活值、通信 buffer、临时张量等。现实训练中，整体需求常更高。
1.2 计算时间约束：单卡训练几乎不可行
即使把模型"塞进"单卡（例如通过 CPU offload），计算量仍会导致训练时间爆炸。
一个简化的经验量级估算：Transformer 的训练 FLOPs 大致与 参数量、序列长度、训练 token 数成正相关。对于百亿到千亿参数模型，单卡训练会是"年"为单位的工程。
因此大模型训练必须依赖多卡并行，核心问题就变成：
怎么切模型、怎么切数据、怎么切序列，才能让 GPU 忙起来且通信不拖后腿？
1.3 并行策略约束：通信模式决定可扩展上限
不同并行策略对应不同的通信模式（AllReduce / AllGather / ReduceScatter / P2P 等），而通信结构通常决定了训练在大规模集群上能否线性扩展。
2. 数据并行（DP）：最简单、最强扩展，但受 batch 限制
2.1 DP 的基本方式
假设有 NNN 张 GPU，每张 GPU 保存完整模型副本。每个 step 的 batch 被分成 NNN 份，每张卡计算本地梯度，然后通过 AllReduce 得到全局平均梯度：
g=1N∑i=1Ngi g = \frac{1}{N}\sum_{i=1}^{N} g_i g=N1i=1∑Ngi
更新参数：
θ←θ−ηg \theta \leftarrow \theta - \eta g θ←θ−ηg
2.2 AllReduce 是什么
AllReduce 的目标是：每张卡输入一段张量 xix_ixi，输出相同的归约结果（例如求和）：
y=∑i=1Nxi y = \sum_{i=1}^{N} x_i y=i=1∑Nxi
然后每张卡都拿到同一个 yyy（或 y/Ny/Ny/N）。
2.3 DP 的两个主要瓶颈
1）batch 变小导致 GPU 利用率下降
当 GPU 数增大、全局 batch 不变时，单卡 micro-batch 过小，kernel 启动开销、通信占比上升，吞吐下滑。
2）设备数受 batch 可扩展性限制
DP 往往需要足够大的 batch 才能把更多 GPU 吃满。否则进入"加卡不加速"的区域。
3. 激活检查点（Activation Checkpointing）：用算力换显存
DP 解决不了"模型塞不进显存"的问题时，常用激活检查点降低激活显存。
3.1 核心思想
前向传播不保存所有中间激活，只保存少量检查点。反向传播时，对缺失的激活进行重算。
假设一段网络 f=fk∘⋯∘f1f = f_k \circ \cdots \circ f_1f=fk∘⋯∘f1，标准反传要保存每层激活 hjh_jhj。检查点策略只保存部分层的激活，反向时重算缺失部分。
3.2 代价
- 优点：显存显著下降
- 缺点：额外的前向重算成本，训练时间上升
4. 模型并行（MP）：把模型切开，解决"单卡放不下"的根本问题
模型并行的两大主流形态：
- 张量并行（TP）：层内切分（Intra-layer）
- 流水线并行（PP）：层间切分（Inter-layer）
Megatron-LM 的核心贡献之一就是把 TP+PP 做得非常工程化，并在大规模训练中表现出色。
5. 张量并行（TP）：把大矩阵拆到多卡上算
张量并行的典型场景：线性层的权重矩阵非常大，单卡放不下，或单卡算太慢。
以一个线性层为例：
Y=XA Y = X A Y=XA
其中 X∈Rbs×hX \in \mathbb{R}^{b s \times h}X∈Rbs×h（把 batch 和 seq 合并），A∈Rh×h′A \in \mathbb{R}^{h \times h'}A∈Rh×h′。
TP 的关键是：把 AAA 沿某个维度切分，让不同 GPU 分担计算。
5.1 行并行（Row Parallelism）
把 AAA 按"行"切成 A1,A2A_1, A_2A1,A2：
A=[A1A2] A = \begin{bmatrix} A_1 \\ A_2 \end{bmatrix} A=[A1A2]
同时把 XXX 按"列"切成 [X1 X2][X_1\; X_2][X1X2]，则：
XA=[X1 X2][A1A2]=X1A1+X2A2 X A = [X_1\; X_2] \begin{bmatrix} A_1 \\ A_2 \end{bmatrix} = X_1A_1 + X_2A_2 XA=[X1X2][A1A2]=X1A1+X2A2
两个 GPU 各自算一部分后，需要把结果做加和（通常通过 AllReduce）得到完整 YYY。
5.2 列并行（Column Parallelism）
把 AAA 按"列"切成 A1,A2A_1, A_2A1,A2：
A=[A1 A2] A = [A_1\; A_2] A=[A1A2]
则输出自然也被切成两段：
Y=XA=[XA1 XA2]=[Y1 Y2] Y = X A = [X A_1\; X A_2] = [Y_1\; Y_2] Y=XA=[XA1XA2]=[Y1Y2]
此时每张卡可以独立计算自己的输出分块，最后只需要在需要"完整向量"的地方再做拼接或通信。
6. Transformer 中的 TP：Attention 和 MLP 都能切
Transformer block 主要是两大计算密集模块：
1）Self-Attention
2）MLP/FFN
Megatron 的经典做法是：对两者都做 TP 切分，并插入必要的同步通信点。
6.1 先看 MLP：为什么常用"第一层列切，第二层行切"
标准 FFN（以隐藏维度 hhh，中间维度 4h4h4h 为例）：
FFN(x)=W2⋅σ(W1x) \text{FFN}(x) = W_2 \cdot \sigma(W_1 x) FFN(x)=W2⋅σ(W1x)
其中 W1∈Rh×4hW_1 \in \mathbb{R}^{h \times 4h}W1∈Rh×4h，W2∈R4h×hW_2 \in \mathbb{R}^{4h \times h}W2∈R4h×h，σ\sigmaσ 常是 GeLU。
关键事实：非线性不可分
如果把第一层做行切，需要在 GeLU 前把各卡的部分结果合并，否则：
σ(u1+u2)≠σ(u1)+σ(u2) \sigma(u_1 + u_2) \ne \sigma(u_1) + \sigma(u_2) σ(u1+u2)=σ(u1)+σ(u2)
这会引入额外同步点，通信更频繁。
因此常见设计：
- 第一层 W1W_1W1：列并行 （各卡独立算出自己那部分 4h/N4h/N4h/N 的中间激活，并各自做 GeLU）
- 第二层 W2W_2W2：行并行 （把 4h4h4h 维拆开乘回 hhh，最后 AllReduce 汇总得到完整输出）
这能最大化 GeLU 前后的独立计算，减少不必要的通信。
6.2 Self-Attention 的 TP：按 head 维度天然可并行
多头注意力可以写成（忽略 mask 细节）：
Attn(X)=Concat(head1,...,headH)WO \text{Attn}(X)=\text{Concat}(\text{head}_1,\ldots,\text{head}_H)W_O Attn(X)=Concat(head1,...,headH)WO
每个 head：
headi=Softmax(QiKi⊤d)Vi \text{head}_i=\text{Softmax}\left(\frac{Q_iK_i^{\top}}{\sqrt{d}}\right)V_i headi=Softmax(d QiKi⊤)Vi
天然并行点：不同 head 之间互不依赖。
因此可以把 head 分配到不同 GPU 上，使每个 GPU 负责一部分 head 的 Q,K,VQ,K,VQ,K,V 投影和注意力计算。最后在输出投影或后续层需要完整向量时做合并。
一个常见要求是：head 数尽量能被 TP 的 GPU 数整除，以平衡负载。
7. 输出层 softmax / cross-entropy 的并行：词表太大时必须优化通信
输出层 logits：
z=hWvocab z = h W_{\text{vocab}} z=hWvocab
如果词表大小为 vvv，logits 张量规模约为 bsvb s vbsv。当 vvv 很大时，直接 AllGather 拼完整 logits 再做 softmax，会产生巨大通信开销。
softmax：
pj=ezj∑k=1vezk p_j = \frac{e^{z_j}}{\sum_{k=1}^{v} e^{z_k}} pj=∑k=1vezkezj
关键在于分母的归一化项：
Z=∑k=1vezk Z = \sum_{k=1}^{v} e^{z_k} Z=k=1∑vezk
如果词表被分片到不同 GPU 上，每个 GPU 先算自己分片的部分和：
Zi=∑k∈Viezk Z_i=\sum_{k \in \mathcal{V}_i} e^{z_k} Zi=k∈Vi∑ezk
再 AllReduce 得到全局：
Z=∑i=1NZi Z = \sum_{i=1}^{N} Z_i Z=i=1∑NZi
这样就不需要 AllGather bsvb s vbsv 级别的大张量，而只需要 AllReduce bsb sbs 规模的归一化量，通信从"随词表线性增长"降到"随 batch*seq 增长"。
8. 流水线并行（PP）：按层切分，降低单卡显存，但会引入空泡
8.1 PP 的基本方式
把网络层按顺序切成多个 stage，每张卡存一段层。例如 8 层、2 卡：
- GPU0：L0L_0L0-L3L_3L3
- GPU1：L4L_4L4-L7L_7L7
前向时激活从 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 等方案更合适