第十章分布式训练¶

导读¶

预训练前沿模型必须用多 GPU，甚至多机多卡。本章解释三种并行策略（数据并行 / 张量并行 / 流水并行）、ZeRO/FSDP 显存优化、3D 并行的工程组合，以及生产训练框架（Megatron-LM、DeepSpeed、FSDP）。

10.1 为什么需要分布式¶

单卡瓶颈：

显存：H100 80GB，70B 模型 fp16 = 140GB 装不下；
算力：单 H100 ~670 TFLOPS bf16，训 70B+1T tokens 要 14 EFLOPS = 21000 H100·天 → 单卡要 57 年；
数据吞吐：单卡 SSD 读 ~3GB/s，训练数据 IO 跟不上。

所以**必须并行**——但具体怎么并行，决定了训练效率能否达到 MFU 40-55%。

10.2 三种并行策略¶

三种并行策略对比

数据并行（Data Parallelism, DP）¶

每个 GPU 完整复制模型；
数据切分到各 GPU；
前向各算各的；
反向后 all-reduce 同步梯度；
简单，但每卡装完整模型 → 显存受限。

张量并行（Tensor Parallelism, TP）¶

权重按维度切到多 GPU；
每层前向各算一部分；
用 all-reduce 合并结果；
通信频繁 → 需 NVLink 等高带宽机内互联；
突破单卡显存上限。

流水并行（Pipeline Parallelism, PP）¶

模型按层切到多 GPU（如 80 层切 4 段）；
数据像流水一样流过各段；
激活在段间点对点传输；
通信量小 → 适合跨机；
有 bubble（气泡）浪费。

10.3 ZeRO：显存切分优化¶

ZeRO三阶段

ZeRO（Zero Redundancy Optimizer, Microsoft 2020）观察到一个反直觉事实：

训练显存主要被三块占用——参数、梯度、优化器状态。DP 默认情况下，这三块在每卡上完整复制——这是巨大冗余。

ZeRO 分三阶段切分：

阶段	切分对象	显存节省
ZeRO-1	优化器状态	4×
ZeRO-2	+ 梯度	8×
ZeRO-3	+ 参数	与 GPU 数线性

FSDP（Fully Sharded Data Parallel） 是 PyTorch 原生版 ZeRO-3，已是当前主流训练默认选择。

10.4 3D 并行¶

生产训练 = DP × TP × PP 组合：

                    ┌── PP stage 1 (4 GPUs, TP=4) ──┐
DP group 1 (8 nodes)├── PP stage 2 (4 GPUs, TP=4) ──┤
                    ├── PP stage 3 (4 GPUs, TP=4) ──┤
                    └── PP stage 4 (4 GPUs, TP=4) ──┘
DP group 2 (8 nodes) ... (重复 N 次)

典型组合：

模型规模	DP	TP	PP	总 GPU
7B	8	1	1	8
70B	64	8	1	512
405B	256	8	8	16384

经验法则：

TP 在机内（NVLink 带宽高）；
PP 跨机或跨机内（带宽要求低）；
DP 跨机（最易扩展）。

10.5 通信原语¶

理解分布式训练必须熟悉几个集合通信原语：

原语	说明	用途
broadcast	一卡数据发给所有卡	参数同步
all-reduce	所有卡求和后发给所有卡	DP 梯度同步
all-gather	拼接所有卡数据发给所有卡	ZeRO-3 取参数
reduce-scatter	求和 + 切分到各卡	ZeRO-2 梯度
send/recv	点对点传输	PP 激活传递

底层库：NCCL（NVIDIA）、RCCL（AMD）、MPI（CPU）。

**通信成本**与算力的比例决定了并行效率——如果通信占比 > 30%，MFU 会显著下降。

10.6 流水并行与 1F1B¶

朴素流水并行（GPipe）：所有 micro-batch 前向，再反向 → bubble 大。

1F1B（One Forward One Backward）：

调度：1 个前向紧跟 1 个反向；
反向立即释放激活，节省显存；
bubble 减小但仍存在。

Interleaved 1F1B（Megatron-LM）：

每个 GPU 负责不连续的多段（chunk）；
进一步压缩 bubble；
工程复杂，但训练前沿模型必备。

10.7 训练框架¶

框架	维护方	特点
Megatron-LM	NVIDIA	工业级 3D 并行，前沿 lab 用
DeepSpeed	Microsoft	ZeRO 系列，易用
FSDP	PyTorch 原生	ZeRO-3 等价，主流默认
Tor Titan	Meta	分布式训练栈
Colossal-AI	HPC-AI Tech	中文社区友好
Ray Train	Anyscale	与 Ray 集群集成

工程师选型建议：

中小规模（≤70B）：PyTorch FSDP 足够；
大规模（70B+）：Megatron-LM 或 DeepSpeed；
中文生态：Colossal-AI。

10.8 训练稳定性¶

前沿模型训练动辄几十万 step，稳定性是工程命门：

loss spike：突然飙升。原因：数据异常、数值精度、优化器状态异常；
NaN / Inf：bf16 比 fp16 稳；grad clip 必开；
通信死锁：NCCL hang，常见于 PP + DP 复杂拓扑；
硬件故障：GPU 损坏率不可忽略，需 checkpoint 频繁保存。

实践 recipe：

bf16 + grad_clip=1.0 + warmup 2000 step；
每 30 分钟存 ckpt；
训练前做小型 sanity check；
监控 loss / grad_norm / tokens/s / MFU；
红队：人为注入 NaN、OOM 测试恢复流程。

10.9 Checkpoint 管理¶

70B 模型 fp32 优化器状态 = 280GB，每 30 分钟存一次，一次存数分钟——挑战：

分布式存储：高吞吐（GB/s 级）；
异步保存：避免阻塞训练；
校验和：防止静默损坏；
版本管理：超参 + 数据版本 + ckpt 三元组；
恢复速度：从 ckpt 重启时间影响有效利用率。

10.10 训练成本估算¶

粗算公式：

\[\text{GPU·时} \approx \frac{6 \times N_{\text{params}} \times N_{\text{tokens}}}{\text{FLOPS per GPU} \times \text{MFU}}\]

例：Llama-3 70B + 15T tokens + H100（990 TFLOPS bf16）+ MFU 0.45：

\[\frac{6 \times 70 \times 10^9 \times 15 \times 10^{12}}{990 \times 10^{12} \times 0.45} \approx 1.4 \times 10^7 \text{ GPU·时} \approx 1600 \text{ GPU·年}\]

按 $3/GPU·时计：$42M——Llama-3 70B 公开估算约 $30-50M，吻合量级。

工程实战要点¶

FSDP 是默认：除非训练 ≥70B，不需要 Megatron 复杂度；
TP 在机内、PP 跨机：根据带宽匹配；
bf16 + grad_clip：基本稳定性保障；
通信占比 < 30%：超过则 MFU 显著下降；
checkpoint 30 分钟一次：恢复时间 < 5 分钟；
MFU 40-55% 算优秀：低于 30% 说明配置有问题；
故障是常态：要有完整恢复流程；
成本估算是基础能力：能在事前判断是否可行。

小结¶

三种并行：DP（数据切分）、TP（权重切分）、PP（层切分）；
ZeRO/FSDP 切分参数/梯度/优化器状态，突破显存上限；
生产 = DP × TP × PP 3D 并行；
TP 在机内（NVLink），PP 跨机，DP 跨机扩展；
1F1B / Interleaved 减小流水 bubble；
训练框架：FSDP（默认）、Megatron（前沿）、DeepSpeed（中间）；
稳定性是命门：bf16 + grad_clip + 频繁 ckpt + 完整监控；
成本估算公式：GPU·时 ≈ 6·N·tokens / (FLOPS·MFU)。

练习题¶

三种并行策略各切什么？通信频率与带宽要求如何？
解释 ZeRO-1/2/3 分别切分什么，节省多少显存。
FSDP 与 ZeRO-3 的关系是什么？
一个 70B 模型用 8 张 H100 训练，应该用什么并行组合？
为什么 TP 通常在机内、PP 跨机？
1F1B 调度相对 GPipe 的改进是什么？
估算训 7B 模型 + 1T tokens 需要多少 H100·时（假设 MFU 0.45）。
训练时遇到 loss spike 应该怎么处理？

下一章：第十一章向量数据库与RAG

第十章 分布式训练¶

导读¶