分布式训练 · 图文笔记
21:24 原片时长
23 静止画面还原
10 重绘矢量图解
16 字节 / 参数
这一章在面试里真正会被追问的只有两件事:每个算法切的是什么、代价在哪 ,以及你实际会选哪一个 。后者的答案很干脆 —— 中小规模模型用 ZeRO 就够,简历上直接写 ZeRO;其余的四种要能讲清原理,但不必死记实现细节。 判断口径一句话:数据并行是为了提速,其余四种是因为一张卡太小、放不下模型。
01 先把五个策略摆到一张图上参数量与数据规模的增速早已超过单卡显存与算力的增速,分布式训练要同时解决两件事:把计算摊到多设备上以缩短训练时间,以及把多设备的存储拼起来突破单卡容量上限。主流做法只有五种,区别全在切什么 —— 而切法与代价是一一对应的。
图 1 · 五种并行策略各自切分什么:PP 切层、TP 切层内、EP 切专家、DP 切数据、ZeRO 切状态
同一堆 Transformer 层,四个切分维度 流水线并行 PP · 沿深度纵切 第 4 段(若干层) 第 3 段(若干层) 第 2 段(若干层) 第 1 段(若干层) 切在层与层之间 显存取自「层」,每卡只放一段;代价是段间串行 → 气泡 张量并行 TP · 沿层内横切 GPU0 GPU1 GPU2 GPU3 同一个权重矩阵被列切开,每卡只持 1/4 的列 显存取自「一层的权重」,切得最细;代价是每层都要通信 专家并行 EP · 按专家分设备 MoE Layer 专家 1 专家 2 专家 3 专家 4 GPU 0 专家 1 专家 2 All-to-All GPU 1 专家 3 专家 4 专家之间本就无共享参数,切分天然合法;代价是 token 路由通信 数据并行 DP · 复制模型,切分数据 完整模型副本 GPU 0 数据 1 完整模型副本 GPU 1 数据 2 完整模型副本 GPU 2 数据 3 每卡一份完整模型,各跑不同 mini-batch,再同步梯度 四者的目的要分清:只有 DP 是为了提速,PP / TP / EP 都是因为一张卡放不下模型;ZeRO 不切模型,它收拾的是 DP 留下的状态冗余 前四种是切模型 ,彼此正交、可以叠着用:真实预训练里典型组合是机内做张量并行(通信最重,必须走 NVLink)、跨机做流水线并行(段间只传激活值)、最外层再叠数据并行。只有在一层都放不下时才有必要上多维并行 —— 而且这已经属于超大模型的领域。实际工程里,十几 B、二十几 B 的模型基本就到头了,那种规模用 ZeRO 完全够。
参数 bf16 14 GB
梯度 bf16 14 GB
优化器状态 fp32 84 GB
合计单卡 112 GB
按 Adam 混合精度的通用口径:参数 2 字节、梯度 2 字节、优化器状态 12 字节(fp32 参数副本 4 + 一阶矩 4 + 二阶矩 4),合计 16 字节 / 参数 ,7B 模型即 112 GB —— 一张 80 GB 的 A100 也装不下。注意优化器状态独吞 12 字节、占 3/4 ,这个比例决定了后面所有优化的方向:要省显存,先省优化器状态。
02 数据并行 DP:唯一为了提速的那一个每张卡各存一份完整模型,各跑不同的数据子集,再把梯度对齐。它解决吞吐,不解决容量 —— 卡再多,单卡显存占用一点都不会降。
图 2 · 数据并行的一轮迭代:本地求梯度 → 汇总求和再平均 → 各自用同一个梯度更新
数据并行的三步循环 ① 各算各的梯度 GPU k 取 mini-batch k 前向 + 反向 → 本地梯度 g_k ② 梯度全局求和平均 AllReduce:Σg_k 再除以 N 每卡拿到完全相同的 g ③ 各自更新参数 用同一个 g 更新本地副本 N 份副本永远保持一致 通信只发生在第 ② 步,且每轮都是一份完整的梯度张量:通信量 ≈ 2 × 参数量 × 字节数,与卡数 N 无关 → 这就是 DP 能线性扩展的底层原因 为什么不能各更新各的:每张卡只看了一部分数据,loss 不同、梯度就不同,真让它们各自更新,最后会得到 N 份互不相同的权重,哪一份都没有用上全部数据 。所以必须先把梯度汇总求和、再除以卡数,把同一个全局梯度发回每张卡,各副本才会永远一致。
N = 3 时用掉的数据 3 份 mini-batch
耗时 1 份的时间
单卡显存变化 0
通信量 ≈ 2 × 参数量
收益可以用一句话概括:用一批数据的时间,实现了三批数据的更新效果 。这就是数据并行的全部意义 —— 提速。代价是第 ② 步每轮都要搬运一份完整的梯度张量,通信量约 2 × 参数量 × 字节数 ;它与卡数 N 无关 ,这正是数据并行能近乎线性扩展的底层原因,也是必须先解决通信瓶颈的原因。
03 AllReduce:为什么不能「都发给第一张卡」汇总梯度最直白的做法是:所有卡把自己的梯度发给 GPU0,GPU0 加完再广播回去。小规模下完全能用,卡一多就废 —— GPU0 的 IO 吞吐是有上限的,同时接收十几张卡的梯度只能排队,速度直接塌掉。要解决的是「把一张卡的压力均摊到每张卡」,这就是 AllReduce。
图 3 · 环形 AllReduce 的两个阶段:ReduceScatter 让每卡拿到一块完整归约值,AllGather 把块补齐
环形 AllReduce:两步走,通信量只有「每卡发 2 份数据」 阶段 初始 (每卡全量) RS-1 RS-2 RS-3 (每卡 1 块完整) AG-1 AG-2 AG-3 (每卡全量) GPU 0 1 2 3 4 2 3 4 3 4 4 4 1 3 4 1 4 3 2 1 GPU 1 1 2 3 4 1 3 4 1 4 1 1 4 2 1 4 4 3 2 1 GPU 2 1 2 3 4 1 2 4 1 2 2 2 1 3 2 1 4 3 2 1 GPU 3 1 2 3 4 1 2 3 2 3 3 3 2 4 3 2 4 3 2 1 阶段 ReduceScatter AllGather 浅蓝块=尚未归约的原始数据片;彩色块=已经累加出完整和的块。3 步 ReduceScatter 后每卡恰好完整持有 1 块,3 步 AllGather 后每卡集齐 4 块。 关键数字:每卡发送/接收的数据量 = 2 × (N-1)/N × 数据量 ≈ 2 倍数据量,与卡数 N 几乎无关(朴素 AllReduce 是 2(N-1) 倍)。 朴素 AllReduce :每张卡发给其他 N−1 张 —— 单卡通信量 2(N−1) 倍数据量,卡越多越慢。环形 AllReduce :把数据切成 N 块按环错位传递,先做归约再做广播 :ReduceScatter 三步步步累加,让每张卡完整持有其中一块;AllGather 三步把块补齐。单卡通信量 ≈ 2 × (N−1)/N 倍,退化成约 2 倍 ,与 N 基本无关。工程含义:环形对带宽敏感、延迟随步数线性增长,因此集群里环 / 树 / 分层拓扑的选择本身就是大头。 04 流水线并行 PP:横着切,用气泡换显存模型太大放不下,就该切模型了。第一种切法是横向按层切 :前面若干层放一张卡、中间若干层一张卡、最后几层一张卡。前向时每张卡算完自己那段就把激活值跨设备传给下一张卡 ,反向则倒序把梯度跨设备回传 ,整条链路照样能完成一次前后向。
图 4 · 同样 4 段流水线:m=1 时利用率仅 25%,切成 4 个 micro-batch 后升到约 73%
同样 4 段流水线:切不切 micro-batch,利用率差 3 倍 未切 micro-batch(m = 1) 段 0 F B 段 1 F B 段 2 F B 段 3 F B 每段 8 个时隙只有 2 个在算 → 利用率 25% 切成 4 个 micro-batch(m = 4) 段 0 F F F F B B B B 段 1 F F F F B B B B 段 2 F F F F B B B B 段 3 F F F F B B B B 每段 11 个时隙有 8 个在算 → 利用率 ≈ 73% 气泡率与「段数 × micro-batch 数」的关系 利用率 ≈ m / (m + P − 1),其中 P 为流水段数、m 为 micro-batch 数。 P = 4 时:m=4 → 57%…换成 m=8 → 68%…m=16 → 76% —— m 越大气泡越小,但 micro-batch 太小会让每卡矩阵乘的并行度不足,是纯工程取舍。 F = 前向(含发起下一段),B = 反向。气泡的本质是「同一时刻只有部分段有活干」,切 micro-batch 就是用更细的任务填缝 问题出在等待:第一批数据在最后一张卡上算的时候,前面的卡全在空转。改法就是让每张卡交出激活值后立刻去取下一批数据 ,用相邻任务填掉自己的空档 —— 把 batch 切成 micro-batch 的目的正在于此。
流水线利用率 ≈ m / (m + P − 1) (P = 段数,m = micro-batch 数)
P = 4 时:m = 4 约 57%、m = 8 约 68%、m = 16 约 76% —— m 越大越好。 但 micro-batch 太小会让每卡矩阵乘的并行度不足,是纯工程取舍。 反直觉的一点 :micro-batch 不减少总通信量,它只是把原本空转的时间填满;真正买到的是利用率,不是通信。05 张量并行 TP:竖着切,切得最细也最贵第二种切法是纵向按张量切 。当某层里那个线性变换的权重矩阵太大、单卡放不下时,看矩阵乘法的结构:结果的每一列只由激活值和权重的对应那几列 算出来。所以完全可以把权重按列拆到多张卡,激活值分别广播给每张卡,各自算自己那部分结果,最后拼起来。
图 5 · 张量并行的两种切法:列并行前向免通信,行并行必须做一次 AllReduce 求和
张量并行:切权重矩阵的两种方式,只有一种要通信 ① 列并行:按列切 W,前向零通信 a11 a12 a13 a14 a21 a22 a23 a24 A (2×4) × w11 w21 w31 w41 v11 v21 v31 v41 W1 W2 = C1 C2 c11 c12 c21 c22 C (2×2) GPU0 持 W 的左半列、GPU1 持右半列,A 完整广播给两卡; 各卡算出自己那半列结果,直接拼起来就是 C。 代价:一次 AllGather(把 A 广播出去)。 ② 行并行:按行切 W,反向必须 AllReduce a11 a12 a21 a22 a13 a14 a23 a24 A1 A2 × w11 w12 w21 w22 v31 v32 v41 v42 W1 / W2(上下切) = 两个部分积 Σ c11 c12 c21 c22 C = C1 + C2 按行切时每卡只累加了部分和,必须靠 AllReduce 把两个部分积逐元素相加 —— 这是 TP 里通信的由来。 Transformer 里通常成对出现:列并行 + 行并行。 TP 的通信量约 2 × 激活值大小,远小于 DP 的 2 × 参数量 → 所以 TP 只放在单机内 NVLink 上,绝不跨机 列并行 :按列切 W,各卡算出自己那半列结果,拼接即得完整 C,前向不需要通信。行并行 :按行切 W,每卡算出的只是部分和 ,必须 AllReduce 逐元素相加才得到 C —— TP 的通信来源就在这里。Transformer 中两者成对出现(MLP 的上投影列切、下投影行切),通信量约 2 × 激活值大小 ,比 DP 的 2 × 参数量小好几个数量级。 结论:TP 只放在单机 NVLink 内 ,绝不跨机 —— 跨机带宽撑不住每层都同步。 06 专家并行 EP:MoE 的天然配套MoE 把 FFN 换成多个专家,每个 token 只激活少数几个。专家之间没有任何共享参数、也互不依赖 ,所以把不同专家分到不同设备上是天然合法的切分 —— 一张卡里放几个 FFN 就行,不需要额外的聚合运算。
图 6 · 专家并行:Router 决定 token 去哪个专家,跨设备靠 All-to-All 派发与回收
专家并行:把一个 MoE 层里的专家拆到不同设备,用 All-to-All 派活 输入 token batch 内全部位置 Router / Gate 为每个 token 选专家 GPU 0 专家 1 FFN 专家 2 FFN 显存只装 2/8 个专家的权重 + 全部 attention / 共享层 All-to-All token 派出去 / 结果收回来 GPU 1 专家 3 FFN 专家 4 FFN 显存只装 2/8 个专家的权重 + 全部 attention / 共享层 All-to-All token 派出去 / 结果收回来 为什么 EP 划算:专家是彼此独立的 FFN,没有共享参数,切分不破坏数学等价性。 为什么 EP 麻烦:token 要走哪几个专家是动态的,All-to-All 的通信量取决于路由分布,负载不均就会互相等。 EP 通常与 DP 组合:EP 划分专家,DP 复制「非专家部分」,一个 batch 内所有专家都被用到 省钱的地方:每张卡只需存自己分到的专家权重,不必复制整层 FFN。 麻烦的地方:token 去哪些专家是动态路由 的,All-to-All 通信量取决于路由分布;一旦负载不均(热门专家挤在一张卡上),快的卡就得等慢的。 工程上通常把 EP 与 DP 叠加:EP 切专家,DP 复制非专家部分,保证一个 batch 内所有专家都被用到。 07 ZeRO:不切模型,切状态核心思想只有两句 :它是增强的数据并行 ;做法是分片存储、按需加载 。把这两句讲透,ZeRO 就算说清楚了。
回到开头的账:数据并行要求每张卡都完整保存「参数 + 梯度 + 优化器状态」,这份跨设备冗余 就是显存最大的浪费。ZeRO 顺着数据并行的流程走,只改一件事 —— 状态不再每卡一份,而是切成 N 片分散存;某层要做前向或反向时,再从其他设备按需临时拉取 ,用完立刻释放。
图 7 · ZeRO-3 在单层内的执行时序:参数拉取、激活值留存、释放、反向再拉一次
ZeRO-3 单层内的执行时序:一次前向到底发生了几次通信 ① 静态:只存分片 第 k 层参数切成 4 份,分别放在 GPU0–3 上 每卡显存 = 1/4 该层参数 ② 前向:先拉全参数 AllGather 凑齐该层 4 个分片,所有卡拿到完整参数 通信 ① AllGather ③ 前向:算激活值 用完整参数算出该层输出;激活值必须留存,反向要用 激活值不可释放 ④ 释放:只留自己那份 其余 3 份立即丢弃,显存回到分片水平;下一层重复 ②③④ 显存回到低点 ⑤ 反向:再拉一次 反向算梯度要用原始权重 → 再拉一次;梯度只回分片所在的卡 AllGather + ReduceScatter 对比:ZeRO-1 只分片优化器状态,前向反向完全不通信;ZeRO-3 把参数也切了,于是每层要拉两次参数。 省下的显存,精确地换成了层数与通信次数的乘积 —— 这就是「阶段越高越慢」的来源。 走一遍 ZeRO-3 的前向:第 k 层的参数只有 GPU_k 存着它的分片,其余卡都要 AllGather 拉过来;算完这一层的激活值后,非本卡的分片立即释放。这里有个不能省的约束 :参数可以丢,激活值不能丢 ,反向算梯度时还要用它。反向则要再来一遍 —— 因为算梯度同样需要原始权重;但这次梯度不需要广播给所有卡,而是只汇总到持有该分片的那张卡 ,由它负责更新这部分参数。
图 8 · ZeRO 分片后的单卡模型状态显存:瓶颈在优化器状态,ZeRO-1 一步就砍掉大半
7B 模型 · Adam 混合精度(16 字节/参数)单卡模型状态显存 参数 2 + 梯度 2 + 优化器状态 12(fp32 参数副本 4 + 一阶矩 4 + 二阶矩 4)= 16 字节/参数 0 20 40 60 80 100 120 1 2 4 8 16 32 64 24 GB 消费卡 80 GB A100/H100 单卡模型状态显存 DP(无分片)= 112 GB ZeRO-1 优化器状态分片 ZeRO-2 + 梯度分片 ZeRO-3 +参数分片 N = 8 卡时 DP 112 GB ZeRO-1 38.5 GB ZeRO-2 26.3 GB ZeRO-3 14.0 GB 结论:瓶颈是优化器状态(12/16 字节)—— ZeRO-1 一步就砍掉 3/4;ZeRO-3 把参数也分片才能让 7B 单卡跑起来,代价是每层前后向都要拉参数 阶段 分片的对象 7B / N=8 单卡状态显存 代价 ZeRO-1 优化器状态 38.5 GB 通信量与原数据并行基本持平 ZeRO-2 + 梯度 26.3 GB 每次反向多一轮 ReduceScatter ZeRO-3 + 模型参数 14.0 GB 每层前后向都要拉参数,通信明显上升
「阶段越高越好」是错的。 阶段越高,要按需加载的东西就越多,而按需加载是跨设备通信。ZeRO-1 只切优化器状态,前后向根本不额外通信,所以它是性价比最高的那一档;越往上越是用通信换显存。
08 选型:优先 ZeRO-1,其次才往上走原则很直接 —— 能用 ZeRO-1 就用 ZeRO-1 :它切得最少,按需加载的也最少,通信成本最低。ZeRO-1 不够用,再考虑 ZeRO-2、ZeRO-3。
租卡训练:优先 ZeRO-1 租卡按时间 计费,训练快就是省钱。 ZeRO-3 可以用更便宜的小卡,但跨设备传输变多、训练更慢。 单卡价格降下来的钱,会被更长的时间吃回去 —— 总账差不多。 既然总价接近,就选更快的那条路。 公司自有卡:可能必须上 ZeRO-3 卡是已经买好的,只能用这几张,没有「换更强的卡」这个选项。 ZeRO-1/2 放不下时,只能靠 ZeRO-3 把参数也切了。 此时省显存的意义大于通信开销。 同一份代码换 stage 即可,改的是配置不是架构。 图 9 · 五个工具与它们真正覆盖的并行方式
五个工具,覆盖五条路线 分布式训练工具 主要支持的并行方式 什么时候用它 PyTorch DDP 数据并行 最轻量;模型状态单卡放得下时的默认解 PyTorch FSDP 全分片数据并行(≈ ZeRO-3) 原生 PyTorch 生态;不想引 DeepSpeed 就用它 DeepSpeed ZeRO-1/2/3、流水线并行、专家并行 显存吃紧、要 Z3 或跨机 PP;ZeRO 事实标准实现 Megatron-LM 张量并行、流水线并行、专家并行 千亿级预训练;TP×PP×DP 三维并行的参考实现 Accelerate 封装 DDP / FSDP / DeepSpeed 单卡脚本改多卡,改几行 + 一份配置 图 10 · 选型顺序:按显存瓶颈在哪一项,从最省事的一档往下走
选型顺序:从最省事的一档往下走 单卡放得下完整状态 DP / DDP 最省事,吞吐随卡数线性涨 优化器状态是瓶颈 ZeRO-1 / ZeRO-2 分片状态与梯度,通信几乎不加 参数本身单卡放不下 ZeRO-3 / FSDP 参数也分片,用通信换显存 一层都放不下 TP + PP 组合 必须切模型本身,TP 锁在机内 MoE 模型 EP(叠加 DP) 专家天然独立,可分到不同卡 面试常问的落点:先问“显存瓶颈在哪一项”,再给方案;一上来就答 ZeRO-3 会被追问通信代价 取材与还原方式 骨架来源 :本机课程录像第 11 讲配套的 Word 讲义第 6 章《分布式训练算法》, 含 6.1 概述、6.2 数据并行、6.3 流水线并行、6.4 张量并行、6.5 专家并行、6.6 ZeRO 六节。讲授范围 :原片 21 分 24 秒,覆盖 6.1 至 6.6.2 分片策略;期间打开过一次 ZeRO 官方博客核对实现细节, 第 7 章《分布式训练实战》不在本片范围内。画面还原 :以 1 fps 采样切出 23 屏静止画面,用于确定讲授顺序与页面停留分布 (其中 816–1170 秒连续停留在 6.6.1 概述页,该段的结论全部来自口述)。口述补足 :正文中标出的面试口径、ZeRO-3 逐层执行流程、梯度只回分片所在卡、 「激活值不可释放」、租卡与自有卡的选型差异,均来自口述讲解,讲义正文未展开。同音错字校正 :转写稿中的「数据并形 / 体度 / 惨势 / 机货值 / 按虚加载 / 分片存住 / 反应传播」 已按上下文与画面原文校正为「数据并行 / 梯度 / 参数 / 激活值 / 按需加载 / 分片存储 / 反向传播」。补充推算 :112 GB / 38.5 GB / 26.3 GB / 14.0 GB 按「16 字节 / 参数」通用口径计算,口述与讲义未给具体数值; 流水线利用率公式为标准形式,用于解释气泡趋势。图解方式 :全部 10 幅配图为按内容重绘的矢量图,未使用原视频画面或讲义插图。