从0训练LLM · 现代大模型体系 阅读 + 实验 ≈ 35 分钟 页内实验 ×3
C10 并行:一张卡不够怎么办 第 2 课

每张卡都放不下怎么办?

ZeRO:不再每卡存一份
数据并行时 8 张卡存着 8 份一模一样的优化器状态、梯度和参数。ZeRO 把它们依次切成 8 份,每张卡只存、只更新自己那份:锚点每卡从 444.5 GB 降到 152.8、104.2、55.6 GB。前两级的通信和数据并行一样多,切参数多 50%;激活一点没切。
❓

本课一问:数据并行解决了「慢」,可每张卡仍要放下完整的 444.5 GB。8 张卡上存着 8 份一模一样的东西——能不能每张卡只存一份的 1/8?

这一课先不碰:激活怎么切(第 3 课的序列并行、岔路 1);把状态挪到 CPU 内存或硬盘(只在 Dig 里提一句);层内怎么切矩阵(第 3 课)。

从已有出发

从已有出发 · 训练一步摊到多张卡 已学本课还没学
取一批
前向 ZeRO
loss
反向
梯度同步 DP
优化器更新
每课多切一维:
数字口径:构建时由锚点 Qwen3.8-27B 与 DeepSeek-V3 的官方 config 按公式现算;卡数按 8 张计。本章没有我们自己的多卡实测。

C9 第 1 课的账,这一课原样搬过来:训练时一个参数要 16 字节——16 位权重 2、16 位梯度 2,再加 FP32 主权重、Adam 的 m、v 各 4(ZeRO 论文把后三项合称「优化器状态」,K = 12)。上一课的数据并行让 8 张卡各存一整份。这一课讨论的是「一张卡上的账」,和句子内容无关,所以贯穿例句暂时不出场。

先自己设计:照数据并行的样子,每卡一份

  1. ① 朴素方案
  2. ② 它在哪失败
  3. ③ 改进方案
朴素方案数据并行:每张卡一整份模型状态
权重 2 字节 梯度 2 字节 优化器状态 12 字节 激活 替别的卡存的副本

改进:ZeRO 的三级切分

关键的观察来自上一课:数据并行里,8 张卡用同一份梯度做同一个更新,每张卡都把全部参数更新了一遍——同一件事做了 8 次。ZeRO(Rajbhandari 等,2019)的做法是分工:卡 i 只负责第 i 份参数的更新,于是只需要存第 i 份的优化器状态;更新完,再把新参数发给大家。逐项看每张卡要存多少:

每张卡的模型状态 · ZeRO 论文 §3 的记账(混合精度 + Adam,Ψ = 参数量,N = 数据并行卡数)
+ + → → →
👆 公式里带下划线的每一项都可以点——点哪项，讲哪项。
把 N 调到 64、Ψ 换成 7.5B,三级依次是 120 → 31.4 → 16.6 → 1.9 GB,正是 ZeRO 原文图 1 的四个数——上面的组件里选「ZeRO 原文的 7.5B」就能复现。

ZeRO 原文图 1 用 7.5B 参数、64 张卡给出了 120 / 31.4 / 16.6 / 1.9 GB 四档 ,同一个公式代入锚点就是上面的四个数。它不压缩任何一个数:每个数的位数和 C9 一模一样,省下的全是副本。

🎯

试三件事:① 在改进一侧从「纯数据并行」逐级点到 ZeRO-3,读每卡合计怎么越过 80 GB 线;② 把卡数 N 从 8 拖到 64、1,024:ZeRO-1、ZeRO-2 很快降不动了(被没切的 2Ψ、4Ψ 卡住),只有 ZeRO-3 一直按 1/N 往下走;③ 在 ZeRO-3 下把「激活」切到全量重计算、选择性重计算,看剩下的 24.4 GB 够不够——激活这一段不随 N 变。

通信:前两级白送,切参数多 50%

ZeRO 切得这么狠,通信会不会爆?回到上一课的 ring all-reduce:它本来就是两半,前一半 reduce-scatter 让每张卡拿到一块的全和,后一半 all-gather 把块发给所有人。

数字口径:每块数据是固定种子生成的整数测试输入,不是真实梯度;通信量与时间由公式现算(推演),参数量取自官方 config 与项目导出。

PyTorch 的 FSDP 就是 ZeRO-3 这一级的工程实现 :把模型按层(或若干层)包成单元,算到哪个单元就 all-gather 哪个单元的参数,算完释放,并提前预取下一个单元,让拼参数的通信和上一个单元的计算重叠。所以 FSDP 的显存峰值 ≈ 本卡的分片 + 正在用的那一两个单元的完整参数:锚点最大的一层(全注意力层 + FFN)约 3.72 亿个参数,16 位下 0.7 GB。

⚠

常见误解之一:「ZeRO 是把数压缩了,所以会损失精度。」

ZeRO 不改变任何一个数的格式:FP32 主权重仍是 FP32,16 位权重仍是 16 位。它省的只是「N 张卡各存一份」的冗余,结果与数据并行逐位相同(不计通信中加法顺序带来的舍入)。

常见误解之二:「ZeRO-3 的通信是数据并行的 3 倍。」

是 1.5 倍:数据并行本来就要 2Ψ(reduce-scatter + all-gather),ZeRO-3 多一次参数 all-gather,变成 3Ψ。

常见误解之三:「ZeRO 把激活也切了。」

没有。每张卡处理自己那份样本,激活就是自己那份,和 N 无关。s = 4,096 时,C9 的选择性重计算要 45.6 GB(C9 写作 42.5 GiB),ZeRO-3 之后剩下的 24.4 GB 放不下;全量重计算只要 2.7 GB,但多花约一次前向的算力。

🔍 DeepSeek-V3 为什么只用 ZeRO-1?

DeepSeek-V3 的报告写明它用的是「ZeRO-1 数据并行」,同时有 16 路流水线并行和跨 8 个节点的 64 路专家并行 。报告没有解释原因;能按记账推出来的是:参数已经被流水线(按层)和专家并行(按专家)切得很碎,每张卡剩下的权重和梯度本来就不多,占大头的优化器状态用 ZeRO-1 再切一刀,就够了——而不切参数,就省掉了 ZeRO-3 那 50% 的额外通信。这是按本课公式的推论,不是报告原话。几种切法怎么组合,见实验课。
🔍 还放不下,能不能挪到 CPU 内存?

可以,代价是走 PCIe。ZeRO 的后续工作 ZeRO-Offload(arXiv 2101.06840)把优化器状态和更新挪到 CPU,ZeRO-Infinity(arXiv 2104.07857)再往 NVMe 硬盘挪。DeepSeek-V3 也把一小块东西放在 CPU:参数的指数滑动平均(EMA)存在 CPU 内存里、每步之后异步更新(报告 §3.2.3)。本章的主线不展开 offload:它换的是带宽更低的链路,原理和本课一样——不是少存,而是换个地方存。

证据三问

问本课的回答
① 它解决什么问题?数据并行让每张卡都存一整份模型状态;ZeRO 让 N 张卡分工,每张卡只存、只更新 1/N
② 数学上为什么可能有效?更新是逐参数的:第 i 份参数的更新只需要第 i 份的梯度与优化器状态,所以可以切开 ;每卡显存与通信都是记账算术
③ 实验上为什么最后选了它?ZeRO 原文在 400 张 GPU 上训了超过 100B 参数的模型、吞吐 15 PFLOPs(摘要自报);FSDP 进了 PyTorch;DeepSeek-V3 用 ZeRO-1,Llama 3 的 DP 维就是 FSDP

演进地图:你在这里

并行的演进:每切一维,都是被上一代的某个真数字逼出来的
  1. 起点 一张卡 C9 的结论
    解决了训练一步的全部东西都在一张卡上,不需要任何通信。
    卡在锚点的模型状态 16 字节 × 27,781,427,952 = 444.5 GB,是 80 GB 的 5.6 倍;就算放得下也太慢:DeepSeek-V3 预训练用了 2,664K H800 GPU 小时,一张卡约 304 年。
  2. 第一代 数据并行 + ring all-reduce ring 来自 HPC(2009),2017–2018 进入深度学习
    解决了每卡一份完整模型、算 1/N 的数据,每步把梯度加起来。ring 让每卡只收发 2(N−1)/N 倍梯度:锚点 8 卡 97.2 GB、64 卡 109.4 GB(汇总到 0 号卡的朴素做法:388.9 GB → 3,500 GB)。
    卡在只解决「慢」:每张卡仍要放下完整的 444.5 GB。
  3. 第一代 · 工程 分桶,边反向边同步 PyTorch DDP(Li et al. 2020)
    解决了反向逐层产出梯度,攒满一桶(默认 25 MiB)就先发,通信藏到后面几层的反向里。
    卡在DP 的通信量每步固定:锚点 8 卡每卡 97.2 GB,按 IB 50 GB/s 要 1.94 s;每卡 4,096 个 token 的反向只有约 1.05 s(推演)——token 太少就藏不住。
  4. 第二代 ZeRO-1 / 2 / 3 · FSDP Rajbhandari 2019;PyTorch FSDP 2023 你在这里
    解决了不再每卡存一份:8 卡上 444.5 → 152.8 → 104.2 → 55.6 GB(依次切优化器状态、梯度、参数)。
    卡在切参数后每步多一次 all-gather,通信是 DP 的 1.5 倍;激活不切:s = 4,096 的选择性重计算要 45.6 GB,放不进剩下的 24.4 GB。
  5. 第三代 张量并行(层内切矩阵) Megatron-LM(Shoeybi 2019)
    解决了第一个矩阵按列切、第二个按行切,一个 MLP 只要一次 all-reduce;t = 8 配序列并行,激活 45.6 → 5.7 GB。
    卡在每层前向 2 次、反向 2 次同步:每 4,096 个 token 每卡要发 18.8 GB,每层都要等——只能放在机内。ZeRO 原文实测:跨两台 DGX-2,每卡约 5 TFLOPS,不到峰值 5%。
  6. 第四代 流水线并行:GPipe → 1F1B → interleaved 2018–2021
    解决了按层切,跨机只传段与段之间的激活;切 micro-batch 填满流水线;1F1B 把在途激活从 m 份压到 p 份;interleaved 把气泡除以 v。
    卡在首尾必然空等:p = 8、m = 32 时气泡 (p−1)/m = 21.9%(占总时间 17.9%),interleaved v = 4 也还有 5.5%。
  7. 2023–2024 Zero Bubble · DualPipe Qi 2023;DeepSeek-V3 2024
    解决了反向拆成「对输入」B 与「对权重」W,W 拿去填空泡(F = B = W 时 ZB-H1 的气泡是 1F1B 的 1/3);DualPipe 从两头喂 micro-batch,并把 all-to-all 藏到另一批计算后面。
    卡在DualPipe 要存两份参数;调度更复杂;通信只是被藏住了,量一点没少。
  8. 岔路 · 长序列 序列并行 / 上下文并行 2021–2023:SP、Ring Attention、Ulysses
    解决了128K 的序列按序列切给多张卡;zig-zag 头尾配对,把因果负载 10 : 26 : 42 : 58 配平成每卡一样多。
    卡在每个全注意力层都要在卡之间传 K、V:锚点 s = 131,072、8 卡时每层 0.5 GB;Ulysses 按头切,KV 只有 4 个头。
  9. 岔路 · MoE 专家并行 + all-to-all GShard 2020 → DeepSeek-V3 2024
    解决了DeepSeek-V3 的 654B 路由专家参数分到 64 张卡,每卡每层 4 个;每个 token 最多去 4 个节点。
    卡在跨节点 IB 50 GB/s 只有 NVLink 160 GB/s 的约 1/3,跨节点 EP 的计算与通信约 1 : 1;热门专家所在的卡成为掉队者(→ C11)。
  10. 2021 起 组合:3D / 4D / 5D 并行 Narayanan 2021;Llama 3;DeepSeek-V3
    解决了Llama 3 405B:[TP, CP, PP, DP] = 8 × 1 × 16 × 128,16,384 张 H100,BF16 MFU 41%;DeepSeek-V3:PP 16 × EP 64 × ZeRO-1 DP,不用 TP。
    卡在没有万能配置:每种组合都要先算清每卡显存、每一维的通信走哪条链路、气泡多大(实验课)。
数字口径:锚点与 DeepSeek-V3 的数字由 anchor-model.json / frontier-models.json 的官方 config 在构建时按公式现算(GB = 10⁹ 字节,卡按 80 GB 计);「分桶」一行的时间是推演(有效算力 400 TFLOPS、IB 50 GB/s);5 TFLOPS 出自 ZeRO 原文 §1;GPU 小时、并行配置、带宽出自 DeepSeek-V3 报告 §1、§3.1–3.2(本地 PDF);Llama 3 的配置出自其报告表 4。本章没有我们自己的多卡实测。

一句话

📌

ZeRO 不压缩任何数,只是让 N 张卡不再各存一份一样的优化器状态、梯度和参数:锚点在 8 卡上从 444.5 GB 降到 152.8 → 104.2 → 55.6 GB;前两级通信和数据并行一样,切参数多 50%。

🚦

先停在这里:记住「优化器状态占 12/16,最先切」「前两级通信白送,ZeRO-3 多 50%」「ZeRO 不切激活」三句,就可以进下一课。
再深一层:展开两个 Dig(DeepSeek-V3 为什么只用 ZeRO-1、offload),再去读 ZeRO 原文里切激活、临时缓冲与显存碎片的 ZeRO-R 部分。

下一个缺口

ZeRO-3 之后,每张卡仍在一个人算完整的一层、存自己那份激活——一层的矩阵乘能不能拆开算? → 第 3 课:一层的矩阵乘,能不能拆到几张卡上一起算?

回扣先记下:C9 第 5 课的激活公式在下一课被切成 1/t;「ZeRO 不切激活」这条限制,在第 3 课(序列并行)和岔路 1(长序列)兑现。

你可能会问