Cloud TPU 简介

张量处理单元 (TPU) 是 Google 定制开发的应用专用集成电路 (ASIC)，用于加速机器学习工作负载。如需详细了解 TPU 硬件，请参阅 TPU 架构。Cloud TPU 是一种 Web 服务，使 TPU 作为 Google Cloud上的可扩缩计算资源使用。

TPU 使用专为执行机器学习算法中常见的大型矩阵运算而设计的硬件，更高效地训练模型。TPU 具有片上高带宽内存 (HBM)，支持使用更大的模型和批次大小。TPU 可以组成切片（称为切片）来扩展工作负载，而无需更改代码。

在 TPU 上运行的代码必须由加速器线性代数 (XLA) 编译器编译。XLA 是一种即时编译器，可获取机器学习框架应用发出的图，并将图的线性代数、损失和梯度组成部分编译到 TPU 机器代码中。程序的其余部分在 TPU 主机上运行。XLA 编译器是 TPU 主机上运行的 TPU 虚拟机映像的一部分。

如需详细了解张量处理单元，请参阅如何看待 TPU。

何时使用 TPU

Cloud TPU 针对特定工作负载进行了优化。在某些情况下，您可能希望在 Compute Engine 实例上使用 GPU 或 CPU 来运行机器学习工作负载。通常，您可以根据以下准则确定哪种硬件最适合您的工作负载：

CPU

  • 需要最高灵活性的快速原型设计
  • 训练时间不长的简单模型
  • 有效批量大小较小的小型模型
  • 包含许多使用 C++ 编写的自定义 TensorFlow 操作的模型
  • 受主机系统可用 I/O 或网络带宽限制的模型

GPU

  • 具有大量自定义 PyTorch/JAX 操作且必须至少部分在 CPU 上运行的模型
  • 具有不能在 Cloud TPU 上使用的 TensorFlow 操作的模型（请参阅可用的 TensorFlow 操作列表）
  • 有效批量大小较大的中到大型模型

TPU

  • 由矩阵计算主导的模型
  • 在主训练循环内没有自定义 PyTorch/JAX 操作的模型
  • 需要训练数周或数月的模型
  • 有效批量大小较大的大型模型
  • 在高级排名和推荐工作负载中常见的具有超大嵌入的模型

Cloud TPU 不适合以下工作负载：

  • 需要频繁分支或包含许多元素级代数运算的线性代数程序
  • 需要高精度算法的工作负载
  • 主训练循环中包含自定义操作的神经网络工作负载

Google Cloud中的 TPU

您可以通过 Cloud TPU 虚拟机、Google Kubernetes Engine 和 Vertex AI 使用 TPU。下表列出了每项 Google Cloud服务的资源。

Google Cloud 服务 资源
Cloud TPU Cloud TPU 虚拟机使用入门
Google Kubernetes Engine

关于 GKE 中的 TPU

使用 TPU 在 GKE 上运行 Ray
Vertex AI

使用 TPU 在 Vertex AI 上训练

在 Vertex AI 上使用 TPU 进行在线预测

模型开发最佳做法

由非矩阵运算（如 add、reshape 或 concatenate）主导计算的程序可能无法实现较高的 MXU 利用率。以下准则可帮助您选择和构建适用于 Cloud TPU 的模型。

布局

XLA 编译器执行代码转换，包括将矩阵乘法平铺成较小的块，以便高效地在矩阵单元 (MXU) 上执行计算。XLA 编译器使用 MXU 硬件的结构（128x128 脉动阵列）和 TPU 内存子系统的设计（该设计更倾向于尺寸为 8 的倍数）来提高平铺效率。因此，某些布局更有利于平铺，而另一些布局则需要在平铺之前调整形状。调整形状的操作在 Cloud TPU 上通常受内存限制。

形状

XLA 编译器及时为第一批编译机器学习图。如果任何后续批次具有不同的形状，则该模型不起作用。（在每次形状发生变化时重新编译图太慢。）因此，任何具有动态形状的张量的模型都不适合 TPU。

内边距

在高性能 Cloud TPU 程序中，密集计算可以平铺成 128x128 个块。如果矩阵计算不能占满整个 MXU，编译器会用零来填充张量。填充有两个不足：

  • 用零填充的张量不能充分利用 TPU 核心。
  • 填充增加了张量所需的片上内存存储，并且在极端情况下可能导致内存不足错误。

虽然填充是由 XLA 编译器在必要时自动执行的，但您可以通过 op_profile 工具确定填充量。您可以通过选择非常适合 TPU 的张量维度来避免填充。

维度

选择合适的张量维度对于使 TPU 硬件（尤其是 MXU）发挥最高性能很有帮助。XLA 编译器会尝试使用批次大小或特征维度来最大限度地使用 MXU。因此，这两个因素之一必须是 128 的倍数。否则，编译器会将其中某个填充到 128。理想情况下，批量大小和特征维度应为 8 的倍数，这样可以使内存子系统发挥最高性能。

Cloud TPU 使用入门

寻求帮助

请与 Cloud TPU 支持团队联系。 如果您有活跃的 Google Cloud 项目，请准备好提供以下信息：

  • 您的 Google Cloud 项目 ID
  • 您的 TPU 名称（如果存在）
  • 您要提供的其他信息

后续步骤

想要详细了解 Cloud TPU？以下资源可能会对您有所帮助：