与 Google Cloud 博客 由 Tao Wang，软件工程师，Aarush Selvan，产品经理交叉发布。

最新一轮的 MLPerf 基准测试结果已经发布，Google 的 TPU v4 超级计算机在规模上展示了打破纪录的性能。这是一个及时的里程碑，因为大规模机器学习训练使近年来 AI 的许多突破成为可能，最新的模型包含数十亿甚至数万亿个参数（T5，Meena，GShard，Switch Transformer 和 GPT-3）。

Google 的 TPU v4 Pod 部分是为了满足这些扩展的训练需求而设计的，TPU v4 Pod 在 Google 使用 TensorFlow 和 JAX 参加的六个 MLPerf 基准测试中的四个基准测试中创下了性能记录。这些分数比我们去年提交的获奖作品有了显著提高，表明 Google 再次拥有世界上最快的机器学习超级计算机。这些 TPU v4 Pod 已经在 Google 数据中心中广泛部署，用于我们的内部机器学习工作负载，并将在今年晚些时候通过 Google Cloud 提供。

图 1：Google 最佳 MLPerf Training v1.0 TPU v4 提交与任何可用性类别中最快的非 Google 提交的加速比 - 在这种情况下，所有基线提交都来自 NVIDIA。比较按总训练时间标准化，无论系统规模如何。更高的柱状图表示更好。1

让我们仔细看看实现这些突破性结果的一些创新，以及这对 Google 及更广泛范围内的模型训练意味着什么。

Google 持续的性能领先地位

Google 最近提交的 MLPerf 展示了领先的顶级性能（达到目标质量所需的最短时间），在四个基准测试中创下了新的性能记录。我们通过将我们的下一代 TPU v4 ASIC 扩展到 3,456 个，并为多个基准测试提供数百个 CPU 主机来实现这一点。与 去年的结果 相比，我们在顶级提交中平均提高了 1.7 倍。这意味着我们现在可以在几秒钟内训练一些最常见的机器学习模型。

图 2：Google 的 MLPerf Training v1.0 TPU v4 提交与 Google 的 MLPerf Training v0.7 TPU v3 提交的加速比（例外：MLPerf v0.7 中的 DLRM 结果是使用 TPU v4 获得的）。比较按总训练时间标准化，无论系统规模如何。更高的柱状图表示更好。Unet3D 未显示，因为它是 MLPerf v1.0 的新基准。2

我们通过持续投资于硬件和软件堆栈来实现这些性能改进。部分加速来自使用 Google 的第四代 TPU ASIC，与上一代 TPU v3 相比，它提供了原始处理能力的显著提升。4,096 个这种 TPU v4 芯片联网在一起，形成一个 TPU v4 Pod，每个 Pod 的峰值性能为 1.1 exaflop/s。

图 3：1 exaflop/s 计算能力的视觉表示。如果 1000 万台笔记本电脑同时运行，那么所有这些计算能力几乎与 1 exaflop/s 的计算能力相匹配。

与此同时，我们在 XLA 编译器 中引入了一些新功能，以提高在 TPU v4 上运行的任何 ML 模型的性能。其中一项功能提供了将两个（或可能更多）TPU 核心作为单个逻辑设备运行的能力，使用共享统一内存访问系统。这种内存空间统一使核心能够轻松共享输入和输出数据 - 允许在核心之间更有效地分配工作。第二个功能通过精细粒度的计算和通信重叠来提高性能。最后，我们引入了一种技术，可以自动转换卷积运算，使空间维度转换为额外的批处理维度。这种技术提高了在非常大规模下常见的低批量大小下的性能。

使用无碳能源支持大型模型研究

尽管顶级 MLPerf 基准测试中差异的幅度可以用几秒钟来衡量，但这可以转化为在包含数十亿甚至数万亿个参数的最先进模型上进行数十天的训练时间。举个例子，今天我们可以用 2048 个 TPU 核心训练一个 4 万亿参数的 带有 GSPMD 的密集 Transformer。作为参考，这比 OpenAI 去年发布的 GPT-3 模型大 20 倍以上。我们已经在 Google 内部广泛使用 TPU v4 Pod 来开发研究突破，例如 MUM 和 LaMDA，以及改进我们的核心产品，例如搜索、助理和翻译。来自 TPU 的更快的训练时间带来了效率节省以及改进的研究和开发速度。许多这些 TPU v4 Pod 将在 90% 无碳能源 或接近该比例的水平运行。此外，云数据中心可以比典型数据中心 高出 ~1.4-2 倍的能源效率，而运行其中的面向 ML 的加速器（如 TPU）可以比现成的系统 高出 ~2-5 倍的效率。

我们也很快将在 Google Cloud 上提供 TPU v4 Pod，使世界上最快的机器学习训练超级计算机可供全球客户使用，并且我们 最近发布了 一个全新的 Cloud TPU 系统架构，它提供对 TPU 主机机器的直接访问，极大地改善了用户体验。

想了解更多信息？

阅读 如何开始使用 TPU 训练您的模型。我们很高兴看到您将如何通过使用 exaflop 级别的 TPU 计算能力来扩展机器学习的边界！

¹ 所有结果均从 www.mlperf.org 获取，截至 2021 年 6 月 30 日。MLPerf 的名称和徽标为商标。有关更多信息，请访问 www.mlperf.org。图表使用结果 1.0-1067、1.0-1070、1.0-1071、1.0-1072、1.0-1073、1.0-1074、1.0-1075、1.0-1076、1.0-1077、1.0-1088、1.0-1089、1.0-1090、1.0-1091、1.0-1092。

² 所有结果均从 www.mlperf.org 获取，截至 2021 年 6 月 30 日。MLPerf 的名称和徽标为商标。有关更多信息，请访问 www.mlperf.org。图表使用结果 0.7-65、0.7-66、0.7-67、1.0-1088、1.0-1090、1.0-1091、1.0-1092。