Google 在最新的 MLPerf 基准测试中展示了领先的性能

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

谷歌的 TPU v4 Pod 的设计部分是为了满足这些广泛的训练需求，TPU v4 Pod 在 Google 使用 TensorFlow 和 JAX 提交的六个 MLPerf 基准测试中，四个创造了性能记录。 这些分数比我们去年的获奖提交有了显着提高，并表明 Google 再次拥有世界上最快的 ML 超级计算机。这些 TPU v4 Pod 已在 Google 数据中心内广泛部署，用于我们的内部 ML 工作负载，并已于2021年年底通过 Google Cloud 提供。

图 1：Google 最好的 MLPerf Training v1.0 TPU v4 提交速度比任何可用性类别中最快的非 Google 提交速度快。在这种情况下，所有基线提交都来自 NVIDIA。 无论系统大小如何，比较都通过整体训练时间进行标准化。高的柱状图示表示更好。1

让我们仔细看看带来这些突破性成果的一些创新，以及这对 Google 及其他地方的大型模型训练意味着什么。

Google 保持性能领先地位

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

图 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 相比，在原始处理能力方面有着显著的提升。 将这些 TPU v4 芯片中的 4,096 个连接在一起，创建出一个 TPU v4 Pod，每个 Pod 提供 1.1 exaflop/s 的峰值性能。

图 3：1 exaflop/s 计算能力的直观表示。 如果 1000 万台笔记本电脑同时运行，那么所有计算能力几乎可以匹配 1 exaflop/s 的计算能力。

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

使用无碳能源实现大型模型研究

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

我们也很期待很快在 Google Cloud 上提供 TPU v4 Pod，让世界各地的客户都可以使用世界上最快的 ML 训练超级计算机。 Cloud TPU 支持 TensorFlow、PyTorch 和 Jax 等领先框架，我们最近发布了全新的 Cloud TPU 系统架构，可直接访问 TPU 主机，极大地提升了用户体验。

想了解更多？

请联系您的 Google Cloud 销售代表，申请提前使用 Cloud TPU v4 Pod。 我们很期待看到您将如何利用 exaflops 的 TPU 计算能力扩展 ML 前沿！