· 分布式训练

大模型分布式训练算法:五种并行策略与代价

数据并行、流水线并行、张量并行、专家并行与 ZeRO——逐个说清楚切的是什么、代价落在哪,以及面试里该怎么说。

分布式训练 · 图文笔记
21:24原片时长
23静止画面还原
10重绘矢量图解
16字节 / 参数
这一章在面试里真正会被追问的只有两件事:每个算法切的是什么、代价在哪,以及你实际会选哪一个。后者的答案很干脆 —— 中小规模模型用 ZeRO 就够,简历上直接写 ZeRO;其余的四种要能讲清原理,但不必死记实现细节。
判断口径一句话:数据并行是为了提速,其余四种是因为一张卡太小、放不下模型。

01 先把五个策略摆到一张图上

参数量与数据规模的增速早已超过单卡显存与算力的增速,分布式训练要同时解决两件事:把计算摊到多设备上以缩短训练时间,以及把多设备的存储拼起来突破单卡容量上限。主流做法只有五种,区别全在切什么 —— 而切法与代价是一一对应的。

图 1 · 五种并行策略各自切分什么:PP 切层、TP 切层内、EP 切专家、DP 切数据、ZeRO 切状态
同一堆 Transformer 层,四个切分维度流水线并行 PP · 沿深度纵切第 4 段(若干层)第 3 段(若干层)第 2 段(若干层)第 1 段(若干层)切在层与层之间显存取自「层」,每卡只放一段;代价是段间串行 → 气泡张量并行 TP · 沿层内横切GPU0GPU1GPU2GPU3同一个权重矩阵被列切开,每卡只持 1/4 的列显存取自「一层的权重」,切得最细;代价是每层都要通信专家并行 EP · 按专家分设备MoE Layer专家 1专家 2专家 3专家 4GPU 0专家 1专家 2All-to-AllGPU 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 完全够。

参数 bf1614 GB
梯度 bf1614 GB
优化器状态 fp3284 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-1RS-2RS-3(每卡 1 块完整)AG-1AG-2AG-3(每卡全量)GPU 01234234344413414321GPU 11234134141142144321GPU 21234124122213214321GPU 31234123233324324321阶段ReduceScatterAllGather浅蓝块=尚未归约的原始数据片;彩色块=已经累加出完整和的块。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)段 0FB段 1FB段 2FB段 3FB每段 8 个时隙只有 2 个在算 → 利用率 25%切成 4 个 micro-batch(m = 4)段 0FFFFBBBB段 1FFFFBBBB段 2FFFFBBBB段 3FFFFBBBB每段 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,前向零通信a11a12a13a14a21a22a23a24A (2×4)×w11w21w31w41v11v21v31v41W1W2=C1C2c11c12c21c22C (2×2)GPU0 持 W 的左半列、GPU1 持右半列,A 完整广播给两卡;各卡算出自己那半列结果,直接拼起来就是 C。代价:一次 AllGather(把 A 广播出去)。② 行并行:按行切 W,反向必须 AllReducea11a12a21a22a13a14a23a24A1 A2×w11w12w21w22v31v32v41v42W1 / W2(上下切)=两个部分积Σc11c12c21c22C = 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 派活输入 tokenbatch 内全部位置Router / Gate为每个 token 选专家GPU 0专家 1FFN专家 2FFN显存只装 2/8 个专家的权重+ 全部 attention / 共享层All-to-Alltoken 派出去 / 结果收回来GPU 1专家 3FFN专家 4FFN显存只装 2/8 个专家的权重+ 全部 attention / 共享层All-to-Alltoken 派出去 / 结果收回来为什么 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 字节/参数020406080100120124816326424 GB 消费卡80 GB A100/H100单卡模型状态显存DP(无分片)= 112 GBZeRO-1 优化器状态分片ZeRO-2 + 梯度分片ZeRO-3 +参数分片N = 8 卡时DP112 GBZeRO-138.5 GBZeRO-226.3 GBZeRO-314.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 就用它DeepSpeedZeRO-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 幅配图为按内容重绘的矢量图,未使用原视频画面或讲义插图。

本笔记为个人学习用途,对视频内容做结构化整理与图解重绘,未使用原视频画面与讲义图片。

术语口径以配套讲义为准;显存数值为按通用字节数口径的估算,实际取决于实现、优化器与并行配置。

分享:
返回文章列表

相关文章

全部文章 »