Roofline 模型详解

时间都花到哪里了？

我们从一个极其简单的问题开始：为什么一个算法耗时 50ms，而不是 50s 或 5ms？模型内部究竟发生了哪些相当耗时的事情，我们又应当预期它们花费多久？

计算： 深度学习模型实际上就是一系列矩阵乘法，每次矩阵乘法又由浮点乘法和加法“运算”（FLOPs）构成。加速器的速度决定了完成这些计算需要多长时间：

例如，NVIDIA H100 每秒大约可执行 9.89e14 次 bfloat16¹ 浮点运算，而 TPU v6e 每秒可执行 9.1e14 次浮点运算。² 这意味着，在 H100 上执行 1e12 FLOPs 大约需要 1e12 / 9.89e14 = 1.01ms，在 TPU v6e 上则需要 1e12 / 9.1e14 = 1.1ms。³

芯片内部通信： 在一个加速器内部，张量需要在加速器内存（HBM）与计算核心之间传输。你会看到这条链路的带宽被称为“HBM 带宽”。⁴ H100 的这一带宽约为 3.35TB/s，TPU v6e 则约为 1.6TB/s。

芯片之间的通信： 当我们把模型分布到多个加速器上时，张量经常需要在加速器之间传输。硬件通常会提供若干种选择（ICI、DCN 和 PCIe），每种链路的带宽都不同。

无论通信发生在芯片内部还是芯片之间，我们都用 bytes/s 衡量它，并用下面的公式估算总通信时间：

通常情况下（但并非总是如此），单芯片内的计算可以与芯片内部和芯片之间的通信重叠。这意味着，取计算时间与通信时间二者中的较大值，就能得到训练和推理耗时的下界；而将两者相加，则能得到上界。实践中，我们针对二者的最大值进行优化，因为这样代数更简单，而且通过重叠通信与计算，通常可以接近这一界限。如果以最大值为优化目标，那么上下界最多相差 2 倍，因为 $T_\text{math} + T_\text{comms} \leq 2 * \max(T_\text{math}, T_\text{comms})$。在此基础上，我们还可以对“重叠区域”和各种开销建模，从而提高估算精度；这些模型可以由目标系统上具体模型的性能剖析结果来校准。

如果假设通信与计算可以完美重叠，那么当 $T_\text{math} > T_\text{comms}$ 时，硬件会得到充分利用。我们称这种状态为“计算受限”。当 $T_\text{comms} > T_\text{math}$ 时，系统往往“通信受限”⁵，加速器至少有一部分 FLOPs/s 会浪费在等待数据传输上。判断一项运算究竟受计算还是通信限制的一种方法，是考察它的“算术强度（arithmetic intensity）”或“运算强度（operational intensity）”。

定义： 算法的算术强度，等于它执行的 FLOPs 总数与它需要通信的字节数之比；这里的通信既可以发生在芯片内部，也可以发生在芯片之间。

算术强度衡量一次运算的“每字节 FLOPs”。从一阶近似来看，算术强度较高时，$T_\text{math}$ 相对于 $T_\text{comms}$ 较大，通常能够用到大部分可用 FLOPs。反之，我们会把更多时间花在通信上，从而浪费 FLOPs。二者发生转换的位置称为硬件的“临界算术强度”，也就是加速器峰值 FLOPs/s 与加速器带宽之比。

$\text{Intensity}(\text{Accelerator})$ 表示加速器达到峰值 FLOPs/s 时所需的算术强度。对于 TPU v5e 的 MXU，这一数值约为 240 FLOPs/byte，因为该 TPU 每秒可执行 1.97e14 次浮点运算，并以每秒 8.2e11 字节的速度从 HBM 加载数据。⁶ 这意味着，如果一个算法的算术强度低于 240 FLOPs/byte，它就会受字节加载速度限制，因而无法充分利用硬件。⁷ 下面来看这样一个例子：

示例（点积）： 为了以 bfloat16 精度计算两个向量的点积 x • y: bf16[N], bf16[N] → bf16[1]，我们需要从内存加载 $x$ 和 $y$，两者各占 $2 * N = 2N$ 字节；随后执行 $N$ 次乘法和 $N-1$ 次加法，并向 HBM 写回 $2$ 字节。

当 $N\rightarrow\infty$ 时，上式趋近于二分之一。因此，点积的算术强度是 $\frac{1}{2}$；换句话说，每加载一个字节，点积只执行 0.5 次浮点运算。这意味着算法的算术强度低于硬件的算术强度，因此系统会受内存带宽限制。⁸

Roofline 可视化

我们可以用 Roofline 图直观表示内存与计算之间的权衡。图中，纵轴是算法在硬件上理论可达到的峰值 FLOPs/s（吞吐量），横轴则是算法的算术强度。下面是一个双对数坐标示例：

在上图中，随着强度增加（从左向右移动），算法性能（FLOPs/s）起初呈线性增长，直到达到硬件的临界算术强度；对于 TPU v5e，这一数值是 240。强度低于该值的算法会受带宽（BW）限制，其性能由峰值内存带宽决定（红色区域）。右侧的算法则会充分利用可用 FLOPs（绿色区域）。这里，Algo 1 受内存带宽限制，只使用了硬件 FLOPs/s 总量的一部分；Algo 2 受计算限制。通常，我们可以通过提高算法的算术强度，或者增加可用内存带宽（从 BW1 移动到 BW2）来提升算法性能。

矩阵乘法

下面来看一个即将成为我们最喜欢的算法：矩阵乘法（简称 matmul）。我们写作 $X * Y \rightarrow Z$，其中 $X$ 的形状是 $\text{bf16}[B, D]$，$Y$ 的形状是 $\text{bf16}[D, F]$，$Z$ 的形状是 $\text{bf16}[B, F]$。为了执行这次矩阵乘法，需要加载 $2DF + 2BD$ 字节，执行 $2BDF$ FLOPs，并写回 $2BF$ 字节。^{9 10} 因此：

如果假设“批大小” $B$ 相对于 $D$ 和 $F$ 很小，就能得到一个很漂亮的简化：

对于 Transformer 矩阵乘法，这是一个合理假设，因为局部（每副本）批大小通常满足 $B < 1024$ 个 token（不是序列），而 $D$ 和 $F > 8000$。因此，当每个模型副本的批大小¹¹超过 240 个 token 时，系统通常会变为计算受限——这是一条非常简单的规则！

要点： 要让 bfloat16 矩阵乘法在大多数 TPU 上进入计算受限状态，每个模型副本的 token 批大小必须大于 240。¹²

这条规则有几个值得注意的限定条件，下面的练习会进一步讨论，尤其是量化情形（例如，激活值已量化，但仍执行全精度 FLOPs）。不过，它仍是一条值得记住的经验法则。对于 GPU，这个数值略高一些（更接近 300），但总体结论相同。当我们把一个大矩阵乘法分解成更小的矩阵乘法时，分块大小也很重要。¹³ 下一章将讨论更底层的 GPU 和 TPU 细节，先从如何理解 TPU开始。

网络通信 Roofline

到目前为止，我们讨论的 Roofline 都针对内存带宽，而且全部发生在单个芯片内部。但这并不意味着它是一条普遍规律。事实上，本书最关心的大多数 Roofline 都涉及芯片之间的通信：通常是对分片到多个 TPU 上的矩阵执行矩阵乘法。

举一个略显刻意的例子：假设我们要把两个大矩阵 $X\sim \text{bf16}[B, D]$ 和 $Y \sim \text{bf16}[D, F]$ 相乘，并沿 $D$ 维将它们均匀拆分到 2 个 TPU/GPU 上。为了完成乘法（将在第 3 章看到），可以在每个 TPU 上计算各自的一半矩阵：TPU 0 执行 Z0 = X[:, :D // 2] @ Y[:D // 2, :]，TPU 1 执行 Z1 = X[:, D // 2:] @ Y[D // 2:, :]；随后把得到的“部分和”复制到另一个 TPU，并把它们相加。假设每个方向的复制带宽是 4.5e10 bytes/s，而每块芯片可执行 1.97e14 FLOPs/s，那么 $T_\text{math}$ 和 $T_\text{comms}$ 分别是多少？

$T_\text{math}$ 显然是原先的一半，因为每个 TPU 只完成一半工作，即：¹⁴

那么 $T_\text{comms}$ 呢？现在，它表示芯片之间的通信时间！它就是发送的总字节数除以网络带宽，即：

因此，当 $\text{Intensity}(\text{matmul (2-chips)}) > \text{Intensity}(\text{TPU w.r.t. inter-chip network})$ 时，系统会进入计算受限状态（这次是相对于芯片间网络而言）。等价地，需要满足 $\frac{BDF}{2BF} = \frac{D}{2} > \frac{1.97e14}{4.5e10} = 4377$，也就是 $D > 8755$。请注意，与之前不同，现在的临界阈值取决于 $D$，而不是 $B$！请思考其中的原因。这只是一个例子，但它说明：要判断一项运算能否并行扩展到多个 TPU，这类 Roofline 至关重要。

练习题

问题 1［int8 矩阵乘法］： 假设我们想以 int8 精度（每个参数 1 字节）而不是 bfloat16 精度（每个参数 2 字节）执行矩阵乘法 $X[B, D] \cdot_D Y[D, F] \rightarrow Z[B, F]$¹⁵，因为 TPU/GPU 能以更低精度更快地执行矩阵乘法。

1. 需要从内存加载多少字节？需要向内存写回多少字节？
2. 总共执行多少次运算？
3. 算术强度是多少？
4. $T_\text{math}$ 和 $T_\text{comms}$ 的 Roofline 估算值分别是多少？整个运算运行时间的合理上下界是什么？

假设 HBM 带宽为 8.2e11 bytes/s，int8 峰值运算速率为 3.94e14 OPs/s（约为 bfloat16 的 2 倍）。

问题 2［int8 + bf16 矩阵乘法］： 实践中，我们经常对权重和激活值采用不同的量化方案，因此可能以很低的精度存储权重，却让激活值（以及计算）保持较高精度。假设权重被量化为 int8，而激活值（以及计算）保持 bfloat16。批大小达到多少时系统会进入计算受限状态？假设 bfloat16 性能为 1.97e14 FLOPs/s。

提示：具体来说，就是执行 bf16[B, D] * int8[D, F] -> bf16[B, F]，其中 $B$ 是“批大小”。

问题 3： 沿用问题 2 的设置，绘制峰值 FLOPs/s 随 $B$ 变化的 Roofline 图，其中分别取 $F = D = 4096$ 和 $F = D = 1024$。请使用实际加载的准确字节数，而不是近似值。

import matplotlib.pyplot as plt
import numpy as np
 
bs = np.arange(1, 512)
 
def roofline(B, D, F):
  total_flops = 2*B*D*F
  flops_time = total_flops / 1.97e14
  comms_time = (2*B*D + D*F + 2*B*F) / 8.2e11
  total_time = np.maximum(flops_time, comms_time)
  return total_flops / total_time
 
roofline_big = roofline(bs, 4096, 4096)
roofline_small = roofline(bs, 1024, 1024)
 
plt.figure(figsize=(8, 4))
plt.plot(bs, roofline_big, label='F=D=4096')
plt.plot(bs, roofline_small, label='F=D=1024')
plt.legend()
plt.xlabel('batch size')
plt.ylabel('peak bfloat16 FLOPs/s on TPU v5e')
plt.grid()

import matplotlib.pyplot as plt
import numpy as np
 
bs = np.arange(1, 512)
 
def roofline(B, D, F):
  total_flops = 2*B*D*F
  flops_time = total_flops / 1.97e14
  comms_time = (2*B*D + D*F + 2*B*F) / 8.2e11
  total_time = np.maximum(flops_time, comms_time)
  return total_flops / total_time
 
roofline_big = roofline(bs, 4096, 4096)
roofline_small = roofline(bs, 1024, 1024)
 
plt.figure(figsize=(8, 4))
plt.plot(bs, roofline_big, label='F=D=4096')
plt.plot(bs, roofline_small, label='F=D=1024')
plt.legend()
plt.xlabel('batch size')
plt.ylabel('peak bfloat16 FLOPs/s on TPU v5e')
plt.grid()

问题 4： 如果我们要执行 $\text{int8}[B, D] \cdot_D \text{int8}[B, D, F] \rightarrow \text{int8}[B, F]$，也就是设想每个批元素都有一个不同的矩阵，那么这项运算的算术强度是多少？

问题 5［GPU 的内存 Roofline］： 根据 NVIDIA 提供的 H100 SXM 规格表，计算 bfloat16 矩阵乘法进入计算受限状态时的批大小。请注意，表中的 Tensor Core FLOPs 数字是实际值的两倍，因为只有使用结构化稀疏时才能达到该数值。

第一部分的第 1 章到此结束！第 2 章将考察真实 TPU 如何处理 FLOPs 与通信，点击这里继续。