系统架构

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

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

一个 TPU 板包含四个 TPU 芯片。每个芯片包含两个 TPU 核心。TPU 核心上可用的资源因版本而异。每个 TPU 核心都有标量、矢量和矩阵乘法单元 (MXU)。MXU 在 TPU 芯片中提供大部分计算能力。每个 MXU 能够在每个周期中执行16K 乘法累加运算,且按降低的 bfloat16 精度执行相乘。bfloat16 是一种 16 位浮点表示法,可比 IEEE 半精度表示法提供更好的训练和模型准确率。TPU 板上的每个核心均可独立执行用户计算。高带宽互连可实现芯片间直接通信。如需详细了解技术信息,请参阅用于训练深度神经网络的领域专用超级计算机

TPU 的设计使其可扩容为 TPU Pod。TPU Pod 是一个超级计算机,最多可以有 2048 个 TPU 核心,允许您跨多个 TPU 板分配处理负载。在 TPU Pod 配置中,专用高速网络接口将多个 TPU 设备连接在一起,为您的机器学习工作负载提供更多的 TPU 核心和更大的 TPU 内存池。

每个 TPU 板都连接到一个基于 CPU 的高性能主机,用于执行加载和预处理要提供给 TPU 的数据等任务。

Cloud TPU

Cloud TPU 服务可让您通过 Google Cloud Platform (GCP) 访问 TPU。您可以使用 Cloud TPU 在 Google 的 TPU 加速器硬件上运行机器学习工作负载。Cloud TPU 旨在实现最高性能与灵活性,帮助研究人员、开发者和企业训练机器学习工作负载。

您可以使用 Cloud TPU API 来自动管理您的 TPU。如此一来,您就可以轻松扩容为大规模计算集群、运行工作负载,并在工作负载完成时缩减这些集群。芯片中内置的硬件支持可为各种深度学习工作负载实现高效的线性性能扩缩。

TPU 配置

Cloud TPU 提供不同的 TPU 配置:

  • TPU Pod
  • TPU 切片
  • TPU 板

TPU Pod

在一个 TPU Pod 中,TPU Pod 中的所有 TPU 芯片通过高速互连直接连接,避免了经过 CPU 主机导致的通信延迟。芯片以二维环面的形式连接,每个芯片可直接与相邻的四个芯片通信。这种架构可使机器学习工作负载中的常见通信模式(例如全缩减)实现超高性能。下表显示了 TPU v2 和 v3 的架构详情和性能特征,详细信息请参阅用于训练深度神经网络的领域专用超级计算机

特征 TPU v2 TPU v3
网络链接 x Gb/s/芯片 4 x 496 4 x 656
对分带宽 Tb/全 Pod 15.9 42

TPU 切片

TPU 切片是 TPU Pod 的一部分。如果您不需要整个 Pod 的资源,则可以使用 Pod 的一部分。如需了解不同切片配置的详情,请参阅 Cloud TPU 中的“Cloud TPU 价格”部分。

单个 TPU 板

单板 TPU 配置是一个独立的板,包含 4 个 TPU 芯片(8 个 TPU 核心),与其他 TPU 板没有网络连接。单板 TPU 不属于 TPU Pod 配置,也不占用 TPU Pod。如需了解可用的单板 TPU 配置,请参阅 TPU 类型页面。

TPU 版本

TPU 版本定义了每个 TPU 核心的架构和高带宽内存 (HBM) 的数量、每个 TPU 板上核心之间的互连以及可用于设备间通信的网络接口。可用的 TPU 版本包括 v1 和 v2。

TPU v2

TPU v2 板包含四个 TPU 芯片,每个芯片有两个核心。每个 TPU 核心有 8 GiB HBM 和一个 MXU。一个 TPU v2 Pod 最多可具有 512 个 TPU 核心和 4 TiB 内存。

图片

TPU v3

一个 TPU v3 电路板包含四个 TPU 芯片,每个芯片有两个核心。每个 TPU 核心有 16 GiB HBM 和两个 MXU。一个 TPU v3 Pod 最多可具有 2048 个 TPU 核心和 32 TiB 内存。

图片

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 中也可能会受限于输入。流水线性能指南可以帮助解决馈入问题。

Cloud TPU 虚拟机架构

每个 TPU 板物理连接到一个主机(TPU 主机)。

图片

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

图片

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

TPU 节点

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

图片

TPU 虚拟机

使用 TPU 虚拟机时,您可以通过 SSH 直接连接到在 TPU 主机上运行的 Google Compute Engine 虚拟机。您可以获取机器的根访问权限,以便运行任何所需的代码。您可以直接从 TPU 编译器和运行时访问调试日志和错误消息。TPU 虚拟机可支持 TPU 节点无法实现的新使用场景。例如,您可以在输入流水线中执行自定义操作,并且可以使用本地存储。

图片

由于没有用户虚拟机,因此您的代码与 TPU 虚拟机之间无需网络、Virtual Private Cloud 或防火墙,从而提高了输入流水线的性能。此外,TPU 虚拟机更便宜,因为您无需为用户虚拟机付费。

支持的框架

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 图的用户代码。用户代码在单独的远程虚拟机中运行,该虚拟机通过网络与 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 类似,每个主机虚拟机都需要一个单独的用户虚拟机。

图片

后续步骤