分布式训练与显存计算¶
本章导读：算法岗面试里，“70B 模型怎么在 8×A100 上训练”“全参微调 7B 要多少显存”“ZeRO-3 和 FSDP 什么关系”几乎是必问。本章把显存拆成可计算的账，把并行策略讲成可选择的方案，并给出两个完整算例。
阅读建议：先背第 1 节的显存公式，再理解第 3 节各并行方式的通信代价，最后用第 6、7 节的算例练习白板推导。
核心要点¶
- 混合精度 Adam 训练时，每个参数约需 16 bytes 模型状态：BF16 参数 2 + BF16 梯度 2 + FP32 master 4 + Adam m 4 + Adam v 4。
- 7B 全参训练仅模型状态约 112GB，加激活轻松超过单卡 80GB；70B 约 1.1TB。
- 数据并行复制模型、切分数据；张量并行切分层内矩阵；流水线并行切分层；专家并行切分 MoE 专家；序列并行切分序列维度。
- ZeRO-1/2/3 分别切分优化器状态、梯度、参数；Stage 3 通信量约为 DP 的 1.5–3 倍，但显存收益最大。
- FSDP 可以理解为 PyTorch 原生的 ZeRO-3，支持分层分片与混合分片。
- 训练算力下限 \(C\approx6ND\)；实际时间还要除以 GPU 峰值、MFU、卡数，并加上通信与故障开销。
- TP 适合节点内高带宽（NVLink），PP 适合跨节点，DP/ZeRO 适合大规模扩展；3D 并行是组合拳。
- 激活重计算用约 33% 的额外前向计算换显存，是长序列/大模型训练的必备手段。
- 面试时先给假设（精度、优化器、并行配置），再给数量级；不要背一个没有条件的数字。
1. 显存构成¶
1.1 四本账¶
训练显存主要由四部分组成：
\[
M_{\text{total}}=M_{\text{param}}+M_{\text{grad}}+M_{\text{optim}}+M_{\text{act}}+M_{\text{temp}}
\]
| 组成 | 内容 | 混合精度 Adam 的每参数字节数 |
|---|---|---|
| 参数 | BF16/FP16 权重 | 2 |
| 梯度 | BF16/FP16 梯度 | 2 |
| 优化器 | FP32 master weight + m + v | 12 |
| 合计（模型状态） | 参数+梯度+优化器 | 16 |
| 激活 | 每层前向保存的中间张量 | 与 batch、序列、层数、隐藏维度相关 |
| 临时 | attention 矩阵、通信 buffer、碎片 | 与实现相关 |
有的框架把梯度也存 FP32，模型状态会到 18–20 bytes/param；面试时说明假设即可。
1.2 快速估算表¶
| 模型 | 参数量 | 模型状态（16 B/param） | 仅权重（BF16） |
|---|---|---|---|
| 1B | \(10^9\) | 16GB | 2GB |
| 7B | \(7\times10^9\) | 112GB | 14GB |
| 13B | \(1.3\times10^{10}\) | 208GB | 26GB |
| 70B | \(7\times10^{10}\) | 1.12TB | 140GB |
| 671B（MoE 总参） | \(6.71\times10^{11}\) | 10.7TB | 1.34TB |
注意：
- MoE 的总参数决定显存，激活参数决定计算；不要用激活参数估显存。
- 推理只需权重 + KV Cache，不需要梯度与优化器；训练和推理的账要分开算。
- 上表未含激活；长序列训练时激活可能超过模型状态。
1.3 激活显存¶
激活显存与以下因素成正比：
\[
M_{\text{act}}\propto L\times B\times S\times d\times \text{每层保存的张量数}
\]
粗略经验：
- 不做重计算：每层每 token 可能保存十余个中间张量，7B、B=1、S=2048 时可达数十 GB。
- 做激活重计算：只存每层输入，反向时重算，显存降到约 1/5–1/10，计算增加约 30%。
- 序列并行可以进一步把激活按序列切分到多卡。
面试表达：不要试图背精确的激活公式，先说“与 L、B、S、d 成正比，重计算能大幅降低”，再给数量级。
2. 训练算力与时间¶
2.1 FLOPs 估算¶
\[
C\approx6ND
\]
- 前向：每个 token 约 \(2N\) FLOPs（每个参数一次乘加）；
- 反向：约 \(4N\)（对输入求梯度 + 对参数求梯度）；
- \(N\) 是参数量，\(D\) 是训练 token 数。
更精确地，注意力部分还有 \(O(L S^2 d)\) 的额外开销；当 \(S\) 很大时不可忽略。
2.2 训练时间¶
\[
T\approx\frac{C}{\text{GPU 数}\times\text{峰值 FLOPS}\times\text{MFU}}
\]
MFU（Model FLOPs Utilization）是实际达到的模型 FLOPs 与硬件峰值的比值：
- 大语言模型训练常见 35%–55%（H100 BF16）；
- 受矩阵规模、通信、重计算、数据加载影响；
- 小模型/短序列的 MFU 通常更低。
2.3 算例：7B、1T token、64×A100¶
- \(C=6\times7\times10^9\times10^{12}=4.2\times10^{22}\) FLOPs。
- A100 BF16 峰值约 312 TFLOPS = \(3.12\times10^{14}\) FLOP/s。
- 64 卡、MFU 40%：有效算力 \(64\times3.12\times10^{14}\times0.4\approx8.0\times10^{15}\) FLOP/s。
- 时间 \(T\approx4.2\times10^{22}/8.0\times10^{15}\approx5.25\times10^6\) 秒 ≈ 61 天。
这说明：大模型训练靠的是数千卡并行。同样任务用 512 张 A100 约 7.6 天（理想线性扩展）。
2.4 算例：70B、15T token、4096×H100¶
- \(C=6\times7\times10^{10}\times1.5\times10^{13}=6.3\times10^{24}\) FLOPs。
- H100 BF16 峰值约 989 TFLOPS = \(9.89\times10^{14}\) FLOP/s。
- 4096 卡、MFU 40%：\(4096\times9.89\times10^{14}\times0.4\approx1.62\times10^{18}\) FLOP/s。
- 时间 \(T\approx6.3\times10^{24}/1.62\times10^{18}\approx3.9\times10^6\) 秒 ≈ 45 天。
实际还要加 checkpoint、故障重启、数据加载、评测，周期通常更长。
3. 并行策略¶
3.1 数据并行（DP）¶
- 每张卡保存完整模型，切分 batch；
- 反向传播后 all-reduce 梯度，保证各卡模型一致；
- 通信量与模型参数量成正比，适合模型能放进单卡的情况；
- PyTorch DDP 是标准实现；梯度 bucketing 可以重叠通信与计算。
优点：实现简单、扩展性好；缺点：显存不省，模型必须单卡可放。
3.2 张量并行（TP）¶
把单层矩阵乘切到多卡：
- 列并行：\(Y=XW\)，\(W\) 按列切，各卡输出部分列，拼接后得到完整 \(Y\)；
- 行并行：\(Y=XW\)，\(W\) 按行切，各卡输入部分行，输出需要 all-reduce 求和；
- Transformer 中通常 attention 的 QKV 用列并行、输出投影用行并行；FFN 第一层列并行、第二层行并行。
通信：每层多次 all-reduce，量大且频繁；需要 NVLink/NVSwitch，跨节点会严重拖慢。
优点：单卡显存显著下降；缺点：通信密集，TP 度一般不超过单节点 GPU 数（如 8）。
3.3 流水线并行（PP）¶
按层切分到多卡，前向逐段传递激活，反向逐段回传梯度：
- 朴素实现有大量空闲（bubble）；
- 1F1B（one forward one backward）调度减少气泡；
- interleaved 1F1B 把模型切成更多虚拟段，进一步降低气泡比例。
气泡比例（1F1B，\(p\) 个 stage、\(m\) 个 microbatch）：
\[
\text{bubble}\approx\frac{p-1}{m+p-1}
\]
优点：跨节点通信少（只传激活），适合大模型；缺点：需要足够多的 microbatch，调度复杂，显存受在途激活影响。
3.4 序列并行（SP）¶
把序列维度切分到多卡，配合 TP 使用：
- LayerNorm/Dropout 等非矩阵操作按序列切分；
- Attention 可用 Ring Attention 在不同卡间传递 KV；
- 降低长序列的激活显存。
3.5 专家并行（EP）¶
MoE 专用：不同专家放在不同卡上，token 通过 all-to-all 路由到对应专家，计算后返回。
- 通信量与 batch、序列、hidden、top-k 成正比；
- 负载不均会导致部分卡等待；
- 需要与 DP/TP/PP 组合，形成 4D 并行。
3.6 3D 并行组合¶
典型配置：
- 节点内：TP=8（NVLink）；
- 节点间：PP（少量跨节点通信）；
- 其余维度：DP/ZeRO（数据并行）；
- MoE 模型再加 EP。
选择原则：
| 约束 | 优先策略 |
|---|---|
| 模型单卡放不下 | TP（节点内）+ PP（跨节点） |
| 显存仍不够 | ZeRO-3/FSDP、offload、重计算 |
| 想扩 batch/吞吐 | DP |
| 长序列 | SP/Ring Attention、重计算 |
| MoE 专家多 | EP |
4. ZeRO 与 FSDP¶
4.1 ZeRO 三个阶段¶
| Stage | 切分对象 | 每卡保留 | 显存收益 | 通信量（相对 DP） |
|---|---|---|---|---|
| 1 | 优化器状态 | 参数、梯度 | 约 4× | ~1× |
| 2 | 优化器状态 + 梯度 | 参数 | 约 8× | ~1× |
| 3 | 优化器状态 + 梯度 + 参数 | 仅当前层参数 | 约 N× | 1.5–3× |
解释：
- Stage 1 把 Adam 的 m/v 和 master weight 按卡切分，参数量和梯度仍复制；
- Stage 2 再把梯度切分，反向时用 reduce-scatter；
- Stage 3 连参数也切分，前向/反向按层 all-gather，计算完再释放。
4.2 通信量直觉¶
- DP 每步 all-reduce 梯度，通信量约 \(2N\)（ring）；
- ZeRO-3 每层前向 all-gather 参数、反向 all-gather 参数 + reduce-scatter 梯度，约 \(3N\)；
- 通信能否被计算掩盖，取决于网络带宽与计算密度；
- 跨节点用 ZeRO-3 时，万兆网可能成为瓶颈，优先 NVLink/InfiniBand。
4.3 FSDP¶
FSDP（Fully Sharded Data Parallel）是 PyTorch 原生的全分片数据并行：
- 把参数、梯度、优化器状态按 flat parameter 分片；
- 前向 all-gather、反向 reduce-scatter；
- 支持
FULL_SHARD、SHARD_GRAD_OP、HYBRID_SHARD、_HYBRID_SHARD_ZERO2； - 与 ZeRO-3 思路一致，但实现与生态不同。
面试表达：FSDP ≈ ZeRO-3 的 PyTorch 实现；ZeRO 是 DeepSpeed 的体系，二者不是严格等价，但可以类比。
4.4 Offload¶
- 把优化器状态/参数换到 CPU 内存，甚至 NVMe；
- 省显存，但 PCIe 带宽成为瓶颈，训练速度显著下降；
- 适合显存极度受限、吞吐不敏感的微调场景；
- ZeRO-Offload/ZeRO-Infinity 是典型实现。
5. 混合精度与重计算¶
5.1 FP16 vs BF16 vs FP8¶
| 格式 | 指数位 | 尾数位 | 动态范围 | 特点 |
|---|---|---|---|---|
| FP16 | 5 | 10 | 较小 | 需 loss scaling，易溢出 |
| BF16 | 8 | 7 | 与 FP32 相同 | 训练稳定，精度略低 |
| FP8（E4M3/E5M2） | 4/5 | 3/2 | 较小 | 需缩放，H100 支持，训练/推理加速 |
混合精度训练通常维护 FP32 master weights，前向/反向用 BF16，更新时用 FP32。
5.2 激活重计算¶
- 前向只保存每层输入（或部分关键张量），反向时重新计算中间激活；
- 显存从 \(O(L)\) 降到 \(O(1)\)/层（配合分块可到 \(O(\sqrt{L})\)）；
- 计算开销约增加一次前向，即总 FLOPs 增加约 33%（前向:反向 ≈ 1:2）；
- 可选择只重计算 attention 或 FFN，平衡显存与速度。
6. 算例：70B 全参训练怎么配¶
6.1 假设¶
- 70B 稠密模型，BF16 混合精度，Adam；
- 64 张 80GB H100（8 节点 × 8 卡）；
- 序列长度 4096，microbatch 1，梯度累积；
- 使用 TP=8、PP=4、DP=2（共 64 卡）。
6.2 显存账¶
- 模型状态：\(70\times10^9\times16=1.12\)TB；
- 若用 ZeRO-1 只切优化器：参数+梯度仍需 \(280\)GB，单卡放不下 → 需要 TP/PP；
- TP=8 把单层参数切 8 份，PP=4 把层切 4 段：单卡权重约 \(70\text{B}/32\times2\approx4.4\)GB；
- 模型状态（参数+梯度+优化器）若按 64 个 TP/PP/DP 分片计，单卡约 \(1.12\text{TB}/64\approx17.5\)GB（其中权重约 4.4GB）；
- 加激活、临时 buffer、通信 bucket，单卡约 40–60GB，可放入 80GB。
6.3 通信与效率¶
- TP=8 在节点内，走 NVLink；
- PP=4 跨节点，只传激活，通信量相对小；
- DP=2 的梯度 all-reduce 跨节点，需 InfiniBand；
- 气泡比例约 \((4-1)/(m+3)\)，需要足够多 microbatch；
- 实际 MFU 可能 35%–45%，低于纯 DP 的小模型。
6.4 如果只有 8×A100 80GB¶
- 70B 全参训练基本不可行：模型状态 1.12TB，8 卡每卡 140GB，已超过 80GB；
- 只能做 LoRA/QLoRA（冻结基座，优化器状态大幅减少），或用 ZeRO-3 + CPU offload（极慢）；
- 推理可以：70B BF16 需 140GB，8×80GB 可放下；INT4 量化后单机 2–4 卡可跑。
面试时如果被问“8 卡训 70B”，先反问“全参还是 LoRA、什么精度、多长序列”，再给结论。直接说“能”或“不能”都不严谨。
7. 算例：8×A100 做 7B GRPO¶
7.1 显存¶
- 7B 全参模型状态 112GB，8 卡平均 14GB；
- GRPO 需要 policy、reference、reward（若用 RM），reference 可量化/共享；
- rollout 用 vLLM 会额外占 KV Cache；
- 训练侧用 LoRA + ZeRO-3 + 重计算，单卡可控制在 60GB 内。
7.2 吞吐¶
- rollout 是主要瓶颈：生成长度越长、组越大，采样越慢；
- 训练与推理分时复用（colocate）省显存，但会互相等待；
- 训练/推理分池（disaggregate）吞吐高，但需要更多卡；
- 用 prefix caching、动态 batch、长度控制提升 rollout 吞吐。
7.3 常见配置¶
# 示意配置，不是某个框架的完整配置
train:
precision: bf16
lora: {r: 16, alpha: 32, target: [q_proj, k_proj, v_proj, o_proj]}
zero_stage: 3
gradient_checkpointing: true
micro_batch_size: 1
gradient_accumulation: 8
rollout:
engine: vllm
n_samples: 8
max_new_tokens: 2048
temperature: 0.8
8. 通信与故障¶
8.1 通信原语¶
- All-reduce：求和后广播，DP 梯度同步；
- All-gather：收集分片，ZeRO-3 前向取参数；
- Reduce-scatter：求和后切分，ZeRO-3 反向回传梯度；
- All-to-all：MoE 路由、专家并行；
- Point-to-point：PP 的激活/梯度传递。
8.2 网络拓扑¶
- 节点内：NVLink/NVSwitch，带宽数百 GB/s；
- 节点间：InfiniBand/RoCE，带宽数十到数百 GB/s；
- 跨节点通信要尽量少、尽量重叠；
- TP 放节点内，PP/DP 放节点间，是常见的拓扑感知策略。
8.3 容错¶
- 定期 checkpoint，保存模型、优化器、数据加载器状态；
- 异步 checkpoint 减少停顿；
- 故障检测与自动重启，记录失败 rank 与 step；
- 大规模训练的有效利用率往往受故障影响，MFU 之外还要看 goodput。
9. 面试高频问题¶
Q1：7B 全参微调需要多少显存？¶
- 模型状态 16 B/param ≈ 112GB；加激活可能 150GB+。
- 单卡 80GB 放不下，需要 ZeRO-3/FSDP、offload 或多卡。
- 若用 LoRA，优化器状态只覆盖低秩矩阵，显存大幅下降。
Q2：ZeRO-1/2/3 分别切什么？通信量如何？¶
- 1 切优化器，2 加梯度，3 加参数。
- 通信量：1/2 约与 DP 同量级，3 约 1.5–3 倍。
- Stage 3 显存最省，但 all-gather 频繁，跨节点要谨慎。
Q3：TP 和 PP 怎么选？¶
- TP 通信密集，放节点内；PP 通信少，放跨节点。
- TP 度通常 ≤ 单节点 GPU 数；PP 需要足够 microbatch 降气泡。
- 大模型常用 TP+PP+DP 组合。
Q4：为什么 ZeRO-3 通信量大？¶
- 每层前向要 all-gather 参数，反向再 all-gather 并 reduce-scatter 梯度。
- 参数被反复取用，通信无法完全被计算掩盖。
- 可用分层分片、通信重叠、更大 batch 提升计算密度。
Q5：FSDP 和 ZeRO-3 的区别？¶
- 思路相同：全分片参数/梯度/优化器。
- FSDP 是 PyTorch 原生，ZeRO 是 DeepSpeed 体系；API、分片粒度、生态不同。
- 面试答“可类比”即可，不必强行说完全等价。
Q6：激活重计算为什么能省显存？代价是什么？¶
- 只存层输入，反向重算中间激活，显存大幅下降。
- 代价是额外一次前向，FLOPs 增加约 33%，训练变慢。
- 可选择部分重计算，平衡显存与速度。
Q7：MFU 是什么？怎么提升？¶
- 实际模型 FLOPs / 硬件峰值 FLOPs。
- 提升：增大矩阵、减少通信、重叠计算通信、优化 attention kernel、减少重计算、提高 batch。
- 常见 35%–55%，小模型/长序列偏低。
Q8：70B 模型 8×A100 能全参训练吗？¶
- 不能。模型状态 1.12TB，8×80GB 仅 640GB。
- 只能 LoRA/QLoRA、ZeRO-3+offload（极慢），或换更多卡。
- 推理可以，量化后更宽裕。
Q9：训练 loss spike 与并行配置有关吗？¶
- 有。TP/PP 的数值差异、梯度累积、混合精度溢出、某个 rank 数据异常都可能引发。
- 排查要看单 rank、单 step、梯度范数、数据分布。
- 处理：回滚、跳 batch、降 LR、BF16、梯度裁剪。
Q10：如何估算训练成本？¶
- \(C=6ND\)，除以有效算力得 GPU 小时，再乘单价。
- 加通信、重计算、故障、调参开销，通常乘 1.3–1.5。
- 给出假设与区间，不要给一个“精确”数字。
本文档为面试复习资料，公式与数字建议结合原始论文与硬件规格复核。