系统架构

张量处理单元 (TPU) 是 Google 设计的机器学习加速器。Cloud TPU 将 TPU 作为一种可伸缩的 GCP 云资源提供。您可以使用机器学习框架(如 TensorFlowPytorchJAX)在 Cloud TPU 上运行机器学习工作负载。

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

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

TPU 版本

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

TPU 版本 4

最小的 TPU v4 配置包含 4 个 TPU 芯片、32 GiB 的 HBM 和 128 MiB 的共享公共内存。每个 TPU 芯片包含两个核心。每个核心有四个 MXU,一个矢量单元和一个标量单位。下图展示了 TPU v4 芯片。

映像

TPU v3

最小的 TPU v3 配置包含 4 个 TPU 芯片和 32 GiB 的 HBM。每个 TPU 芯片包含两个核心。每个核心有两个 MXU,一个矢量单元和一个标量单位。下图展示了 TPU v3 芯片。

映像

TPU v2

最小的 TPU v2 配置包含 4 个 TPU 芯片和 16 GiB 的 HBM。每个 TPU 芯片包含两个核心。每个核心都有一个 MXU、矢量单位和标量单位。下图展示了 TPU v2 芯片。

映像

Cloud TPU 提供以下 TPU 配置:

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

与 v3 相比,TPU v4 性能优势

内存系统:

  • v4 TPU 芯片在整个芯片上具有统一的 32 GiB HBM 内存空间,能够更好地实现两个芯片上 TPU 核心之间的协调。
  • 使用最新的内存标准和速度提高了 HBM 性能。
  • 改进了 DMA 性能配置文件,以 5120 亿粒度的精细步长原生支持。

TPU 核心:

  • 提供 275 个 TFLOPS 上限的 2 倍,同时时钟频率更高。
  • 2 倍的转置和排列带宽。
  • 通用内存 (Cmem) 的加载存储内存访问模型。
  • 更快的 MXU 权重加载带宽和 8 位模式支持,以减少批次大小并缩短推断延迟时间。

芯片间互连:

每个芯片提供 6 个互连链接,以实现网络直径较小的网络拓扑。

其他:

  • 用于主机的 x16 PCIE gen3 接口(直接连接)。
  • 改进了安全模型。
  • 提高了能源效率。

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

TPU v4 Pod 由 4096 个芯片组成,通过可配置的高速链路互连。TPU v4' 灵活的网络支持以多种方式连接相同大小的切片组件。您可以使用 AcceleratorConfig 来配置 v4 TPU。AcceleratorConfig 可让您通过拓扑字段来指定芯片大小的切片大小。

使用 3 元组指定 TPU v4 拓扑,该元组描述 TPU 芯片之间的数字、形状和互连。以下插图显示了一些常见的 TPU v4 拓扑。

映像

更大的切片可以从一个或多个 4x4x4 立方体块构建。

可以按不同方式配置给定数量的芯片的 TPU 切片。例如,带有 AcceleratorType 为 v4-1024 的 TPU 切片现在可以配置为:4x4x32、4x8x16、8x8x8。v4-4096 的 TPU 更加灵活:4x4x128、4x8x64、4x16x32、8x16x16、8x16x16_twisted。如需了解详情,请参阅 TPU 类型和拓扑

TPU v2 和 v3 配置

TPU 支持以下配置:

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

单个 TPU 板

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

TPU Pod 和 Slice

在 TPU Pod 或 TPU Pod 切片中,TPU 芯片使用高速互连进行连接,每个 TPU 芯片直接与 TPU 设备上的其他芯片进行通信。TPU 软件自动处理将数据分发到 Pod 或 Slice 中的每个 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 节点架构由通过 gRPC 与 TPU 主机通信的用户虚拟机组成。使用此架构时,您无法直接访问 TPU 主机。这可能会使调试训练和 TPU 错误变得困难。

映像

TPU 虚拟机

TPU 虚拟机架构不再需要用户虚拟机,您可以直接通过 SSH 连接到以物理方式连接到 TPU 设备的虚拟机。您对虚拟机具有 root 访问权限,因此可运行任意代码。您可以访问编译器和运行时调试日志以及错误消息。

映像

支持的框架

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 主机,并运行 TensorFlow 服务器的无头副本。它们不运行 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 类似,每个主机虚拟机都需要一个单独的用户虚拟机。

映像

后续步骤