跳到主要内容

性能 — 编译、硬件与多 GPU 的分布式训练(把一个模型分摊到多张卡、多台机器上一起训)

这一章讲三件事: 代码怎么被执行(命令式/符号式,以及异步——入队即返回、不等算完就干下一件); 硬件的关键数字(为什么「顺序读、批量传」是铁律); 多 GPU 与多机怎么协同训练(数据并行与参数服务器)。 作者自己定调:这不是硬件课,而是**「让统计建模者能做对设计决策」的生存包**—— 好的设计轻松差出一个数量级,那是「一周训完」和「错过死线」的差别1。

1. 命令式 vs 符号式:Python 解释器也是瓶颈​

我们一路写的都是命令式代码:一句一句执行,状态随语句改变。 它好写好调试,但有两个性能问题2: Python 解释器逐句调度,GPU 一快,解释器本身成了瓶颈; 而且逐句执行时框架不敢释放中间变量(不知道后面还会不会用)。

符号式编程反过来:先把整个计算定义成一张图, 编译优化后再执行。编译器能看到全貌: 可以跳过解释器、把 print((1+2)+(3+4)) 直接折成 print(10)、 及时释放不再需要的内存3。 代价是写法绕、调试难(断点打不进图里)。 实践中两派合流为混合模式:开发时用命令式,部署时一键转成符号式—— Theano/早期 TensorFlow 是符号派,PyTorch 是命令派出身,后来互相同化。

2. 异步:调用不等于执行​

GPU 操作的默认行为是入队即返回: 你调一个矩阵乘,Python 立刻拿回控制权继续跑,真正的计算在 GPU 后端排队进行4。 好处:CPU 不等 GPU,可以继续派活,流水线塞满; 坑:计时会骗你——测出来的「执行时间」可能只是入队时间, 强制等后端算完才能看到真实耗时。 框架还从计算图里读依赖:没有依赖关系的操作自动并行 (比如两个独立张量的初始化在两张卡上同时跑),不用你写并行代码。

3. 硬件生存包:记住六个数字​

这一节是全书最「硬」的部分,取对我们最有用的5。

① 内存:第一次读比后续贵 500 倍。 读内存要先发地址(~100ns),之后突发连续读每次只要 ~0.2ns—— 随机访问是头号性能杀手,顺序读写是铁律。 (GPU 显存同理,只是位宽更大、带宽更高。)

② HDD:机械时代结束了。 7200 转/分钟,等盘片转到位平均要 8ms—— 约 100 IOPS,这个数字二十年没涨过;带宽 100-200MB/s。 今天的定位:归档。

③ SSD:快,但有脾气。 比 HDD 快三个数量级(几十万 IOPS),但只能按块(256KB+)整体擦写, 随机写很差;存储单元会磨损(几千次写入)—— 大日志、swap 别放 SSD;NVMe 直挂 PCIe(CPU 与高速设备之间的点对点总线),能到 8GB/s。

④ CPU 缓存(CPU 旁边比内存快得多的小容量暂存)的暗坑:false sharing。 两个线程写同一条缓存行里的不同变量,会互相把对方的缓存踢掉—— 多核代码因此可能比单核还慢。

⑤ 训练卡 ≠ 推理卡。 推理只跑前向、不存中间值,FP16/INT8 精度就够用(T4 这类卡); 训练要存全部中间值反传、梯度累加要防下溢, 至少 FP16 混合精度(训练里高低精度混用:前向半精度省显存、累加单精度防下溢),显存更快更大(V100/A100 这类卡)—— 买错卡,钱白花6。

⑥ tensor core:为矩阵乘而生的电路。 GPU 里专门加速 4×4 到 16×16 小矩阵乘的单元; Google TPU 走了另一个极端——用 systolic array 做超大规模的矩阵乘。 代价是通用性:GPU 不擅长中断和稀疏数据7。

作者给的工程建议同样值钱:把算法对到硬件上 (参数能装进缓存时,提速可以是数量级的); 写新算法先在纸上估算性能,和实测差一个数量级以上就要警惕。

4. 数据并行与 allreduce:多 GPU 的最小可行方案​

多张 GPU 怎么一起训?三种切法:按网络结构切、按层切、按数据切。 显存够时,数据并行最方便8:

  每张 GPU 各存一份完整的模型参数(值保持同步)
  每个 minibatch 切成 k 份,每张卡算自己那份的损失与梯度
  各卡的梯度汇总(allreduce)成全 batch 的梯度
  聚合梯度广播回每张卡,各自更新——参数重新同步

  图说:每张卡干的活和单卡训练一模一样,
  只是每步多一次「梯度汇总-广播」的通信。

allreduce 是这里唯一的原语:把多张卡上的向量逐元素求和, 结果广播回所有卡——因为是可交换、可结合的操作,各卡到达顺序无所谓9。 通信是这里的新成本:梯度总量等于参数总量,每步都要过一遍—— 计算与通信的比值,决定了加卡的收益。 模型小、卡间互联慢(NVLink 100GB/s vs PCIe 32GB/s vs 100GbE 网卡 10GB/s), 加卡就不划算。

5. 参数服务器:跨机时的带宽账​

单机内 allreduce 简单;跨机时,中央聚合点成为瓶颈: m 台 worker 都把梯度发给一台服务器,它的时间是 O(m)—— 人越多,等得越久10。 解法:服务器也从 1 台扩到 n 台,每台只存 1/n 的参数, 总时间降到 O(m/n)——台数匹配,就是常数级扩展。 实践中 worker 和 server 就是同一批机器,一机两用。

给统计建模者的抽象是 push/pull 两个动词11:

  • push(key, value):worker 把本地梯度推到公共存储,在那里聚合(求和);
  • pull(key, value):worker 拉回聚合后的值。

所有同步的脏活(谁先发、谁后到、哪台挂了)都藏在这两个动词后面—— 建模者写优化,系统工程师管同步,互不打扰。 这个设计和 Dynamo 这类键值存储神似,不是巧合: 分布式参数管理,本质就是带自定义更新语义(每次写入该怎么合并的规则)的键值存储。

6. 作者的判断与证据​

书里给了证据的: 延迟数字表(引 Jeff Dean 2010 与 Colin Scott 的更新); 内存 500 倍、HDD 100 IOPS、SSD 擦写与磨损的机制解释; 数据并行的完整流程;参数服务器的 O(m)→O(m/n) 账。

经验判断: 「数据并行几乎总是首选」(前提是显存够); 「push/pull 抽象值得」是系统设计判断; 混合编程「开发命令式、部署符号式」是框架演进的共识路线。

判断(我们的,不是书里的): 这一章和第 13 章的 scaling law 是同一件事的两侧:「把规模做大」在算法侧是幂律, 在系统侧就是这一章的 allreduce 与参数服务器。 今天训练万亿参数模型,瓶颈早已不是单卡算力, 而是这一章的通信账——这也解释了为什么英伟达的互联(NVLink 与 InfiniBand:前者连机内的显卡,后者连机与机,都是高速互联线路) 和它的 GPU 一样值钱。 如果错,会错在: 新的并行策略(流水并行、张量并行、ZeRO 分片) 已经超出「数据并行 + 参数服务器」的框架——本章是入门地图, 超大模型要另学一套,但通信是硬成本这一点不变。

7. 边界与局限​

  • 硬件数字 2023 年前后为准,量级正确、具体型号会过时;
  • 只讲数据并行;模型并行、流水并行原书未展开;
  • 异步与自动并行各框架实现差异大(PyTorch/MXNet/JAX 行为不同);
  • 容错(机器挂了怎么办)只开了个头。

8. 可带走的​

  1. 命令式好写、符号式好优化;混合模式:开发命令式、部署符号式;
  2. GPU 调用入队即返回;计时先等后端算完;
  3. 六数字:内存首读贵 500 倍、HDD ~100 IOPS、SSD 按块擦写会磨损、 false sharing、训练卡≠推理卡、tensor core 专精矩阵乘;
  4. 顺序读、批量传,是存储层的一切;
  5. 数据并行:参数同步 + 梯度 allreduce;通信决定加卡收益;
  6. 参数服务器:push 推梯度、pull 拉聚合;多 server 分片把 O(m) 降到 O(m/n)。

9. 原文地图​

主题原书章原文位置
命令式的两个瓶颈Compilers and Interpreterstext/89-compilers-and-interpreters.txt:25(搜「overwhelming」)
符号式的优化Compilers and Interpreterstext/89-compilers-and-interpreters.txt:35(搜「skip the Python interpreter」)
异步入队Asynchronous Computationtext/90-asynchronous-computation.txt:6(搜「enqueued」)
内存 500 倍Hardwaretext/92-hardware.txt:35(搜「500 times」)
HDD 100 IOPSHardwaretext/92-hardware.txt:50(搜「100 IOPs」)
SSD 磨损与 NVMeHardwaretext/92-hardware.txt:58(搜「wear out」) · text/92-hardware.txt:59(搜「NVMe」)
false sharingHardwaretext/92-hardware.txt:105(搜「false sharing」)
训练卡 vs 推理卡Hardwaretext/92-hardware.txt:114(搜「training or inference」)
tensor core 与 TPUHardwaretext/92-hardware.txt:131(搜「tensor cores」)
数据并行Training on Multiple GPUstext/93-training-on-multiple-gpus.txt:43(搜「partition data」) · text/93-training-on-multiple-gpus.txt:73(搜「aggregated」)
allreduceTraining on Multiple GPUstext/93-training-on-multiple-gpus.txt:177(搜「allreduce」)
O(m) 瓶颈与 O(m/n)Parameter Serverstext/95-parameter-servers.txt:74(搜「mathcal{O}(m/n)」)
push/pull 抽象Parameter Serverstext/95-parameter-servers.txt:100(搜「push(key, value)」) · text/95-parameter-servers.txt:103(搜「decouple the concerns」)

Footnotes​

  1. 出处:「Hardware」第 4 段(text/92-hardware.txt:4,搜「order of magnitude」)。 ↩

  2. 出处:「Compilers and Interpreters」第 25 段(text/89-compilers-and-interpreters.txt:25,搜「overwhelming」)。 ↩

  3. 出处:「Compilers and Interpreters」第 35 段(text/89-compilers-and-interpreters.txt:35,搜「skip the Python interpreter」)。 ↩

  4. 出处:「Asynchronous Computation」第 6 段(text/90-asynchronous-computation.txt:6,搜「enqueued」)。 ↩

  5. 出处:「Hardware」第 35 段(text/92-hardware.txt:35,搜「500 times」)、第 50 段(text/92-hardware.txt:50,搜「100 IOPs」)、第 58 段(text/92-hardware.txt:58,搜「wear out」)与第 105 段(text/92-hardware.txt:105,搜「false sharing」)。 ↩

  6. 出处:「Hardware」第 114 段(text/92-hardware.txt:114,搜「training or inference」)。 ↩

  7. 出处:「Hardware」第 131 段(text/92-hardware.txt:131,搜「tensor cores」)与第 137 段(text/92-hardware.txt:137,搜「sparse data」)。 ↩

  8. 出处:「Training on Multiple GPUs」第 55 段(text/93-training-on-multiple-gpus.txt:55,搜「most convenient way」)。 ↩

  9. 出处:「Training on Multiple GPUs」第 177 段(text/93-training-on-multiple-gpus.txt:177,搜「allreduce」);可交换性见「Parameter Servers」第 88 段。 ↩

  10. 出处:「Parameter Servers」第 74 段(text/95-parameter-servers.txt:74,搜「mathcal{O}(m/n)」)。 ↩

  11. 出处:「Parameter Servers」第 100 段(text/95-parameter-servers.txt:100,搜「push(key, value)」)与第 103 段(text/95-parameter-servers.txt:103,搜「decouple the concerns」)。Smola & Narayanamurthy 2010 提出,Li et al. 2014 系统化。 ↩