系统架构

张量处理单元 (TPU) 是 Google 设计的机器学习加速器。Cloud TPU 使 TPU 可用作可扩缩的 GPC 云资源。您可以使用 TensorFlowPytorchJAX 等机器学习框架在 Cloud TPU 上运行机器学习工作负载。

一个 TPU 设备包含 4 个芯片,每个芯片均包含 2 个 TPU 核心。TPU 核心包含一个或多个矩阵乘法单元 (MXU)、矢量处理单元 (VPU) 和标量单元。

MXU 由脉动阵列中的 128 x 128 乘法/累加器组成。MXU 可在 TPU 芯片中提供大部分计算能力。每个 MXU 能够在每个周期中使用 bfloat16 数字格式执行 16K 乘法累加运算。

VPU 用于激活、softmax 等常规计算。标量单元用于控制流、计算内存地址和其他维护操作。

TPU 版本

TPU 设备的确切布局取决于您使用的 TPU 版本。如需了解 TPU v2 和 v3 的架构详情和性能特征,请参阅用于训练深度神经网络的领域专用超级计算机

TPU v2

TPU v2 板包含四个 TPU 芯片和 16 GiB HBM。每个 TPU 芯片均包含两个核心。每个核心均包含 MXU、矢量单元和标量单元。

映像

TPU v3

TPU v3 板包含四个 TPU 芯片和 32 GiB HBM。每个 TPU 芯片均包含两个核心。每个核心均包含 MXU、矢量单元和标量单元。

映像

Cloud TPU 提供以下 TPU 配置:

  • 单个 TPU 设备
  • TPU Pod - 通过高速互连连接的一组 TPU 设备
  • TPU 切片 - TPU Pod 的细分

TPU v3 相对于 v2 的性能优势

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 配置

TPU 支持以下配置:

  1. 单个 TPU 板
  2. TPU Pod
  3. TPU Pod 切片

单个 TPU 板

单板 TPU 配置是一个独立的板,包含 4 个 TPU 芯片(8 个 TPU 核心),与其他 TPU 板没有网络连接。单板 TPU 不属于 TPU Pod 配置,也不占用 TPU Pod。

TPU Pod 和切片

在 TPU Pod 或 TPU Pod 切片中,TPU 芯片使用高速互连进行连接,每个 TPU 芯片直接与 TPU 设备上的其他芯片通信。TPU 运行时自动将数据分布到 Pod 或切片中的每个 TPU 核心。Pod 切片可以具有 32、128、512、1024 或 2048 个核心。

Cloud TPU 虚拟机架构

TPU 只能执行矩阵操作,因此每个 TPU 板都会连接到基于 CPU 的主机,以执行无法在 TPU 上执行的操作。主机负责从 Cloud Storage 加载数据、预处理数据并将数据发送到 TPU。

映像

在 TPU Pod 中,每个 TPU 板都有一个 TPU 主机。

图片

与 TPU 主机(和 TPU 板)的交互方式取决于您使用的 TPU 虚拟机架构:TPU 节点或 TPU 虚拟机。

TPU 节点

TPU 节点是原始的 TPU 体验。它们需要一个额外的用户虚拟机,该虚拟机通过 gRPC 与 TPU 主机进行通信;不能直接访问 TPU 主机。

映像

TPU 虚拟机

当您使用 TPU 虚拟机时,您可以通过 SSH 直接连接到以物理方式连接到 TPU 设备的 Google Compute Engine 虚拟机。您可以获取对该虚拟机的根访问权限,以便运行任意代码。您可以访问编译器和运行时调试日志以及错误消息。

映像

支持的框架

JAX、PyTorch 和 TensorFlow 等框架通过存在于每个 TPU 上的名为 libtpu 的共享库访问 TPU。此库包含用于编译 TPU 程序的 XLA 编译器、用于运行已编译程序的 TPU 运行时,以及该运行时用于对 TPU 的低级层访问的 TPU 驱动程序。

图片

TensorFlow

TPU 虚拟机

使用 TPU 虚拟机时,Python 代码可以直接在 TPU 主机上运行,而不是在用户虚拟机上运行。

图片

如需详细了解 TensorFlow 和 Cloud TPU,请参阅在 Cloud TPU 上运行 TensorFlow 模型

TPU 节点

Cloud TPU 节点系统架构最初是针对 TensorFlow 构建的。用户无法访问 TPU 主机,并且 TPU 主机运行 TensorFlow 服务器的无头副本。TPU 主机不会运行 Python 或任何未表示为 TensorFlow 图的用户代码。用户代码在单独的远程虚拟机中运行,该虚拟机通过 gRPC 网络与 TPU 主机进行通信。

映像

PyTorch

TPU 虚拟机

使用 TPU 虚拟机时,PyTorch 代码直接在 TPU 主机上运行。

图片

如需详细了解 PyTorch 和 Cloud TPU,请参阅在 Cloud TPU 上运行 PyTorch 模型

TPU 节点

PyTorch 使用名为 XRT 的库在 Cloud TPU 节点架构上运行,该库允许通过 TensorFlow gRPC 连接发送 XLA 图和运行时指令以及在 TensorFlow 服务器上执行这些指令。每个 TPU 主机都需要一个用户虚拟机。

图片

JAX

TPU 虚拟机

使用 TPU 虚拟机时不需要用户虚拟机,因为您可以直接在 TPU 主机上运行代码。

图片

如需详细了解如何在 Cloud TPU 上运行 JAX,请参阅 JAX 快速入门

TPU 节点

JAX 在 Cloud TPU 节点上的运行方式与 PyTorch 类似,每个主机虚拟机都需要一个单独的用户虚拟机。

图片

后续步骤