加速器与网络

一次训练运行的成败取决于带宽，而非峰值 FLOPs。更有用的视角是一套带宽层级：紧贴加速器裸片的高带宽显存（HBM）供给单块芯片，NVLink 与 NVSwitch 连接节点内加速器，InfiniBand 或 RoCE 在节点之间搬运字节。正是这套层级解释了为什么张量并行组几乎不越过 NVLink 边界，为什么数据并行与流水线并行可以跨节点，以及模型 FLOPs 利用率如何衡量硬件底座有没有被真正用上。

一次规格表无法解释的运行

设想一次前沿运行铺在数千张加速器上连续跑数周。规格表许诺了一个峰值算术速率，可这次运行只交付了其中一小部分。缺失的那一块，很少是矩阵乘法单元因为没有工作可做而闲置。缺失的是这些单元所等待的字节：一块芯片访问自身显存有多快、机箱内两块芯片交换数据有多快、跨机箱交换数据又有多快。这三个数字每往外走一步就相差一个数量级以上，而一个并行轴一旦需要的通信量超出其层级所能负担，就会拖住整个集合通信，也就是那种跨设备协同的通信操作。

所以只看规格表上的峰值 FLOPs 会把问题看偏。更合适的起点是它下面那套带宽层级，因为每一个布局决策，本质上都在问某次交换应该落在哪一层。这里的硬件底座包括加速器、互连层级，以及说明这一底座有没有被用上的核心数字，MFU。建立在其上的并行算法属于 第 10 章；集群编排、检查点与数据层则属于 第 65 章。

三层带宽

这台机器是一套带宽层级。顺着一个字节从裸片一路向外，就能看清每一层为何存在、又能容纳什么。

最内的一层是芯片及其显存。一块加速器把巨大的矩阵乘法吞吐能力，与紧贴裸片的HBM配在一起。前向与反向的矩阵乘法都从 HBM 读入数据，因此往往是 HBM 带宽、而非原始 FLOPs，决定了一个内核可达的速率。

拖动曲线，看看为何一个每取一字节只做很少工作的内核会停在带宽受限处，达不到规格表许诺的峰值 FLOPs。

下一层是节点内：一个机箱里的那几块加速器。在 NVIDIA 硬件上，它们用 NVLink 这种专用芯片互连，再配上 NVSwitch 连接起来；NVSwitch 让节点内每块加速器都能以完整的 NVLink 带宽访问其他任一块，而不必走更慢的 PCIe 总线。这一层最适合安放通信需求最高的轴，因为 NVLink 带宽足够高，可以把一次集合通信藏在计算之下。一次集合通信究竟能不能藏住，取决于一个比值，正如 图 62.2 所勾画：只要一步的计算多于通信，集合通信就完全消失在计算之下；而节点间网络之所以力有不逮，恰恰是因为它的字节昂贵到足以留下一部分通信暴露在外。

最外的一层是节点间：机箱通过数据中心网络彼此通信。这里通常使用 InfiniBand 或 RoCE（RDMA over Converged Ethernet），以胖树或 rail-optimized 这样的拓扑连接，每块加速器配一张网卡（NIC）来设定每设备的注入带宽。这一层每字节比 NVLink 慢一个数量级，因此只有那些通信节俭的轴才负担得起跨越它。

把并行映射到层级上

层级一旦固定，每个并行轴该放在哪里也就随之确定，而且取决于通信量，而非偏好。张量并行（TP）把单个矩阵乘法切分到各设备上，每层要用两次 all-reduce 交换激活；激活是设备间传递的中间值，all-reduce 则指每块设备把自己的部分结果与所有其他设备相加。因此它是对带宽需求最高的轴，通常受限于节点内的 NVLink 层 (Shoeybi et al. 2019)。数据并行与流水线并行每步通信少得多：DP 每步只做一次梯度归约，PP 每个阶段边界只做一次激活交接，因此它们是被允许跨越节点间网络的轴 (Narayanan et al. 2021)。MoE 的专家并行使用自己的 all-to-all 组，留到 第 9 章 再谈。由此得出的经验法则，TP 在节点内、DP/PP 跨节点，并不是一种约定，而是带宽鸿沟的必然后果，正如 图 62.5 所展示。

TPU 的岔路：环面而非交换机

张量处理器（TPU） pod 是 Google 一侧加速器集群的对应物，成本模型基于另一套 JAX/XLA，恰好说明上述带宽层级只是一种设计选择，而非一条自然定律。这些芯片通过芯片间互连（ICI），也就是 pod 内部的芯片间互连，与一台光路交换（OCS）连成一个环面拓扑，而不是交换式胖树，于是邻居到邻居的带宽充裕，切分也是通过 GSPMD 风格的注解与环面形状协同设计出来的，而不是事后再映射到一棵胖树上 (Xu et al. 2021)。同一个字节，在 图形处理器（GPU） 集群，也就是用图形处理器作加速器的集群上，会穿过一台交换机，在 TPU pod 上却走向一个物理邻居，这正是那里带宽随局部性扩展的原因。

MFU：究竟用上了多少硬件

说明这一切有没有奏效的核心数字，是模型 FLOPs 利用率。MFU 是模型实际所需的 FLOPs，除以硬件在同一段实际耗时内本可交付的峰值 FLOPs：

$$\mathrm{MFU} = \frac{\text{model FLOPs per step}}{\text{peak hardware FLOPs} \times \text{step time}}$$

它只计入模型有用的那部分算术。相近的指标 硬件 FLOPs 利用率（HFU），还把硬件确实做了、但模型并非严格需要的工作也算进来，其中最显著的是激活重算带来的那些额外前向传播；这也正是为何只要开了重算，HFU 就会高于 MFU。MFU 回答的是预算关心的问题：已经购买和供电的算术能力中，有多少真正推动了模型训练。

这些层级为何长成这样

互连层级之所以演化成今天这样，是因为带宽鸿沟扩大的速度，超过了任何单一互连所能弥合的速度。早期的多 GPU 训练倚靠 PCIe，可它对切分一个矩阵乘法已经太慢，于是 NVLink、继而 NVSwitch 被引入，正是为了让节点内这一层快到足以容纳张量并行 (Shoeybi et al. 2019)。加速器族系也在扩展原生精度集合：每一代都拓宽矩阵乘法单元可以直接执行的格式，Hopper 级及更新的硬件加入了 fp8，让同一块硅每搬运一字节交付更多矩阵乘法吞吐 (Micikevicius et al. 2022)（其内核层面的后果在 第 34 章）。网络一侧，InfiniBand 与 RoCE 收敛到 远程直接内存访问（RDMA），也就是绕过 CPU 直接读写远端内存的远程直接内存访问，以及 rail-optimized 拓扑上，后者给每块加速器一条通往其它节点对端的专用路径，使节点间这一层得以增长，而不必让每条流争抢同一批交换端口。TPU pod 则沿一条完全不同的路线演化，选了环面而非交换式底座，使带宽随局部性扩展。

每个轴各自的代价

每个并行轴缓解一种不同约束，也支付一种不同成本；互连层级决定这些成本是否承受得起。

• 张量并行对带宽。 TP 削减每设备的显存与延迟，代价是每层两次 all-reduce 消耗节点内带宽。这笔开销只有在 NVLink 上才承受得起，因此 TP 度以 NVLink 域为上限。一旦把它推过节点，all-reduce 就落到节点间网络上，在那里藏不进计算之下，只会表现为 MFU 损失。图 62.6 展示了这个临界点：TP 仍处在 NVLink 域之内时，MFU 维持一段平台；一旦其度溢出到节点间网络上，就立刻下降。

改动 nvlink_size 与 slowdown，观察平台在何处结束、下降有多快。

import numpy as np
import matplotlib.pyplot as plt
nvlink_size = 8       # 以 NVLink 带宽可达的加速器数
peak = 0.55           # 通信被计算隐藏时的 MFU 平台
slowdown = 10.0       # 节点间字节比 NVLink 贵这么多倍
tp = np.arange(1, 33)
# 一旦 TP 溢出，每步因暴露（未隐藏）的 all-reduce 而损失的比例
comm = np.where(tp <= nvlink_size, 0.0, (tp - nvlink_size) / tp * (1 - 1 / slowdown))
mfu = peak * (1 - comm)
plt.plot(tp, mfu, marker='o')
plt.axvline(nvlink_size, ls='--', color='gray', label='NVLink 边界')
plt.xlabel('张量并行度'); plt.ylabel('MFU')
plt.ylim(0, 0.6); plt.legend(); plt.title('TP 跨出节点后陡降')
print('TP=8 的 MFU:', round(mfu[7], 3), ' TP=16 的 MFU:', round(mfu[15], 3))
plt.show()

• 流水线并行对气泡。 PP 在带宽上廉价，每个阶段边界只有一次激活交接，因此能从容跨越节点，但要在每步的填充与排空处付出一个流水线气泡。
• MFU 对显存。 重算在反向传播里重做前向工作，让一个更大的模型或更长的序列能够放进显存，这会提高 HFU 却拉低 MFU。究竟优化哪个数字，取决于约束落在显存还是吞吐上。

MFU 如何出现在时间线里

在这些层级之间搬运字节的集合通信，由 NVIDIA 硬件上的 NCCL（在 AMD 上为 RCCL）提供，它把每个集合通信映射到拓扑上，根据各 rank 所在之处选择环或树，再把它路由到 NVLink 或网络上 (NVIDIA 2024)。从业者很少直接调用它，调用的是并行框架。真正要读的是每步的剖析，因为 MFU 是从时间线里诊断出来的，而不是从损失曲线。没能藏进计算之下的暴露通信、过大的流水线气泡、过多的重算，还有未融合的内核，各自在达成与峰值 FLOPs 之间的差距里占去一份，而时间线正是每一份可见之处。

图 62.7 把时间线读作一次从峰值 FLOPs 下降到 MFU 的过程。空闲时间，即一块设备等待暴露通信、或停在流水线气泡里的那段，是纯粹的损失。在硬件随后实际跑出的 FLOPs 中，重算做了模型并非严格需要的工作：这些工作把 HFU 抬到 MFU 之上，但只有有用的算术才计入那个真实数字。

延伸阅读