了解TPU的系统架构非常有趣。

https://t.co/3zHvGYAzRm https://t.co/ggn4ovpfc5

TPU 深度解析

来源：https://henryhmko.github.io/posts/tpu/tpu.html 脉动阵列图更新：在 Hacker News 上被推荐，感谢阅读！最近我大量接触 TPU，发现它们与 GPU 的设计理念截然不同，这非常有趣。

TPU 的主要优势在于其可扩展性。这是通过硬件（例如能效和模块化）和软件（例如 XLA 编译器）的协同设计实现的。

背景

简单概括一下，TPU 是 Google 的 ASIC，专注于两个因素：极高的矩阵乘法吞吐量和能效。

其起源可追溯到 2006 年，当时 Google 正在评估是否使用 GPU、FPGA 还是定制 ASIC。当时只有少数应用需要专用硬件，他们判断这些需求可以通过利用大型数据中心中过剩的 CPU 计算能力来满足。但在 2013 年，当 Google 的语音搜索功能运行在神经网络上，并且内部预测显示如果该功能流行起来，计算需求将大幅增长时，情况发生了变化。

时至今日，TPU 为 Google 的大部分 AI 服务提供动力。当然，这包括 Gemini 或 Veo 的训练和推理，也包括部署其推荐模型（DLRM）。

让我们深入探讨，从底层到上层来了解 TPU 的内部结构。

TPU 单芯片层面

我将以 TPUv4 的示意图为主，但这个布局大致也适用于最新一代的 TPU（例如 TPUv6p “Trillium”；TPUv7 “Ironwood” 的细节在撰写本文时的 2025 年 6 月尚未公布）。

以下是单个 TPUv4 芯片的布局：

芯片图单个 TPU 芯片 + TensorCore每个芯片有两个 TPU TensorCore，负责计算。（注意：推理专用的 TPU 只有一个 TensorCore）。两个 TensorCore 共享内存单元：CMEM（128MiB）和 HBM（32GiB）。

在每个 TensorCore 内部，有我们的计算单元和更小的内存缓冲区：

如果你从 NVIDIA GPU 切换过来，可能会发现一些初步的观察点让人困惑：

1. TPU 上的片上内存单元（CMEM、VMEM、SMEM）远大于 GPU 上的 L1、L2 缓存。
2. TPU 上的 HBM 也远小于 GPU 上的 HBM。
3. 负责计算的“核心”数量似乎少得多。

这与 GPU 完全相反，GPU 拥有较小的 L1、L2 缓存（H100 分别为 256KB 和 50MB）、更大的 HBM（H100 为 80GB）以及成千上万个核心。

在我们继续之前，回顾一下，TPU 和 GPU 一样能够实现极高的吞吐量。TPU v5p 每个芯片可达到 500 TFLOPs/秒，而一个包含 8960 个芯片的完整 Pod 可实现约 4.45 ExaFLOPs/秒。最新的 “Ironwood” TPUv7 据称每个 Pod（9216 个芯片）可达 42.5 ExaFLOPS/秒。

要理解 TPU 如何实现这一点，我们需要了解其设计理念。

TPU 设计理念

TPU 通过两大支柱和一个关键假设来实现惊人的吞吐量和能效：脉动阵列 + 流水线、提前编译，以及假设大多数操作都能以适合脉动阵列的方式表达。幸运的是，在我们现代的深度学习时代，矩阵乘法构成了计算的主体，而这正好适合脉动阵列。

TPU 设计选择 #1：脉动阵列 + 流水线

问：什么是脉动阵列？

脉动阵列是一种硬件设计架构，由相互连接的处理单元网格组成。每个处理单元执行一个小型计算（例如乘积累加），并将结果传递给相邻的处理单元。

脉动阵列图这种设计的好处是，一旦数据传入脉动阵列，就无需额外的控制逻辑来决定如何处理数据。此外，在给定足够大的脉动阵列时，除了输入和输出之外，没有内存读写操作。

由于其固定的组织结构，脉动阵列只能处理具有固定数据流模式的操作，但幸运的是，矩阵乘法和卷积完全符合这种模式。

此外，流水线提供了明显的机会来重叠计算和数据移动。下面是一个在 TPU 上执行的流水线逐点运算示意图。

TPU 数据移动可视化流水线逐点操作（来自 “How to Scale Your Model”[4]）### 附：脉动阵列的缺点——稀疏性

你可以看到脉动阵列喜欢稠密矩阵（即几乎每个周期所有处理单元都处于活动状态）。然而，缺点是对于同样大小的稀疏矩阵，性能没有提升：仍然需要执行相同数量的周期，即使对于值为零的元素，处理单元仍然在工作。

如果深度学习社区更倾向于不规则的稀疏性（例如 MoE），那么处理脉动阵列这一系统性限制将变得更加重要。

TPU 设计选择 #2：提前编译 + 减少对缓存的依赖

本节讲解 TPU 如何通过 TPU 与 XLA 编译器的软硬件协同设计，在避免使用缓存的情况下实现高能效。

首先，传统缓存设计用于处理不可预测的内存访问模式。一个应用程序的程序可能与其他应用程序的程序有着截然不同的内存访问模式。本质上，缓存使硬件能够灵活地适应广泛的应用。这也是 GPU 相对于 TPU 而言非常灵活的一个重要原因。

然而，缓存访问（以及一般的内存访问）消耗大量能量。以下是芯片上操作能耗的粗略估计（45nm，0.9V；[18]）。关键结论是，内存访问和控制消耗了大部分能量，而算术运算消耗的能量要少得多。

TPU 数据移动可视化但是，如果你的应用非常特定，并且其计算/内存访问模式高度可预测呢？

举一个极端的例子，如果我们的编译器可以提前计算出所有所需的内存访问，那么我们的硬件就可以只用一个暂存器作为缓冲区，而无需缓存。

这正是 TPU 理念的目标，也是 TPU 与 XLA 编译器协同设计以实现这一点的原因。XLA 编译器通过提前分析计算图来生成优化的程序。

问：但 JAX 也能很好地与 TPU 配合使用，但它们使用 @jit？

JAX 和 XLA 在 TPU 上处于即时编译和提前编译的混合状态，因此造成了混淆。当我们第一次在 JAX 中调用一个 jitted 函数时，JAX 会追踪它来创建一个静态计算图。这个图被传递给 XLA 编译器，在那里被转化为一个完全静态的、面向 TPU 的二进制文件。正是在最后的转换阶段，TPU 特定的优化（例如最小化内存访问）得以执行，以使其适应 TPU。

但有一点需要注意：如果使用不同的输入形状运行 jitted 函数，它们必须被重新编译和缓存。这就是为什么 JAX 在处理任何动态填充或依赖于输入的、具有不同长度 for 循环的层时表现不佳。

当然，这种方法听起来很好，但也存在不便的缺点。缺乏灵活性，并且对编译器的高度依赖是一把双刃剑。

但 Google 为什么仍然坚持这种设计理念呢？

TPU 与能效（TPUv4）

之前展示的能耗图并不能准确代表 TPU 的情况，因此以下是 TPUv4 的能耗分解。请注意，TPUv4 采用 7nm 工艺，而 45nm 数据用于对比（[3], [16]）。

左侧的条形图是数值的可视化，但需要注意的一点是，现代芯片使用 HBM3，其能耗远低于此处显示的 DDR3/4 DRAM。尽管如此，这显示了内存操作消耗的能量高出几个数量级。

这与现代缩放定律有很好的联系：我们非常乐意增加 FLOPS 以换取减少内存操作。因此，减少内存操作具有双重优化优势，因为它们不仅能使程序更快，而且消耗的能量也少得多。

TPU 多芯片层面

让我们往上走，看看 TPU 在多芯片环境下如何工作。

托盘级别（又称“板卡”；4 个芯片）

TPU 数据移动可视化一个 TPU 托盘由 4 个 TPU 芯片或 8 个 TensorCore（简称“核心”）组成。每个托盘都有自己的 CPU 主机（注意：对于推理 TPU，一个主机访问两个托盘，因为每个芯片只有一个核心）。

主机与芯片之间的连接是 PCIe，但芯片与芯片之间的连接是核间互连（ICI），其带宽更高。

但 ICI 连接可以扩展到多个托盘。为此，我们需要上升到机架级别。

机架级别（4x4x4 芯片）

TPU 特别令人兴奋的一点在于其可扩展性，我们从机架级别开始就能看到这一点。

一个 TPU 机架由 64 个 TPU 组成，它们以 4x4x4 的 3D 环面方式连接。如果你看过 Google 的 TPU 宣传材料，比如下图，这是 8 个 TPU 机架的图像。

TPU 数据移动可视化8 个 TPU 机架（TPUv4）但在深入探讨机架之前，我们需要澄清一些容易混淆的术语：机架 vs Pod vs 切片。

问：“TPU 机架”、“TPU Pod”和“TPU 切片”有什么区别？

不同的 Google 资料使用方式略有不同，有时会将“TPU Pod”和“TPU 切片”互换使用。但在本文中，我们将沿用 Google 的 TPU 论文和 GCP TPU 文档中的定义（[3], [7], [9]）。

1. TPU 机架：包含 64 个芯片的物理单元。也称为“立方体”。
2. TPU Pod：通过 ICI 和光纤连接的最大 TPU 单元。也称为“Superpod”或“完整 Pod”。例如，TPUv4 的一个 Pod 由 4096 个芯片或 64 个 TPU 机架组成。
3. TPU 切片：介于 4 个芯片和 Superpod 大小之间的任何 TPU 配置。

关键区别在于，TPU 机架和 TPU Pod 是物理度量单位，而 TPU 切片是一个抽象度量单位。当然，设置 TPU 切片有重要的物理属性，但我们暂时将其抽象化。

现在，我们将使用物理度量单位：TPU 机架和 TPU Pod。这是因为了解 TPU 系统的物理连接方式有助于我们更好地理解 TPU 的设计理念。

现在回到 TPU 机架（TPUv4）：

一个 TPU 机架由 64 个芯片组成，它们通过 ICI 和光电路交换（OCS）连接在一起。本质上，我们将多个托盘连接在一起，模拟出一个 64 芯片的系统。这种将较小部分组合成超级计算机的主题将在后面继续出现。

下面是一个 TPUv4 机架的示意图。它是一个 4x4x4 的 3D 环面，每个节点是一个芯片，蓝色的箭头是 ICI，而面上的线是 OCS（根据 [7] 重绘）。

简单线程层次结构单个 TPU 机架与 OCS然而，这个示意图引发了一些问题。为什么 OCS 只用在面上？ 换句话说，使用 OCS 有什么好处？ 有三个主要好处，另外两个我们稍后会讲到。

OCS 的好处 #1：环绕

通过环绕实现节点间更快的通信

OCS 也充当了给定 TPU 配置的环绕连接。这将两个节点之间的最坏情况跳数从每轴 N-1 跳减少到 (N-1)/2 跳，因为每个轴都变成了一个环（1D 环面）。

随着我们进一步扩展规模，这种效果变得更加重要，因为减少芯片间通信延迟对于高并行化至关重要。

附：并非所有 TPU 都采用 3D 环面拓扑

注意：较老代的 TPU（例如 TPUv2、v3）和推理 TPU（例如 TPUv5e、TPUv6e）采用 2D 环面拓扑，而不是像下面这样的 3D 环面。然而，TPUv7 “Ironwood” 虽然被宣传为推理芯片，但似乎是 3D 环面（注意：我只是根据其宣传材料推测）。

TPU 数据移动可视化2D 环面拓扑## 完整 Pod 级别（又称“Superpod”；TPUv4 为 4096 个芯片）

就像我们将多个芯片连接在一起组成一个 TPU 机架一样，我们可以将多个机架连接在一起组成一个大型 Superpod。

Superpod 也是指 TPU 能够达到的最大互连芯片配置（仅使用 ICI 和 OCS）。接下来还有多 Pod 级别，但这需要通过较慢的互连，我们稍后会讲到这一点。

这根据世代的不同而变化，但 TPUv4 是 4096 个芯片（即 64 个 4x4x4 芯片的机架）。对于最新的 TPUv7 “Ironwood”，这是 9216 个芯片。

下图显示了 TPUv4 的一个 Superpod。

TPU 数据移动可视化TPUv4 的 Superpod（64 个机架）注意每个立方体（即一个 TPU 机架）如何通过 OCS 相互连接。这也允许我们在 Pod 中取出 TPU 的切片。

使用 OCS 的 TPU 切片

我们可以请求 Pod 内的 TPU 子集，这些就是 TPU 切片。但是，即使你想要 N 个芯片，也有多种拓扑结构可供选择。

例如，假设你想要总共 512 个芯片。你可以请求一个立方体（8x8x8）、一个雪茄形状（4x4x32）或一个矩形（4x8x16）。选择切片的拓扑结构本身就是一个超参数。

你所选择的拓扑结构将影响节点间的通信带宽。这直接影响到不同并行方法的性能。

例如，对于全收集通信（如数据并行或张量并行），立方体（如 8x8x8）是首选，因为它具有最高的对分带宽。然而，对于流水线并行，雪茄形状（如 4x4x32）更好，因为它可以更快地与顺序层通信（假设一层适合 4x4 芯片的子切片）。

TPU 数据移动可视化示例 TPU 拓扑当然，最佳拓扑结构取决于模型，找到它本身也是一项工作。TPUv4 论文 [9] 也对此进行了测量，以展示拓扑变化如何加速吞吐量（注意：我不确定第一行指的是哪种 LLM 架构，因为未明确说明）。

TPU 数据移动可视化不同拓扑带来的吞吐量提升我们已经介绍了 TPU 切片，但还有一个重要的特性有助于 TPU 实现高运行稳定性。

得益于 OCS，这些切片不必是连续的机架。这是我们之前没有讲到的 OCS 的第二个好处——也可能是最大的好处。

OCS 的好处 #2：（可重构的）非连续多节点切片

注意，这与将多个节点硬连接在一起以模拟非连续切片不同。由于 OCS 是一个交换机，而不是硬连接，因此跨节点的物理线缆更少，从而允许更高的可扩展性（即更大的 TPU Pod 大小）。

这使得在大规模情况下可以实现灵活的节点配置。例如，假设我们想在单个 Pod 上运行三个作业。虽然简单的调度无法实现这一点，但 OCS 连接允许我们抽象掉节点的位置，将整个 Pod 视为一个“节点包”（根据 [6] 重绘）。