2347 字
7 分钟
Megatron-LM：张量并行与大规模训练
2019 年，NVIDIA 的 Shoeybi 等人发表了 Megatron-LM，提出了张量模型并行（Tensor Model Parallelism）的系统化方法，使得在多 GPU 上高效训练超大规模 Transformer 成为可能。从 GPT-3（175B）到 PaLM（540B），这些里程碑模型的训练都建立在 Megatron-LM 的并行策略之上。结合 DeepSpeed 的 ZeRO 优化和 FSDP，Megatron-LM 的思想已成为现代分布式训练的基础设施。
没有 Megatron-LM 的并行策略，就没有 GPT-3、PaLM 这些超大规模模型的诞生。
本文要点
- 分布式训练的三大并行策略：数据、模型、流水线
- 张量模型并行：列并行与行并行的数学原理
- MLP 层和 Self-Attention 的张量并行方案
- 流水线并行：层级划分与微批次技术
- ZeRO 优化：三级分片策略
- 3D 并行：张量 + 流水线 + 数据的组合策略
- 混合精度训练：FP16/BF16 与损失缩放
- Megatron-LM 在 GPT-3 和 PaLM 中的应用
- Megatron-DeepSpeed 与 FSDP
一、背景：为什么需要分布式训练？
1.1 模型规模的增长
graph LR
subgraph "模型规模增长"
A["GPT-2<br/>1.5B<br/>2019"]
B["GPT-3<br/>175B<br/>2020"]
C["PaLM<br/>540B<br/>2022"]
D["DeepSeek-V3<br/>671B<br/>2024"]
end
A --> |"117×"| B
B --> |"3×"| C
C --> |"1.2×"| D
style D fill:#4CAF50,color:#fff
一个 175B 参数的模型仅权重就需要约 350GB 显存（FP16），而单张 A100 GPU 只有 80GB。因此，分布式训练不是可选项，而是必需品。
1.2 分布式训练的三大策略
graph TD
A["分布式训练"] --> B["数据并行<br/>Data Parallelism"]
A --> C["模型并行<br/>Model Parallelism"]
A --> D["流水线并行<br/>Pipeline Parallelism"]
B --> B1["相同模型复制到多 GPU<br/>不同数据切片并行训练"]
C --> C1["模型拆分到多 GPU<br/>张量并行（Megatron-LM）"]
D --> D1["模型按层拆分到多 GPU<br/>微批次流水执行"]
style A fill:#1976D2,color:#fff
二、张量模型并行
张量并行是 Megatron-LM 的核心贡献，它将单个张量操作（矩阵乘法）拆分到多个 GPU 上并行计算。
2.1 MLP 层的张量并行
标准的 Transformer MLP 层由两个线性层组成：
MLP: Y = GeLU(XA)B其中 A: [h, 4h], B: [4h, h]列并行（Column Parallelism）
第一个线性层 A 按列拆分到 N 个 GPU：
graph LR
subgraph "列并行（第一个线性层）"
X["X<br/>[b, s, h]"] --> GPU1["GPU 1<br/>A₁<br/>[h, 4h/N]"]
X --> GPU2["GPU 2<br/>A₂<br/>[h, 4h/N]"]
X --> GPUN["GPU N<br/>Aₙ<br/>[h, 4h/N]"]
GPU1 --> Y1["Y₁ = GeLU(XA₁)<br/>[b, s, 4h/N]"]
GPU2 --> Y2["Y₂ = GeLU(XA₂)<br/>[b, s, 4h/N]"]
GPUN --> YN["Yₙ = GeLU(XAₙ)<br/>[b, s, 4h/N]"]
end
style GPU1 fill:#4CAF50,color:#fff
style GPU2 fill:#2196F3,color:#fff
style GPUN fill:#FF9800,color:#fff
每个 GPU 持有 A 的 1/N 列，独立计算自己部分的输出。GeLU 激活函数可以独立应用于每个分片。
行并行（Row Parallelism）
第二个线性层 B 按行拆分：
graph LR
subgraph "行并行（第二个线性层）"
Y1["Y₁<br/>[b, s, 4h/N]"] --> GPU1_B["GPU 1<br/>B₁<br/>[4h/N, h]"]
Y2["Y₂<br/>[b, s, 4h/N]"] --> GPU2_B["GPU 2<br/>B₂<br/>[4h/N, h]"]
YN["Yₙ<br/>[b, s, 4h/N]"] --> GPUN_B["GPU N<br/>Bₙ<br/>[4h/N, h]"]
GPU1_B --> Z1["Z₁ = Y₁B₁<br/>[b, s, h]"]
GPU2_B --> Z2["Z₂ = Y₂B₂<br/>[b, s, h]"]
GPUN_B --> ZN["Zₙ = YₙBₙ<br/>[b, s, h]"]
Z1 --> ALLREDUCE["All-Reduce<br/>Z = Z₁+Z₂+...+Zₙ"]
Z2 --> ALLREDUCE
ZN --> ALLREDUCE
end
style ALLREDUCE fill:#F44336,color:#fff
每个 GPU 独立计算自己分片的矩阵乘法，然后通过 All-Reduce 操作合并结果。
# MLP 张量并行的伪代码def tensor_parallel_mlp(X, A_col_shard, B_row_shard): """ X: [b, s, h] — 在所有 GPU 上相同 A_col_shard: [h, 4h/N] — 每个 GPU 持有 A 的 1/N 列 B_row_shard: [4h/N, h] — 每个 GPU 持有 B 的 1/N 行 """ # 每个 GPU 独立计算（无通信） Y_shard = gelu(X @ A_col_shard) # [b, s, 4h/N] Z_shard = Y_shard @ B_row_shard # [b, s, h]
# All-Reduce 通信（唯一的同步点） Z = all_reduce(Z_shard) # [b, s, h] return Z2.2 Self-Attention 的张量并行
Self-Attention 天然适合并行化：多头注意力中不同的头本身就是独立的计算单元。
graph TD
X["输入 X [b, s, h]"] --> QKV["QKV 投影<br/>[b, s, 3h]"]
QKV --> GPU1["GPU 1<br/>Head 1-N/N<br/>Q₁, K₁, V₁"]
QKV --> GPU2["GPU 2<br/>Head N/N+1-2N/N<br/>Q₂, K₂, V₂"]
QKV --> GPUN["GPU N<br/>Head (N-1)N/N-N<br/>Qₙ, Kₙ, Vₙ"]
GPU1 --> ATT1["Attention(Q₁,K₁,V₁)"]
GPU2 --> ATT2["Attention(Q₂,K₂,V₂)"]
GPUN --> ATTN["Attention(Qₙ,Kₙ,Vₙ)"]
ATT1 --> PROJ1["Output 投影（行并行）"]
ATT2 --> PROJ2["Output 投影（行并行）"]
ATTN --> PROJN["Output 投影（行并行）"]
PROJ1 --> ALLREDUCE["All-Reduce"]
PROJ2 --> ALLREDUCE
PROJN --> ALLREDUCE
style ALLREDUCE fill:#F44336,color:#fff
关键：每个 GPU 处理 N/N 个注意力头，计算完全独立，不需要通信。只有在最后的输出投影后才需要 All-Reduce。
2.3 通信分析
| 组件 | 通信操作 | 通信量 | 通信次数 |
|---|---|---|---|
| Self-Attention | All-Reduce | 2bsh | 每层 1 次 |
| MLP | All-Reduce | 2bsh | 每层 1 次 |
| 总计 | — | 4bsh × L | 每层 2 次 |
其中 b 是批次大小，s 是序列长度，h 是隐藏维度，L 是层数。通信量与计算量相比很小（计算量是 O(bsh²)，通信量是 O(bsh)），因此张量并行的效率很高。
三、流水线并行
当模型太大，无法放入单 GPU 时，需要将模型按层拆分到多个 GPU——这就是流水线并行。
3.1 基本流水线
graph LR
subgraph "4-GPU 流水线并行"
GPU0["GPU 0<br/>Layer 0-5"] --> GPU1["GPU 1<br/>Layer 6-11"]
GPU1 --> GPU2["GPU 2<br/>Layer 12-17"]
GPU2 --> GPU3["GPU 3<br/>Layer 18-23"]
end
style GPU0 fill:#4CAF50,color:#fff
style GPU1 fill:#2196F3,color:#fff
style GPU2 fill:#FF9800,color:#fff
style GPU3 fill:#9C27B0,color:#fff
3.2 气泡问题
简单流水线存在”气泡”（Bubble）：当 GPU 0 处理第一批数据时，GPU 1-3 处于空闲状态。
时间步: 1 2 3 4 5 6 7 8GPU 0: [F0] [F1] [F2] [F3] .. .. [B0] [B1]GPU 1: .. [F0] [F1] [F2] [F3] .. .. [B0]GPU 2: .. .. [F0] [F1] [F2] [F3] .. ..GPU 3: .. .. .. [F0] [F1] [F2] [F3] ..
F = Forward, B = Backward, .. = 气泡（空闲）3.3 微批次技术（Micro-batching）
Megatron-LM 使用微批次技术来减少气泡：
将一个大批次分成 M 个微批次，流水执行：时间步: 1 2 3 4 5 6 7 8GPU 0: [F0] [F1] [F2] [F3] [B0] [B1] [B2] [B3]GPU 1: .. [F0] [F1] [F2] [F3] [B0] [B1] [B2]GPU 2: .. .. [F0] [F1] [F2] [F3] [B0] [B1]GPU 3: .. .. .. [F0] [F1] [F2] [F3] [B0] [B1]
气泡比例 = (P-1) / M，其中 P 是流水线阶段数，M 是微批次数当 M >> P 时，气泡接近 03.4 1F1B 调度
Megatron-LM 采用 1F1B（One Forward One Backward）调度策略，进一步优化内存使用：
1F1B 调度（4 GPU, 4 微批次）：GPU 0: [F0] [F1] [F2] [F3] [B0] [B1] [B2] [B3]GPU 1: .. [F0] [F1] [B0] [F2] [B1] [F3] [B2] [B3]GPU 2: .. .. [F0] [B0] [F1] [B1] [F2] [B2] [F3] [B3]GPU 3: .. .. .. [F0] [B0] [F1] [B1] [F2] [B2] [F3] [B3]
好处：尽早释放中间激活值，减少峰值内存四、数据并行与 ZeRO
4.1 标准 数据并行
标准数据并行将模型复制到每个 GPU，每个 GPU 处理不同的数据分片：
GPU 0: 完整模型 + Batch[0:N/G]GPU 1: 完整模型 + Batch[N/G:2N/G]GPU G-1: 完整模型 + Batch[(G-1)N/G:N]
每步后 All-Reduce 梯度问题：每个 GPU 都要存储完整的模型、梯度和优化器状态。
4.2 ZeRO：零冗余优化器
DeepSpeed 的 ZeRO（Zero Redundancy Optimizer）通过分片消除数据并行中的冗余：
graph TD
subgraph "标准数据并行（每个 GPU）"
DP_All["模型参数 + 梯度 + 优化器状态"]
end
subgraph "ZeRO-1：优化器状态分片"
Z1["模型参数 + 梯度（完整）"]
Z1_S["优化器状态（1/G）"]
end
subgraph "ZeRO-2：+梯度分片"
Z2["模型参数（完整）"]
Z2_S["梯度 + 优化器状态（1/G）"]
end
subgraph "ZeRO-3：+参数分片"
Z3_S["模型参数 + 梯度 + 优化器状态（1/G）"]
end
DP_All --> Z1
Z1 --> Z2
Z2 --> Z3_S
style Z3_S fill:#4CAF50,color:#fff
| ZeRO 阶段 | 分片内容 | 内存节省 | 通信开销 |
|---|---|---|---|
| ZeRO-1 | 优化器状态 | ~4× | 与标准 DP 相同 |
| ZeRO-2 | 优化器状态 + 梯度 | ~8× | 与标准 DP 相同 |
| ZeRO-3 | 优化器状态 + 梯度 + 参数 | ~G×（G 为 GPU 数） | ~1.5× 通信 |
五、3D 并行策略
实际训练超大规模模型时，通常需要组合三种并行策略：
graph TD
A["3D 并行策略"] --> B["张量并行<br/>TP"]
A --> C["流水线并行<br/>PP"]
A --> D["数据并行<br/>DP"]
B --> B1["GPU 组内<br/>同一层的不同分片"]
C --> C1["GPU 组间<br/>不同层的分片"]
D --> D1["GPU 组副本<br/>不同数据分片"]
B1 --> E["示例：GPT-3 175B"]
C1 --> E
D1 --> E
E --> F["TP=8, PP=64<br/>DP=512<br/>总计 1024 GPU"]
style A fill:#1976D2,color:#fff
style F fill:#4CAF50,color:#fff
5.1 GPT-3 的并行配置
| 维度 | 配置 | GPU 数 | 说明 |
|---|---|---|---|
| 张量并行 | TP=8 | 8 | 同一节点内 8 GPU |
| 流水线并行 | PP=64 | 64 | 64 个流水线阶段 |
| 数据并行 | DP=512 | 512 | 512 个数据副本 |
| 总计 | — | ~10,000 | Azure A100 集群 |
5.2 并行策略选择指南
| 模型规模 | 推荐策略 | 说明 |
|---|---|---|
| <1B | 纯数据并行 | 单 GPU 可容纳 |
| 1B-10B | 数据并行 + ZeRO | DeepSpeed ZeRO-3 |
| 10B-100B | 3D 并行 | TP + PP + DP |
| 100B+ | 3D 并行 + ZeRO | Megatron + DeepSpeed |
5.3 通信与计算的权衡
张量并行：通信量大但延迟低（节点内 NVLink）流水线并行：通信量小但延迟高（节点间网络）数据并行：通信量中（All-Reduce 梯度）
最优策略：- TP 用于节点内（NVLink 高带宽）- PP 用于节点间（InfiniBand）- DP 用于扩展到更多节点六、混合精度训练
Megatron-LM 使用混合精度训练来加速计算和减少显存占用。
5.4 FP16/BF16 混合精度
# 混合精度训练流程def mixed_precision_step(model, optimizer, batch): # 1. 前向传播（FP16） with autocast(dtype=torch.float16): loss = model(batch)
# 2. 损失缩放（防止梯度下溢） scaled_loss = loss * scale_factor
# 3. 反向传播（FP16） scaled_loss.backward()
# 4. 梯度缩放回 FP32 更新 optimizer.step() # FP32 主权重更新
# 5. 动态调整缩放因子 if has_inf_or_nan(gradients): scale_factor /= 2 # 减小缩放因子 else: scale_factor *= 2 # 增大缩放因子5.5 BF16 vs FP16
| 特性 | FP16 | BF16 |
|---|---|---|
| 指数位 | 5 位 | 8 位 |
| 尾数位 | 10 位 | 7 位 |
| 数值范围 | 小（需要损失缩放） | 大（与 FP32 相同） |
| 精度 | 高 | 低 |
| 需要 Loss Scaling | 是 | 否 |
| 硬件要求 | Volta+ | A100+ |
BF16 的优势是不需要损失缩放（数值范围与 FP32 相同），简化了训练流程。Ampere 及以上架构的 GPU 支持 BF16。
七、Megatron-LM 在 GPT-3 和 PaLM 中的应用
5.6 GPT-3（175B）的训练配置
| 参数 | 值 |
|---|---|
| 模型参数 | 175B |
| 层数 | 96 |
| 隐藏维度 | 12,288 |
| 注意力头数 | 96 |
| 上下文长度 | 2,048 |
| 训练数据 | 300B Tokens |
| GPU | ~10,000 × V100 |
| 训练成本 | 估计 $4-5M |
5.7 PaLM（540B）的训练配置
| 参数 | 值 |
|---|---|
| 模型参数 | 540B |
| 层数 | 118 |
| 隐藏维度 | 18,432 |
| 注意力头数 | 48 |
| 上下文长度 | 4,096 |
| 训练数据 | 780B Tokens |
| TPU | 6,144 × TPU v4 |
| 训练成本 | 估计 $10-20M |
八、Megatron-DeepSpeed 与 FSDP
5.8 Megatron-DeepSpeed
NVIDIA 和 Microsoft 合作将 Megatron-LM 的张量并行与 DeepSpeed 的 ZeRO 优化结合：
Megatron-DeepSpeed = Megatron 张量并行 + DeepSpeed ZeRO + 流水线并行这是目前训练超大规模模型最成熟的方案之一。
5.9 FSDP（Fully Sharded Data Parallel）
Meta 提出的 FSDP 是 ZeRO-3 的 PyTorch 原生实现：
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
# 使用 FSDP 包装模型model = FSDP( model, sharding_strategy=ShardingStrategy.FULL_SHARD, # 等价于 ZeRO-3 mixed_precision=MixedPrecision( param_dtype=torch.bfloat16, reduce_dtype=torch.bfloat16, ),)FSDP 的优势是与 PyTorch 深度集成，使用更简单。但在超大规模场景下，Megatron-DeepSpeed 仍然更成熟。
常见问题 FAQ
5.1 Q1： 张量并行和流水线并行应该用哪个？
取决于硬件拓扑：
- 节点内（NVLink 高带宽）：使用张量并行
- 节点间（网络带宽低）：使用流水线并行
- 两者可以组合使用
5.2 Q2： ZeRO-3 的通信开销有多大？
ZeRO-3 需要在每次前向/反向传播时进行参数的 All-Gather 通信。相比标准数据并行，通信量增加约 1.5 倍，但内存节省可达 G 倍（G 为 GPU 数）。在 NVLink 连接的节点内，通信开销可以接受。
5.3 Q3： 训练 175B 模型需要多少 GPU？
取决于上下文长度和批次大小：
- 最小配置：约 128 × A100-80G（使用 ZeRO-3 + 梯度检查点）
- 推荐配置：约 256-512 × A100-80G
- GPT-3 实际：约 10,000 × V100
5.4 Q4： Megatron-LM 和 Colossal-AI 有什么区别？
Colossal-AI 是另一个分布式训练框架，支持更多并行策略（如序列并行、自动并行搜索）。Megatron-LM 更成熟稳定，被工业界广泛使用。两者可以互补。
5.5 Q5： 流水线并行的气泡问题能完全解决吗？
理论上当微批次数 M 远大于流水线阶段数 P 时，气泡比例趋近于 0。但 M 不能无限增大（受限于总批次大小和梯度累积）。实际中气泡比例通常在 5-15% 之间。
5.6 Q6： FSDP 会替代 Megatron-LM 吗？
FSDP 简化了使用，但在超大规模（>100B）训练中，Megatron-LM 的手动调优仍然更高效。趋势是两者融合——例如在 FSDP 中集成张量并行策略。
小结
Megatron-LM 的核心贡献可以总结为：
- 张量模型并行：将矩阵乘法按列/行拆分到多 GPU，高效且通信量小
- Self-Attention 并行：利用多头注意力的天然并行性
- 1F1B 流水线调度：减少流水线气泡和峰值内存
- 3D 并行策略：张量 + 流水线 + 数据并行的最优组合
- 混合精度训练：FP16/BF16 加速计算、减少显存
Megatron-LM 使得训练 GPT-3、PaLM 等超大规模模型成为可能。虽然后来出现了 DeepSpeed ZeRO、FSDP 等替代方案，但 Megatron-LM 的张量并行思想仍然是分布式训练的基石。
对于想深入了解的读者，建议阅读顺序：
- 本文（Megatron-LM）→ 理解分布式训练基础
- 第 4 篇（PaLM 与 Scaling Law）→ 理解大规模模型设计
- 第 2 篇（GPT 系列）→ 理解 GPT-3 的训练细节
- 第 53 篇（DeepSeek-V3）→ 理解现代 MoE 训练
参考资料
- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism — Shoeybi et al., 2019
- Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM — Narayanan et al., 2021
- Megatron-LM 第3版：Efficient Training of Large Language Models — Smith et al., 2022
- DeepSpeed: System Optimizations for Training Large Deep Learning Models — Rajbhandari et al., 2019
- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models — Rajbhandari et al., 2020
- PyTorch FSDP
- Megatron-LM GitHub
支持与分享
如果这篇文章对你有帮助，欢迎支持作者或分享给更多人
部分信息可能已经过时
相关文章 智能推荐
1
Chinchilla 与训练token数量的秘密
AI 深度解读 DeepMind Chinchilla 论文(2022)——揭示训练tokens数量与模型大小的关系，重新定义 Scaling Law
2
BPE 与 SentencePiece：LLM 分词器原理
AI 深度解读 LLM 分词器——BPE 算法、WordPiece、Unigram LM、SentencePiece 原理与模型对比
3
BERT 与双向预训练：NLP 预训练模型的崛起
AI 深度解读 Google BERT 论文（2018）——双向 Transformer、掩码语言模型（MLM）、下一句预测（NSP），以及预训练+微调范式。
4
Phi-1 论文解读：高质量数据胜过大规模
AI 深度解读 Phi-1 论文——Microsoft 如何用教科书级别的数据和合成数据，用远小于竞争对手的模型达到优秀的代码生成性能。
5
OpenAI o1/o3 与推理模型的崛起
AI 深度解读 OpenAI o1/o3 系统（2024-2025）——推理模型的新范式、RLVR 训练方法、Test-Time Compute Scaling，以及在 ARC-AGI 上的突破性表现。