深度解读谷歌第四代TPU

似乎我们一直在谈论 Google 的 TPUv4 机器学习加速器，但也仅仅是谈论。今天，在 Google Cloud Next 2022 大会上，这家搜索引擎、广告、视频流、企业应用程序和云计算巨头最终将向大众推出其第四代——也是最强大和可扩展的张量处理单元——Tensor Processing Units。

几年来，谷歌一直在泄露有关 TPUv4 架构的信息，但并没有非常具体地说明系统的馈送和速度。但我们与 TPU 背后的一些人进行了交谈，还研究了论文和技术演示，将它们汇集在一起，对四代 TPU 计算引擎和使用它们的系统进行了概述。在我们无法找到信息的地方，我们按照我们的习惯进行了我们认为是明智的猜测。

也许重要的是，我们要提醒大家为什么谷歌首先决定为人工智能训练和推理做定制芯片。早在 2013 年，当谷歌正在为其同名搜索引擎开发语音激活搜索功能时，谷歌研究员Jeff Dean参与了许多关键技术，很难在此一一列举。在一张废纸上计算数学，并发现如果这种人工智能辅助搜索投入生产，那么谷歌将不得不将其数据中心的足迹增加一倍。当然，这对于“免费”的产品是不可行的。支持这种扩张的搜索广告商不会突然增加 2 倍。因此，谷歌开始努力为人工智能训练和推理构建一个更好的数学引擎。

事不宜迟，以下是四代 TPU 的相关进料、速度、规模和成本，包括一个 TPUv4i 子变体，直到我们看到Norman Jouppi 的这篇论文和一个许多其他人（包括谷歌研究员David Patterson)）。

TPU 是实际领域特定架构的一个例子，我们认为谷歌追随 Nvidia 的脚步很重要，因为它创建了一个可以进行训练和推理的通用产品，同时它还创建了一个专门针对推理进行调整的子变体——在 TPUv4i 的情况下，针对其自己的内部推理工作负载进行了调整，而不是那些想要在 Google Cloud 上运行推理作业的人，至少现在还没有。

事实证明，这款 TPUv4i 推理芯片在 TPUv4 通用引擎之前投入生产，这与 Nvidia 的做法大相径庭，这也表明谷歌真的非常需要廉价的推理来驱动十几个大约是利用其 TPU 添加 AI 处理以增强其功能的应用程序。谷歌在实际将使用完整 TPUv4 的系统投入生产时并不准确，我们猜测是在 2021 年第四季度左右。

谷歌在其 TPU 设计上并不羞于落后一两个节点，这绝对是为了降低芯片设计和生产成本。领先计算引擎的大部分成本来自领先工艺的不良良率，虽然您可以通过在设计中添加备用组件来稍微降低成本，但这也是有代价的。通过这种选择，芯片在许多载体上更便宜，也更容易通过供应链进入谷歌手中。

TPUv4i 推理芯片采用台积电的 7 纳米工艺制造，并在一年半前投入生产，我们认为全开的 TPUv4 也使用相同的 7 纳米工艺。我们在之前的报道中一直在讨论这个问题，这就是我们的最终答案。我们认为 TPUv5i 和 TPUv5 计算引擎将使用5 纳米工艺，如果 Google 有 TPUv4i 和 TPUv4 在生产中并且 TPUv4 在谷歌云上普遍可供客户使用，几乎可以肯定它们正在测试中。

TPU 内核的基本结构是它有一个标量和矢量处理器前端，带有一个本地的矩阵数学单元，可以执行 BF16（Google 为其 TPU 发明的格式）乘法和 FP32 累加操作（有时使用 INT8 进行推理，具体取决于在模型上）。

最初的 TPUv1 处理器有一个 256×256 的矩阵数学单元，这个单元非常庞大，而且正如 Google 在此过程中了解到的那样，它并不是最佳配置。而且，事实证明，谷歌自己的人工智能应用程序也不是 64×64 阵列。

矩阵吞吐量和利用率之间的最佳点是至少对于谷歌自己的代码来说是一个 128×128 阵列，一旦谷歌发现了这一点，该公司就通过添加 TPU 内核（其中嵌入了 MXU 单元）来增加 TPU 的规模。它们或将 MXU 单元添加到每个内核，或同时执行这两项操作。

以下是从 TPUv2 到 TPUv4 设计的演变，TPUv4 具有 128×128 阵列 MXU：

下面是使用每一代 TPU 的服务器主板的外观：

在 TPUv4 服务器主板的右侧，您可以看到一些相邻的芯片。这是一个六端口 6 Tb/sec 交换机，用作网络接口卡和 3D 环面网络的基础。

当我们谈论 TPUv4 芯片时，我们不确定 Google 是否会添加内核或扩展 MXU，事实证明，它已将内核数保持在两个，并且再次将每个内核的 MXU 数量翻了一番。除非有什么改变，否则这就是谷歌创建 TPUv5、TPUv6 等的方式。我们认为，在某个时候，谷歌将使用四个内核，这取决于其自主研发的 XLA 机器学习编译器可以做什么，但目前，提供较少数量的内核是扩展应用程序的最简单方法。

您会注意到 TPUv4i 推理芯片本身就相当庞大，它有一个带有四个 128×128 MXU 的 TPU 内核，运行在非常可观的 1.05 GHz 下，但仅占用 400 mm 2的面积。据我们所知，根据一些平面图估计，将单核上的 MXU 数量增加一倍，缓存增加一倍并使用其他一些非核心功能来创建具有单核的完整 TPUv4 可能只会有提升die尺寸约为563 mm 2左右。

但无论出于何种原因——可能与 XLA 编译器中的双核 TPU 优化有关——谷歌做了一个双核 TPUv4 AI 训练引擎，我们认为它的面积约为 780 mm 2。从我们所见，这款两核 TPUv4 芯片将自己呈现为一个逻辑核心，具有一个 32 GB HBM 内存空间。我们怀疑它是一种单片设计，但它可能由两个 TPUv4 芯片在一个插槽中捆绑在一起组成，这会降低成本，但可能会影响性能，具体取决于 Google 可能创建或许可的插槽内互连。我们高度怀疑它是一种单片设计，其逻辑块总是要加倍。

让我们暂时讨论一下过去三代核心 TPU 中所有这些不同向量的尺度。TPUv3 是对 TPUv2 的增量升级，在许多馈送和速度方面提高了 30%，在相同工艺下晶体管数量仅增加了 11%，裸片尺寸增加了 12% （16 纳米），芯片性能提高 2.67 倍，HBM 主存储器容量提高 2 倍（让它处理更大的数据集）。在我们看来，与 TPUv3 最大的不同在于，2D torus 互连从 TPUv2 芯片中的 256 个芯片扩展到 TPUv3 一代中的 1,024 个芯片，这让Pod 处理能力增加了 10.7 倍，从12 petaflops 到 126 petaflops（在 BF16 操作中测量）。

TPUv4 是计算引擎的真正升级，工艺从 16 纳米缩小到 7 纳米，并且具有所有允许的优点。MXU 的数量再次翻了一番，缓存内存增加了 9 倍至 244 MB（我们认为），HBM2 内存带宽增加了 33% 至 1.2 TB/秒，但有趣的是 HBM2 内存容量保持不变，为 32 GB。

Google 能做到这一点的原因是它可以像发疯一样扩展 TPUv4 pod 。英伟达吹嘘它如何将多达 256 个“Hopper”GH100 GPU 加速器与其 NVSwitch 结构紧密耦合，但随着 TPUv4 首次亮相的新 3D torus互联，它具有更多的带宽和更高的基数，它可以紧密耦合4,096 个 TPUv4 引擎，总计 1.126 exaflops 的 BF16 计算。（其中 8 个 TPUv4 pod 位于 Google 的俄克拉荷马州梅斯县数据中心，总共产生 9 exaflops 的原始 AI 计算。）

实际上，在这 4,096 个 TPUv4 芯片上同时处理的 AI 工作负载并不多，这就是 3D  torus有用的原因。它允许将机器切成紧密耦合的块，这些块仍然相互连接，因此它们可以共享工作。

“在很多情况下，工程师使用 64 个甚至更少的芯片，”谷歌云 TPU 产品的对外产品经理 Vaibhav Singh 告诉The Next Platform。“使用 TPUv4 完成的一个有趣的事情是多维数据集级别的环绕式连接，其中多维数据集的单位是 4 x 4 x 4，因为就最常见的工作负载大小而言，这是我们的最佳选择。我们可以将更大的切片构建为这四乘四乘四的立方体的组合。”

在上表中，立方体切片大小是计算核心数，而不是芯片数，因此在解释时要小心。这些切片大小是动态配置的。

“当我们谈论不适合单个芯片并且必须分布在数百个芯片中的模型时，对物理拓扑的认识，以及您的软件实际上尝试在多大程度上进行特定优化计算和通信的重叠对系统的整体性能产生了巨大的影响，”Singh 补充道。

HPC 人群非常清楚 3D 和更高维环面互连的好处——IBM 的 BlueGene 超级计算机有一个 3D 环面，而在 K 和 Fukagu 超级计算机中使用的富士通 Tofu 互连是 6D 环面/网状互连——即使它们在如果添加机器，则重新布线。（因此为需要升级和重新布线的系统选择 Clos 或蜻蜓拓扑。）

TPUv4i 和 TPUv4，尤其是针对其前身以及针对 AI 训练和推理的竞争 GPU，展示了仅做一件事的极简设计的价值：用于 AI 处理的矩阵数学。（尽管我们知道有些研究人员正试图利用 TPU 来运行加速 HPC 仿真和建模应用程序。）

由于 Hopper GPU 现在才变得普遍可用，Google 使用 TPUv4 引擎与其他加速器进行的比较必然要与 Nvidia 的“Ampere”A100 GPU 进行比较。下面的图表由 Google 系统和服务基础设施副总裁 Amin Vahdat 提供，展示了 TPU 如何与 A100 叠加，以及 Google 云上的 TPUv4 pod 如何与 Microsoft Azure 上类似大小的 A100 集群叠加。

这些是 MLPerf 2.0 测试套件的研究提交基准，这是测试的hot rod, nitro-guzzling 版本。这就是人工智能最前沿的地方。以下是 TPUv4 和 A100 在不同计算引擎规模（如条形图所示）下针对各种神经网络模型的直接比较：

在较小规模的模型上，性能差异并不像您在 Meta Platforms 的深度学习推荐模型 (DLRM) 中看到的那么大。很多时候，性能提升似乎是由于代码能够在更多iron上进行扩展（如 MaskRCNN 和较小程度的 RetinaNet），有时是由于在大量iron上更好地扩展（如BERT 和 ResNet）。Vahdat 表示，平均而言，在五个基准测试中，TPUv4 在 MLPerf 测试中的性能比 Nvidia A100 高 40%，有趣的是，谷歌在这五个测试中提交的 MLPerf 2.0 结果比相同的 TPUv4 铁测试高 50%使用 MLPerf 1.0 测试套件。

“这里的重点是，实际上这些性能提升反映了软件工作，以及我们能够在软件上进行的大量优化，以利用我们的硬件功能。相对于我们上次录取的 50% 尤为关键，因为这意味着我们所有的 Cloud TPU 用户基本上都可以免费获得相同的性能优势。”

最终，比硬件馈送和速度更重要的是 AI 训练成本。这一次，谷歌向我们分析了它的 TPUv4 pod 将如何使用相同的 4,096 个 Nvidia A100 GPU 堆叠到类似大小的机器集群：

Vahdat 的团队根据提交的 MLPerf 结果和自己对 BERT 和 ResNet 如何在 GPU 上扩展的知识对这一成本进行了建模——Google 销售这些 GPU，并且 Google 仍在将其用于大量自己的 AI 工作负载。目前尚不清楚与 H100 GPU 加速器的比较会是什么样子。Nvidia 可能会以 2.5 到 3 倍的成本提供 3 到 6 倍的性能（取决于工作负载）。（这就是我们的样子。）因此，Nvidia 可能会稍微缩小规模和性能差距，但不会像您想象的那样改变价格/性能等式。

我们还发现有趣的是，TPUv4 容量的定价正是我们在 5 月份回信时所预期的。TPUv4 芯片的一小时时间成本比 TPUv3 芯片高出 61%，但性能却提高了 2.23 倍。因此，性价比提高了 28%。随着 TPUv2 到 TPUv3 的跃升，性能提升了 2.67%，性价比提升了 33.5%。这与谷歌希望看到的物有所值的改进以及客户在 IT 领域的世代期望是一致的。