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

  • 具有必须至少部分在 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 服务 资源
Cloud TPU Cloud TPU 虚拟机使用入门
Google Kubernetes Engine

关于 GKE 中的 TPU

使用 TPU 在 GKE 上运行 Ray
Vertex AI

使用 TPU 在 Vertex AI 上进行训练

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

模型开发最佳做法

由非矩阵操作（如添加、调整形状或串联）主导计算的程序可能无法实现较高的 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 使用入门

  1. 设置 Google Cloud 账号
  2. 激活 Cloud TPU API
  3. 向 Cloud TPU 授予对 Cloud Storage 存储桶的访问权限
  4. 在 TPU 上运行基本计算
  5. 在 TPU 上训练参考模型
  6. 分析模型

请求帮助

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

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

后续步骤

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