Cloud TPU 简介

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

TPU 使用专为执行机器学习算法中常见的大型矩阵运算而设计的硬件，更高效地训练模型。TPU 具有片上高带宽内存 (HBM)，可让您使用较大的模型和批量大小。TPU 可以连接在名为 slice 的组中，这让您可以在几乎不需要更改代码的情况下对工作负载进行纵向扩容。

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

何时使用 TPU

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

CPU

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

GPU

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

TPU

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

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

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

Google Cloud中的 TPU

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

Google Cloud service 资源
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？以下资源可能对您有所帮助：