系统架构

本文档介绍 Cloud TPU 系统的所有硬件和软件组件的架构。

概览

张量处理单元 (TPU) 是 Google 定制开发的专用集成电路 (ASIC),用于加速机器学习工作负载。Google 结合机器学习领域的丰富经验与领先优势,从零开始打造了这些 TPU。

您可以使用 Cloud TPU 和 TensorFlow 在 Google 的 TPU 加速器硬件上运行您自己的机器学习工作负载。Cloud TPU 旨在实现最高性能与灵活性,帮助研究人员、开发者和企业构建可使用 CPU、GPU 和 TPU 的 TensorFlow 计算集群。高级 Tensorflow API 可让您在 Cloud TPU 硬件上轻松运行复制的模型。

您的 TensorFlow 应用可以从 Google Cloud 上的容器、实例或服务访问 TPU 节点。该应用需要通过您的 VPC 网络连接到 TPU 节点。

TPU 版本

每个 TPU 版本定义 TPU 设备的特定硬件特征。 TPU 版本定义了每个 TPU 核心的架构和高带宽内存 (HBM) 的数量、每个 TPU 设备上核心之间的互连以及可用于设备间通信的网络接口。例如,每个 TPU 版本都具有以下特征:

  • TPU v2:
    • 每个 TPU 核心的 HBM 为 8 GiB
    • 每个 TPU 核心拥有一个 MXU
    • TPU Pod 中拥有最多 512 个 TPU 核心和 4 TiB 总内存
  • TPU v3:
    • 每个 TPU 核心的 HBM 为 16 GiB
    • 每个 TPU 核心拥有两个 MXU
    • TPU Pod 中拥有最多 2048 个 TPU 核心和 32 TiB 总内存

图片

每个 TPU 核心都有标量、矢量和矩阵单位 (MXU)。MXU 可在 TPU 芯片中提供大部分计算能力。每个 MXU 能够在每个周期中执行16K 乘法累加运算,且按降低的 bfloat16 精度执行相乘。Bfloat16 是一种 16 位浮点表示法,可比 IEEE 半精度表示法提供更好的训练和模型准确率。

TPU 设备上的每个核心均可独立执行用户计算(XLA 操作)。高带宽互连可实现 TPU 设备上的芯片间直接通信。在 TPU Pod 配置中,专用高速网络接口将多个 TPU 设备连接在一起,为您的机器学习工作负载提供更多的 TPU 核心和更大的 TPU 内存池。

TPU v3 的性能优势

TPU v3 配置中增加的每个核心的 FLOPS 和内存容量可以通过以下方式提高模型的性能:

  • 对于计算受限的模型,TPU v3 配置可为每个核心提供明显的性能优势。如果采用 TPU v2 配置且内存受限的模型在采用 TPU v3 配置时同样也受内存限制,则可能无法实现同等的性能提升。
  • 如果采用 TPU v2 配置时,内存放不下数据,则 TPU v3 可以提供改进的性能并减少中间值的重计算(重实体化)。
  • TPU v3 配置可以运行批次大小不适合 TPU v2 配置的新模型。例如,TPU v3 可能允许更深的 ResNet 和使用 RetinaNet 的较大图片。

因训练步骤等待输入而在 TPU v2 上几乎成为受限于输入(“馈入”)的模型,在 Cloud TPU v3 中也可能会受限于输入。流水线性能指南可以帮助解决馈入问题。

您可以在不同 TPU 版本上运行基准测试并使用 TensorBoard 工具监控性能来确定是否可以通过 TPU v3 提高模型的性能。

TPU 配置

在 Google 数据中心,TPU 设备均支持以下 TPU v2 和 TPU v3 配置:

  • 单个设备 TPU 是指未通过专用高速网络相互连接的单个 TPU 设备。您无法组合多个单个设备 TPU 类型以协同处理单个工作负载。
  • TPU Pod 是指通过专用高速网络相互相连的 TPU 设备集群。

单个设备 TPU

Google 数据中心内的单个设备 TPU 配置是指一台不与其他 TPU 设备通过专用高速网络连接的 TPU 设备。您的 TPU 节点仅连接到此单个设备。

图片

对于单个设备 TPU,芯片在设备上互连,因此芯片之间的通信无需主机 CPU 或主机网络资源。

创建 TPU 节点时,您需要指定 TPU 类型。例如,您可以指定 v2-8v3-8 来为单个设备配置 TPU 节点。单个设备 TPU 不属于 TPU Pod 配置,也不占用 TPU Pod。如需了解您的 TPU 节点可用的单个设备 TPU 配置,请参阅 TPU 类型页面。

TPU Pod

Google 数据中心中的 TPU pod 配置是通过专用高速网络相互连接的多个 TPU 设备。TPU 节点中的主机在所有 TPU 设备上分配您的机器学习工作负载。

图片

在 TPU Pod 中,TPU 芯片在设备上互连,因此芯片之间的通信无需主机 CPU 或主机网络资源。此外,TPU Pod 中的每个 TPU 设备都通过专用高速网络互相连接,无需主机 CPU 或主机网络资源。

创建 TPU 节点时,请指定该类型的 TPU 是否占用完整的 TPU pod 或该 TPU pod 的较小部分。例如,v2-512 类型的 TPU 占用完整的 v2 TPU Pod,而 v2-128 类型的 TPU 仅占用 v2 TPU Pod 的 1/4。如需了解您的 TPU 节点可用的 TPU Pod 配置,请参阅 TPU 类型页面。

图片

v2 TPU Pod 最多可配置 64 个设备,总共具有 512 个 TPU v2 核心和 4 TiB 的 TPU 内存。

图片

v3 TPU Pod 最多可配置 256 个设备,总共具有 2048 个 TPU v3 核心和 32 TiB 的 TPU 内存。

如需详细了解不同 TPU 版本之间的架构差异,请参阅 TPU 版本部分。

您可以使用 Cloud TPU API 对任意规模的 TPU 节点自动执行 TPU 管理。因此,您可以轻松纵向扩容到大规模计算集群、运行工作负载,并在工作负载完成时缩减这些集群。芯片中内置的硬件支持可为各种深度学习工作负载实现高效的线性性能扩缩。实际上,Cloud TPU 软件堆栈消除了生成、运行和馈送 TPU Cloud 程序的复杂性。

软件架构

当您运行应用时,TensorFlow 会生成计算图并通过 gRPC 将其发送到 TPU 节点。您的工作负载可用的设备数量取决于您为 TPU 节点选择的 TPU 类型。TPU 节点及时编译计算图并将程序二进制文件发送到一个或多个 TPU 设备以供执行。模型的输入通常存储在 Cloud Storage 中。TPU 节点将输入流式传输到一个或多个 TPU 设备以供使用。

以下图表展示 Cloud TPU 软件架构,包括神经网络模型、TPU Estimator 和 TensorFlow 客户端、TensorFlow 服务器和 XLA 编译器。

图片

TPU Estimator

TPU Estimator 是一套基于 Estimator 的高级 API,可简化 Cloud TPU 的构建模型,并使 TPU 发挥最佳性能。如果要编写使用 Cloud TPU 的神经网络模型,您应使用 TPU Estimator API。

TensorFlow 客户端

TPU Estimator 将您的程序转换为 TensorFlow 操作,然后这些操作由 TensorFlow 客户端转换为计算图。TensorFlow 客户端将计算图传送给 TensorFlow 服务器。

TensorFlow 服务器

TensorFlow 服务器在 Cloud TPU 服务器上运行。当 TensorFlow 服务器从 TensorFlow 客户端接收计算图时,它会执行以下操作:

  1. 从 Cloud Storage 加载输入
  2. 将图划分为可以在 Cloud TPU 上运行的部分以及必须在 CPU 上运行的部分
  3. 生成对应于要在 Cloud TPU 上运行的子图的 XLA 操作
  4. 调用 XLA 编译器

XLA 编译器

XLA 是一种即时编译器,它将 TensorFlow 服务器生成的高级优化器 (HLO) 操作当作输入。由 XLA 生成要在 Cloud TPU 上运行的二进制代码,包括从片上存储器到硬件执行设备的数据编排和芯片间通信。 生成的二进制文件通过 Cloud TPU 服务器与 Cloud TPU 之间的 PCIe 连接加载到 Cloud TPU,然后开始执行。

后续步骤