如何理解 TPU

你可能也会喜欢阅读有关 NVIDIA GPU 的新第 12 节！

TPU 是什么？

TPU 基本上就是一个专门进行矩阵乘法的计算核心（称为 TensorCore），外接一组高速内存（称为高带宽内存，即 HBM）[tpu_paper]。 示意图如下：

你可以把 TensorCore 基本理解为一台极其擅长矩阵乘法的机器，不过它还有几个值得注意的功能。TensorCore 包含三个关键单元：

• MXU（矩阵乘法单元，Matrix Multiply Unit）是 TensorCore 的核心。对于大多数代际的 TPU，它会使用脉动阵列，每 8 个周期执行一次 bf16[8,128] @ bf16[128,128] -> f32[8,128] 矩阵乘法¹（详见附录 B）。

  • 对于运行频率为 1.5GHz 的 TPU v5e，每个 MXU 大约能达到 5e13 bf16 FLOP/s。大多数 TensorCore 拥有 2 个或 4 个 MXU，因此，例如 TPU v5e 的 bf16 总吞吐量为 2e14 FLOP/s。
  • TPU 还支持吞吐量更高的低精度矩阵乘法（例如，每块 TPU v5e 芯片可以执行 4e14 int8 OP/s）。

• VPU（向量处理单元，Vector Processing Unit）执行一般数学运算，例如 ReLU 激活，或向量之间的逐点加法、乘法。归约（求和）也在这里执行。附录 A提供了更多细节。

• VMEM（向量内存，Vector Memory）是 TensorCore 内部靠近计算单元的一块片上暂存区。它比 HBM 小得多（例如，TPU v5e 上为 128 MiB），但到 MXU 的带宽高得多。VMEM 的工作方式有些类似于 CPU 上的 L1/L2 缓存，但容量大得多，而且由程序员控制。TensorCore 要对 HBM 中的数据执行任何计算之前，必须先把数据复制到 VMEM。

TPU 执行矩阵乘法的速度非常、非常快。 这基本就是它们的主要工作，而且它们做得很好。TPU v5p 是迄今最强大的 TPU 之一，每个核心可以执行 2.5e14 bf16 FLOP/秒，每块芯片可以执行 5e14 bf16 FLOP/秒。一个包含 8960 块芯片的 Pod 可以达到 4 bf16 exaFLOP/s。这个数字非常庞大。这相当于世界上最强大的超级计算机之一。而且 Google 拥有很多这样的机器。²

上图还包含 SMEM 和标量单元等其他一些部件，它们用于处理控制流，并会在附录 A中简要讨论；不过，理解它们并非关键。另一方面，HBM 很重要，而且相当简单：

• HBM（高带宽内存，High Bandwidth Memory）是一大块高速内存，用来存储供 TensorCore 使用的张量。HBM 的容量通常为数十 GB（例如，TPU v5e 拥有 16GiB HBM）。

  • 计算需要这些张量时，张量会从 HBM 流出，经由 VMEM（见下文）进入 MXU；计算结果则从 VMEM 写回 HBM。

  • HBM 与 TensorCore 之间（途经 VMEM）的带宽称为“HBM 带宽”（通常约为 1–2TB/秒），它会限制内存受限工作负载的计算速度。

一般而言，TPU 上的所有操作都会流水线化并相互重叠。 为了执行矩阵乘法 $X \cdot A \to Y$，TPU 首先需要把矩阵 $A$ 和 $X$ 的数据块从 HBM 复制到 VMEM，再把它们加载到 MXU；MXU 会把 8x128 的数据块（来自 $X$）与 128x128 的数据块（来自 $A$）相乘，之后再把结果逐块复制回 HBM。为了高效完成这一过程，矩阵乘法会被流水线化，使往返 VMEM 的复制与 MXU 工作相互重叠。这样，MXU 可以持续工作，无需等待内存传输，从而让矩阵乘法受计算限制，而不是受内存限制。

下面展示了如何对来自 HBM 的数据执行逐元素乘积：

矩阵乘法看起来几乎完全相同，只不过数据会加载到 MXU，而不是 VPU/向量单元；加载和存储的顺序也会不同，因为同一个权重数据块会供多个激活值数据块使用。你可以看到数据块先流入 VMEM，然后进入 VREG（向量寄存器），再进入向量单元，最后回到 VMEM 和 HBM。我们马上就会看到，如果从 HBM 到 VMEM 的加载速度比向量单元（或 MXU）执行 FLOP 的速度慢，就会变成“带宽受限”，因为 VPU 或 MXU 得不到足够的工作。

关键要点： TPU 非常简单。它们把权重从 HBM 加载到 VMEM，再从 VMEM 加载到脉动阵列；该阵列每秒可执行约 200 万亿次乘加。HBM $\leftrightarrow$ VMEM 和 VMEM $\leftrightarrow$ 脉动阵列的带宽，对 TPU 能够高效执行哪些计算设定了根本限制。

VMEM 与算术强度： VMEM 比 HBM 小得多，但它到 MXU 的带宽高得多。正如第 1 节所见，这意味着，如果某个算法能把全部输入/输出都装入 VMEM，就不太容易遇到通信瓶颈。当计算的算术强度较低时，这尤其有用：VMEM 带宽约为 HBM 带宽的 22 倍，这意味着，从 VMEM 读取或向其写入的 MXU 运算，只需达到 10–20 的算术强度，就能实现峰值 FLOP 利用率。也就是说，如果能把权重放进 VMEM 而不是 HBM，矩阵乘法就可以在小得多的批大小下达到计算受限。而且，即使某些算法从根本上具有较低的算术强度，它们仍然可以高效运行。只是 VMEM 太小，因此这往往是个挑战。³

一块 TPU 芯片通常（但并非总是）由两个共享内存的 TPU 核心组成；它们可以视为一个 FLOP 翻倍的大型加速器，这种配置称为“megacore”。v4、v5 和 v6 TPU 都是如此（TPU v7 取消了 megacore，改为在两个核心之间使用高带宽链路）。较早的 TPU 芯片拥有彼此独立的内存，被视为两个独立加速器（TPU v3 及更早版本）。为推理优化的芯片（如 TPU v5e）每块芯片只有一个 TPU 核心。

芯片会以每 4 块一组的方式安装在“托盘”上，并通过 PCIe 网络连接到 CPU 主机。 这是大多数读者熟悉的形式：通过 Colab 或一台 TPU-VM 暴露 4 块芯片（8 个核心，不过通常视为 4 个逻辑 megacore）。对于 TPU v5e 这类推理芯片，每台主机有 2 个托盘，而不是 1 个；但每块芯片也只有 1 个核心，因此总计为 8 块芯片 = 8 个核心。⁴

PCIe 带宽有限： 与 HBM $\leftrightarrow$ VMEM 链路一样，CPU $\leftrightarrow$ HBM 的 PCIe 连接具有特定带宽，它限制了从主机内存加载到 HBM、或反向传输的速度。例如，TPU v4 的 PCIe 带宽在每个方向上都是 16GB/秒，因此比 HBM 慢近 100 倍。我们可以在主机（CPU）RAM 与加速器之间加载或卸载数据，但速度不会很快。

TPU 网络

在 Pod 内，芯片通过 ICI 网络相互连接。 在较早的代际（TPU v2 和 TPU v3）、推理芯片（例如 TPU v5e）以及 Trillium（TPU v6e）中，ICI（“芯片间互连”）会连接最近的 4 个邻居（并通过边缘链路形成二维环面）。TPU v4 和 TPU v5p 则连接最近的 6 个邻居（形成三维环面）。请注意，这些连接不会经过主机；它们是芯片之间的直接链路。

环面结构把任意两个节点之间的最大距离从 $N$ 降到 $N / 2$，使通信快得多。TPU 还有一种“扭曲环面”配置，它采用类似莫比乌斯带的拓扑来环绕环面，从而进一步缩短节点间的平均距离。

TPU Pod（通过 ICI 连接）可以变得非常庞大： TPU v4 的最大 Pod 大小（称为 SuperPod）是 16x16x16，TPU v5p 则是 16x20x28。这些大型 Pod 由可重新配置的 4x4x4 芯片立方体组成，立方体通过光学环绕链路⁵连接；我们可以重新配置这些链路，以连接非常大的拓扑。

也可以申请较小的拓扑（例如 2x2x1、2x2x2），不过它们没有环绕连接。这是个重要的限制，因为它通常会让大多数通信的时间翻倍。只要拓扑包含完整立方体的整数倍（例如 4x4x4 或 4x4x8），光学交换机就会提供环绕连接。⁶

TPU v5e 和 Trillium Pod 由单个 16x16 二维环面组成；任一大小为 16 的轴都有环绕连接（这意味着 8x16 会在长轴上拥有环绕连接）。TPU v5e 和 v6e（Trillium）无法扩展到超过 16x16 的环面，但不同 Pod 仍可通过标准数据中心网络（DCN）通信；DCN 会把各 TPU 主机相互连接。同样，可以申请较小的拓扑，但大小 $<16$ 的维度没有环绕连接。

这种最近邻连接是 TPU 与 GPU 之间的关键区别。 GPU 通过交换机层级结构连接，近似让每块 GPU 都能点对点连接，而不是像 TPU 那样采用局部连接。通常，同一节点内的 GPU（H100 为 8 块，B200 NVL72 最多为 72 块）会直接连接；更大的拓扑则要求每对 GPU 之间经过 O(log(N)) 跳。一方面，这意味着 GPU 可以在少量跳数内发送任意数据。另一方面，TPU 便宜得多（因为 NVLink 交换机价格高昂）、布线更简单，而且可以扩展到大得多的拓扑，因为每台设备的链路数量和带宽都保持恒定。可在这里阅读更多内容。

相对于 DCN，ICI 非常快，但它仍慢于 HBM 带宽。 例如，一块 TPU v5p 拥有：

• 每块芯片 2.8e12 字节/秒（2.8 TB/s）的 HBM 带宽。
• 每条轴 9e10 字节/秒（90 GB/s）的 ICI 带宽，每块芯片有 3 条轴。⁷
• 每块 TPU 6.25e9 字节/秒（6.25 GB/s）的 DCN（出站）带宽（通过每台主机上的 1–2 个 NIC）。⁸

这意味着，把模型拆分到多块芯片时，需要谨慎避免较慢的跨设备通信成为 MXU 的瓶颈。

多切片训练： 一组通过 ICI 连接的 TPU 称为一个切片。不同切片可以通过 DCN 相互连接，例如用于连接不同 Pod 上的切片。由于 DCN 比 ICI 慢得多，我们应尽量减少计算等待来自 DCN 的数据。DCN 是主机到主机的连接，因此要通过 DCN 把缓冲区从一块 TPU 传到另一块 TPU，首先要经 PCIe 传到主机，然后经网络出站，再从目标主机网络入站，最后经 PCIe 传入 HBM。

关键要点

• TPU 很简单，在大多数情况下可以把它看作一个矩阵乘法单元，它连接着内存（极快）、通过 ICI 连接的其他芯片（相当快），以及通过 DCN 连接的数据中心其余部分（还算快）。

• 通信会受到不同网络带宽的限制；按速度从快到慢排列：

  • HBM 带宽：TensorCore 与其对应 HBM 之间。
  • ICI 带宽：一块 TPU 芯片与最近的 4 个或 6 个邻居之间。
  • PCIe 带宽：CPU 主机与其对应的芯片托盘之间。
  • DCN 带宽：多台 CPU 主机之间，通常是未通过 ICI 连接的主机。

• 在一个切片内部，TPU 仅通过 ICI 连接到最近邻。 这意味着，切片内相距较远的芯片之间通过 ICI 通信时，必须先经过中间的芯片。

• 权重矩阵的两个维度都需要填充到至少 128（TPU v6e 上为 256），才能填满 MXU（实际上，小于 128 的轴会被填充到 128）。

• 低精度矩阵乘法往往更快。 对于支持这些精度的代际，TPU 执行 int8 或 int4 OP 的速度大约是 bfloat16 FLOP 的 2 倍或 4 倍。VPU 运算仍以 fp32 执行。

• 为避免 TPU 计算单元成为瓶颈，需要确保每条通道上的通信量与其速度成比例。

TPU 规格

下面列出这些芯片的一些具体数值：

主机大小指连接到单台主机的 TPU 拓扑（例如，TPU v5e 让单台 CPU 主机连接采用 4x2 拓扑的 8 块 TPU）。关于最新一代的更多细节，请参阅 TPU7x 文档。下面是互连数据：

这里同时列出单向带宽和双向带宽，因为单向带宽更贴近硬件本身，而涉及完整环的公式往往更常使用双向带宽。⁹

PCIe 带宽通常约为每块 TPU 1.6e10 字节/秒（TPU v6e 为 3.2e10），而 DCN 带宽通常约为每块 TPU 6.25e9 字节/秒（TPU v6e 和 TPU7x 为 12.5e9，TPU v5e 为 3.125e9）。

例题

这些数字多少有些枯燥，但可以用它们对模型性能进行基本的 Roofline 估算。我们来做几道题，说明它为什么有用。第 3 部分中还会看到更多例子。

问题 1［确定 LLM 延迟边界］： 假设你想从一个采用 bf16、拆分到 32 块 TPU v4p 上的 200B 参数模型中采样。把全部参数从 HBM 加载到脉动阵列需要多长时间？提示：使用上面的数字。

问题 2［TPU 细节］： 考虑一个完整的 TPU v5e Pod。总共有多少台 CPU 主机？有多少个 TPU TensorCore？整个 Pod 的总 FLOP/s 是多少？HBM 总容量是多少？再对 TPU v5p Pod 完成同样的练习。

问题 3［PCIe 运算强度］： 假设我们不得不把一个大型权重矩阵 $A$（类型为 $\text{bf16}[D, F]$），以及一批激活值 $x$（类型为 $\text{bf16}[B, D]$）存储在主机 DRAM 中，并希望对它们执行矩阵乘法。计算在单台主机上运行，我们使用连接到该主机的一块 TPU v6e 芯片。你可以假设 $B \ll D$，且 $F = 4D$（后续章节会解释为什么这些假设合理）。要在 PCIe 上保持计算受限，所需的最小批大小 $B$ 是多少？假设 PCIe 带宽为 1.6e10 字节/秒。

问题 4［一般矩阵乘法延迟］： 假设我们想让权重矩阵 int8[16384, 4096] 乘以大小为 int8[B, 4096] 的激活矩阵，其中 B 是某个未知批大小。先假设在 1 块 TPU v5e 上执行。

1. 这次乘法的耗时如何随 B 变化？提示：分别计算从 HBM 加载数组所需的时间，以及实际执行乘法所需的时间，可能会有所帮助。哪一个构成瓶颈？
2. 如果希望从 VMEM 执行这一运算，会怎样？耗时如何随 B 变化？

问题 5［ICI 带宽］： 假设有一个 TPU v5e 4x4 切片，希望把一个类型为 bf16[8, 128, 8192] 的数组从 TPU{0,0} 发送到 TPU{3, 3}。假设 TPU v5e 每跳延迟为 $1\mu s$。

1. 第一个字节多久能到达目的地？
2. 整次传输需要多长时间？

问题 6［综合运用，难］： 假设有一个大矩阵 A：int8[128 * 1024, 128 * 1024]，它被均匀分片到 TPU v5e 4x4 切片上，但每块芯片上的分片都卸载到了主机 DRAM。现在希望把整个数组复制到 TPU{0, 0}，再乘以一个向量 bf16[8, 128 * 1024]。这需要多长时间？提示：使用上面的数字。

第 2 部分到此结束！第 3 部分将介绍分区和跨 TPU 通信，请点击这里。

附录

附录 A：进一步了解 TPU 内部结构

这里将更深入地探讨 TPU 的内部操作。除非另有说明，我们会给出 TPU v5p 的规格。

VPU

VPU 是 TPU 的向量算术核心。VPU 包含一台执行 vadd（向量加法）、vmax（逐元素最大值）等逐元素算术运算的二维 SIMD 向量机器（即 VPU），以及一组称为 VREG 的向量寄存器，用于保存供 VPU 和 MXU 使用的数据。

VREG： 每个 TPU v5p 核心拥有 64 个 32 位 VREG（TPU v4 中为 32 个），因此每个核心的 VREG 内存总计约为 64 * 8 * 128 * 4 = 256kB（整块芯片则是其 2 倍，因为有两个核心）。TPU v5p 每个周期可以从 VMEM 加载 3 个寄存器，并向 VMEM 写入 1 个寄存器。

VPU： VPU 是一个形状为 (8, 128) 的二维向量算术单元，其中大小为 128 的维度称为 lane 轴，大小为 8 的维度称为 sublane 轴。在 v5 上，每一对 (lane, sublane) 都包含 4 个彼此独立的标准浮点 ALU。VPU 在各 ALU 中用一个周期执行大多数算术指令（例如 vadd 或向量加法），延迟为 2 个周期；因此，例如在 v5 上，每个周期可以把来自 VREG 的 4 对 f32 值相加。典型的 VPU 指令可能写成 {v2 = vadd.8x128.f32 v0, v1}，其中 v0 和 v1 是输入 VREG，v2 是输出 VREG。

所有 lane 和 sublane 每个周期都以纯 SIMD 方式执行同一个程序，但每个 ALU 可以执行不同运算。因此，例如可以在一个周期内处理 1 次 vadd 和 1 次 vsub；每项运算都对两个完整 VREG 操作，并把输出写入第三个 VREG。

小测验［计算 VPU 吞吐量］： 使用上述信息，计算一块 TPU v5p 每秒可以执行多少向量 FLOP。TPU v5p 的时钟频率约为 1.75GHz。

归约： 一般而言，跨 sublane 维度的通信或归约比跨 lane 维度容易。例如，VPU 支持一种 lane 内 shuffle 操作，可以在大约一个周期内沿大小为 8 的轴滚动数据。这可用于沿 sublane 维度高效归约（只需依次移动 4、2、1 个位置，再执行 3 对逐元素求和）。

跨 lane 归约要困难得多，它会涉及一个称为 XLU（即“cross lane unit”）的独立硬件单元；这个单元很慢，而且成本相当高。

与 GPU 的比较： 对熟悉 NVIDIA GPU 的读者而言，VPU 中的每个 ALU 类似于一个 CUDA core，而单条 VPU lane 类似于“Warp Scheduler”，也就是通常由 32 个执行 SIMD 算术的 CUDA Core 构成的集合。在 lane 内归约相当容易；但如果需要跨 lane，就至少要经过 VMEM/XLU/SMEM，速度慢得多。更多细节请参阅 GPU 一节。

标量核心

标量核心是 TPU 的控制单元。它获取并分派所有指令，执行从 HBM 到 VMEM 的传输，也可以通过编程来处理标量元数据工作。由于标量核心是单线程的，一个副作用是，每个 TPU 核心每周期只能创建一个 DMA 请求。

从整体来看，一个标量核心会控制一个 VPU（由 4096 个 ALU 组成）、4 个 MXU、2 个 XLU 和多个 DMA 引擎。这种控制资源与单元计算资源高度不对称的特性，是硬件效率的一个来源，但它也限制了以任何有实际意义的方式进行数据依赖向量化的能力。

附录 B：脉动阵列如何工作？

TPU MXU 的核心是一个 128x128 脉动阵列（TPU v6e 上为 256x256）。完全饱和时，该脉动阵列每 8 个时钟周期可以执行一次 bf16[8,128] @ bf16[128,128] -> f32[8,128]¹⁰ 乘法。

• 从根本上说，脉动阵列是一个由 ALU 组成的二维 128x128（=16,384）网格，每个 ALU 都能执行一次乘法和加法运算。
• 权重（W，即 128x128 输入）从上方传入（称为 RHS），而输入（X，即 8x128 输入）从左侧传入（称为 LHS）。

下面是一段简化动画，展示一组权重（蓝色）与一组激活值（绿色）相乘。你会注意到，首先以对角方式部分加载权重（RHS），然后同样以对角方式送入激活值。下面的每一帧中，我们会把所有重叠的绿色与蓝色单元相乘，再将结果与从上方传来的任何残差相加，然后依次把结果向下传递一个单元。

下面是这段动画的更一般版本，它展示了输出如何从计算中流式送出：

下图展示了如何跨多个 RHS 和 LHS 数组建立流水线：

权重（RHS）和激活值（LHS）加载时会出现最初的流水线气泡。度过这个初始气泡后，就可以加载新的输入和权重，而不会产生额外气泡。

下面是一段不太精美的动画，展示 bf16[2, 3] x bf16[3, 3] 矩阵乘法；你可以把它想象成一个 2x3 权重矩阵与批大小为 1、维度为 3 的输入激活值相乘。与前面的幻灯片相比，这里旋转了方向，而且输入向右而不是向下流出，不过仍然可以大致看出其结构。

我们可以高效地把这一过程流水线化，从而对大矩阵执行乘法，而不会产生太大的流水线气泡。尽管如此，矩阵形状大于 MXU 的边长维度仍然很重要，而这个维度通常是 128x128。某些 TPU（自 TPU v3 起）拥有多个 MXU：TPU v3 有 2 个，TPU v4/5 有 4 个；因此需要确保分块维度大于 128 * MXU 数量。这里有一段很好的动画。

Trillium（TPU v6e）拥有 256x256 脉动阵列，这意味着它每周期可以执行 4 倍的 FLOP。这也意味着，要充分利用 MXU，张量的维度需要增大一倍。

这篇博客文章还提供了另一段出色的脉动阵列动画，其中对固定权重矩阵执行乘法。