跳到主要内容

07 · 算力:同一台机器,怎么跑出三倍成绩

1. 这一章讲什么​

模型定了、数据好了,轮到训练本身:一批数据放不进显存怎么办?一块卡太慢怎么办?原书第 7 章给三招——混合精度训练、多 GPU(显卡)并行、梯度累积;并把梯度累积单独再讲了一遍(原书 9.2 与 7.3 内容高度重复,本拆解合并到这里)1。

三招的共同点:模型和算法(计算 steps 的固定套路)一行不改,只改「数怎么存、活怎么分」。这也是它排在微调(05)和数据(06)之后的原因——先把能省的省了,再谈买卡。

2. 顶层全景:三招各管一件事​

 显存不够  → 梯度累积:小批次攒梯度,凑成大批次的效果
 显存不够  → 混合精度:权重用 16 位存,显存直接省近半
 速度太慢  → 数据并行:数据切开,多卡同时算
 模型太大  → 模型并行:模型切开,一层层接力
                │
                ▼
        三招可以叠着用(实战框架里常是默认一起开)

3. 核心原理​

3.1 混合精度:16 位干活,32 位守账​

神经网络里的数默认是 32 位浮点(FP32);混合精度训练让不需要高精度的计算用 16 位(FP16)——显存省一半、算力翻倍——但在梯度累积、参数更新这类怕出错的关键步骤保留 FP32。顺带把两个方向词钉死:前向传播(数据从输入流向输出、算出预测的那趟)用低精度,反向(把错误往回传、算梯度的那趟)的账本用 FP322。

FP16 有个天生的坑:它的数值范围小,梯度里那些极小的数会直接「沉底」变成 0(下溢),小到训练等于没训。解法是梯度缩放:把损失先乘一个大数,梯度跟着整体上浮,安全区里算完再缩回去。原书用的是 PyTorch 的自动混合精度(AMP):autocast() 自动决定每个算子用哪种精度,GradScaler 负责缩放损失、反向传播后 scaler.update() 再动态调整3。

补充(不在书里,依据我们的 frontier 书架): 手写混合精度的样板是「autocast 包 forward、scaler 缩放 loss、溢出时跳过 step」;绕过这套直接 loss.backward(),fp16 不走 scaler 时「梯度可能下溢」,loss 不除累积步数时「梯度偏大 N 倍」——两个坑 Accelerate 的拆解里都点了名。 依据: shelf=ai-frontier-reference/accelerate#04-mixed-precision-and-accumulation.md 事实=「fp16 没走 scaler(梯度可能下溢)、loss 没除累积步数(梯度偏大 N 倍)」。

3.2 两种并行:切数据,还是切模型​

多卡并行有两个方向,原书用餐厅打了个整套的比方4:

方案切什么怎么干适用
数据并行(DP)切数据每张卡放完整模型,各算各的一小份,梯度汇总回主卡更新5数据多、模型单卡装得下
模型并行(MP)切模型模型按层拆开,数据依次流过各卡6模型单卡装不下

餐厅版:数据并行是「一大盘菜分给多个服务员,各管一桌,最后汇总反馈」;模型并行是「一道大菜分到切菜区、炒菜区、蒸煮区,按顺序接力」7。原书总结得很准:前者「单独完成」但每卡都要装完整模型,后者「分步完成」但卡间要传数据、通信开销大8。

分布式训练(把一次训练摊到多张卡、多台机器上)的标准姿势是 DDP(DistributedDataParallel):每个进程一张卡,init_process_group 建通信组(world_size=总进程数),DistributedSampler 保证各卡数据不重不漏,DDP 自动同步梯度9。注意别把 DP 和 DDP 混为一谈——前者是老接口(主卡汇总,负载不均),后者是分布式正解。

补充(不在书里,依据我们的 frontier 书架): DDP 快的秘密是把「梯度同步」藏进了反向传播本身:每个参数的梯度一算好就拷进「梯度桶」,桶一满立刻发起跨卡求平均,通信和剩下的反向计算同时进行——不是等全部算完才傻等。 依据: shelf=ai-frontier-reference/pytorch#05-distributed-ddp.md 事实=「每张卡拿不同数据、跑同一个模型,反向后把梯度跨卡求平均」;「梯度桶(聚小成大)+autograd 钩子(边算边传)」。

3.3 梯度累积:小份攒出大份的效果​

训练不是一次吞下所有数据,而是切成小份——每一份就叫「批」(batch);大批(每份切大一点)训练更新更稳,但显存装不下。梯度累积的办法:算完一小份的梯度先不更新,攒够 N 份再统一更新一次。

等效批量(一次更新相当于喂了多少条)= 单份 × N,显存却只按单份算10。

主走查:四个小批次怎么攒成一次更新。(原书 9.2 的设置:单批 16,累积 4 步,等效批量 64。)

 批次 1: forward → loss₁/4 → backward(梯度进账,不更新)
 批次 2: forward → loss₂/4 → backward(继续攒)
 批次 3: forward → loss₃/4 → backward(继续攒)
 批次 4: forward → loss₄/4 → backward
         此时 (i+1) % 4 == 0 成立 → optimizer.step()  ← 四份梯度合成一步更新
         → optimizer.zero_grad()                      ← 清零,开始下一轮攒

走查说明:两个细节不能省。损失除以 4(accumulation_steps)——否则攒出来的梯度是单批的 4 倍,等效学习率跟着翻了 4 倍;更新后必须清零梯度——否则上一轮的旧账混进下一轮。原书明确写了「若单批量大小为 16,则实际模拟 64 的批量大小」11。

原书还给了调参的两头:累积步数加大会更稳但更新变慢;设成 1(不累积)则可能训练不稳12。比喻是「攒够再出手」的购物——小篮子分几趟买,账本记着,凑够一起结账13。

3.4 三招怎么叠​

实战里三招不是三选一:混合精度管「每个数省一半」、梯度累积管「批次凑大」、DDP 管「多卡摊数据」,叠起来才是一轮大模型训练的标配。原书把它们分散在三节讲,每次只叠两招(比如 7.1 的 AMP 演示是单卡,7.2 的 DDP 演示是全精度)——想看三招全叠的成品,11 章的综合案例最接近。

4. 作者的判断与证据​

  • 有演示数据的: DDP 两个 Rank 各自的损失输出(7.2.1)、DP/MP 的每轮平均损失(7.2.2/7.2.3)、梯度累积的每步损失序列(7.3)。
  • 是作者判断的: 「大批量训练能更稳定地更新模型」是经验判断,书里没给对照;「Model Parallel 对计算速度的提升不如 Data Parallel 明显」也是定性比较——根因(卡间通信)书里说了,但没测。
  • 要分开的: 餐厅比喻、购物比喻、搬运工比喻(9.2 那版)都是原书自己的教学装置,别当机制。

5. 边界与局限​

  • 原书的模型并行是最朴素的手动分层:前半层放 cuda:0、后半层放 cuda:1。现代大模型训练的并行方案复杂得多(流水线并行、张量(多维数表)并行、ZeRO 切优化器状态),原书一概没讲。补充(不在书里,依据我们的 frontier 书架):DeepSpeed 的 ZeRO 把「切什么」做成了三档开关(优化器状态/梯度/参数分片),是今天更常用的路子。
  • 混合精度只讲了 FP16 一支;BF16(指数位更多、更稳,不需要损失缩放)没有出现。
  • DDP 演示用 gloo 后端;GPU 上生产环境一般用 NCCL——原书 11.2 提了一句,7 章正文没提。
  • 7.3 与 9.2 重复(搬运工比喻 vs 购物比喻讲的是同一件事),原书未作整合,读书时按一章看即可。

6. 可带走的​

  1. 显存爆了先想梯度累积(等效批量=单批×N),再想混合精度,最后才是换卡。
  2. 梯度累积两个不能省的细节:损失除以累积步数、更新后清零梯度。
  3. FP16 会下溢,GradScaler 是必备件,不是可选项。
  4. 数据并行切数据(每卡完整模型),模型并行切模型(数据接力);模型装得下就别用后者。
  5. DDP 不是 DP:前者每卡一个进程+梯度桶高效同步,后者主卡汇总的老接口。
  6. 三招可叠加;真实训练框架里它们常是配置项而不是手写代码。
  7. 「等效批量 64」这种数字要学会换算:它决定学习率的合理范围。

7. 原文地图​

主题原书章原文位置
章导语第7章text/42-ch07.txt:24(搜「混合精度训练和分布式训练」)
混合精度定义7.1text/43-ch07-01-7-1.txt:24(搜「半精度浮点（FP16）」)
GradScaler 与 autocast7.1text/43-ch07-01-7-1.txt:45(搜「GradScaler」) · text/43-ch07-01-7-1.txt:48(搜「避免数值下溢的风险」)
DP/MP 定义7.2text/44-ch07-02-7-2.txt:24(搜「将数据划分到多个GPU上」)
DataParallel / Model Parallel7.2text/44-ch07-02-7-2.txt:97(搜「torch.nn.DataParallel接口」) · text/44-ch07-02-7-2.txt:149(搜「划分到多个GPU上分别处理」)
餐厅比喻7.2text/44-ch07-02-7-2.txt:204(搜「分菜给多个服务员」) · text/44-ch07-02-7-2.txt:216(搜「分工合作的大菜」)
DP vs MP 总结7.2text/44-ch07-02-7-2.txt:228(搜「单独完成」)
梯度累积原理与损失缩放7.3text/45-ch07-03-7-3.txt:33(搜「累积后进行一次参数更新」) · text/45-ch07-03-7-3.txt:51(搜「accumulation_steps缩放」)
购物比喻7.3text/45-ch07-03-7-3.txt:79(搜「攒够再出手」)
等效批量 16→649.2text/57-ch09-02-9-2.txt:48(搜「模拟64的批量大小」)
累积步数两头9.2text/57-ch09-02-9-2.txt:252(搜「更新变得缓慢」)

Footnotes​

  1. 出处:「第7章 模型性能优化：混合精度训练与分布式训练」第 24 段(text/42-ch07.txt:24,搜「混合精度训练和分布式训练」);重复处见「9.2 基于梯度累积的优化技巧」第 24 段(text/57-ch09-02-9-2.txt:24,搜「内存限制」)。 ↩

  2. 出处:「7.1 混合精度训练的实现」第 24 段(text/43-ch07-01-7-1.txt:24,搜「半精度浮点（FP16）」)。 ↩

  3. 出处:「7.1 混合精度训练的实现」第 46-52 段(text/43-ch07-01-7-1.txt:45,搜「GradScaler」;text/43-ch07-01-7-1.txt:48,搜「避免数值下溢的风险」;text/43-ch07-01-7-1.txt:51,搜「自动调整梯度」)。 ↩

  4. 出处:「7.2 多GPU并行与分布式训练的实现」第 201 段(text/44-ch07-02-7-2.txt:201,搜「分菜」)。 ↩

  5. 出处:「7.2 多GPU并行与分布式训练的实现」第 97 段(text/44-ch07-02-7-2.txt:97,搜「torch.nn.DataParallel接口」)。 ↩

  6. 出处:「7.2 多GPU并行与分布式训练的实现」第 149 段(text/44-ch07-02-7-2.txt:149,搜「划分到多个GPU上分别处理」)。 ↩

  7. 出处:「7.2 多GPU并行与分布式训练的实现」第 204-219 段(text/44-ch07-02-7-2.txt:204,搜「分菜给多个服务员」;text/44-ch07-02-7-2.txt:216,搜「分工合作的大菜」)。 ↩

  8. 出处:「7.2 多GPU并行与分布式训练的实现」第 228 段(text/44-ch07-02-7-2.txt:228,搜「单独完成」)。 ↩

  9. 出处:「7.2 多GPU并行与分布式训练的实现」第 33 与 51-57 段(text/44-ch07-02-7-2.txt:33,搜「torch.distributed包」;text/44-ch07-02-7-2.txt:51,搜「world_size」;text/44-ch07-02-7-2.txt:54,搜「DistributedSampler」)。 ↩

  10. 出处:「7.3 梯度累积的实现」第 33 段(text/45-ch07-03-7-3.txt:33,搜「累积后进行一次参数更新」)。 ↩

  11. 出处:「9.2 基于梯度累积的优化技巧」第 48 段(text/57-ch09-02-9-2.txt:48,搜「模拟64的批量大小」)。 ↩

  12. 出处:「9.2 基于梯度累积的优化技巧」第 252-256 段(text/57-ch09-02-9-2.txt:252,搜「更新变得缓慢」)。 ↩

  13. 出处:「7.3 梯度累积的实现」第 79 段(text/45-ch07-03-7-3.txt:79,搜「攒够再出手」)。 ↩