系统架构

张量处理单元 (TPU) 是 Google 设计的应用专用集成电路 (ASIC),用于加速机器学习工作负载。Cloud TPU 是一项 Google Cloud 服务,它可将 TPU 作为可伸缩资源提供。

TPU 旨在快速执行矩阵操作,因此非常适合机器学习工作负载。您可以使用 TensorFlowPytorchJAX 等框架在 TPU 上运行机器学习工作负载。

Cloud TPU 术语

如果您刚开始接触 Cloud TPU,请查看 TPU 文档首页。 以下部分介绍了本文档中使用的术语和相关概念。

批量推理

批量或离线推理是指在生产流水线之外进行推理,通常对大量输入进行推理。批量推理可用于离线任务(如为数据添加标签)以及评估经过训练的模型。延迟时间 SLO 不是批量推理的优先级。

TPU 芯片

TPU 芯片包含一个或多个 TensorCore。TensorCore 的数量取决于 TPU 芯片的版本。每个 TensorCore 由一个或多个矩阵乘法单元 (MXU)、一个矢量单元和一个标量单元组成。

MXU 由脉动阵列中的 128 x 128 乘法累加器组成。 MXU 在 TensorCore 中提供大部分计算能力。每个 MXU 能够在每个周期执行 16K 乘法累加操作。所有乘法均接受 bfloat16 输入,但所有累加均以 FP32 数字格式执行。

该矢量单位用于一般计算,例如激活和 softmax。标量单位用于控制流、计算内存地址和其他维护操作。

TPU 立方体

4x4x4 拓扑。这仅适用于 3D 拓扑(从 v4 TPU 版本开始)。

推理

推断是使用训练好的模型对新数据进行预测的过程。供传送进程使用。

多切片与单切片

多切片是一组切片,将 TPU 连接扩展到芯片间互连 (ICI) 连接之外,并利用数据中心网络 (DCN) 在切片之外传输数据。每个切片中的数据仍然由 ICI 传输。利用这种混合连接,Multislice 可实现多个切片的并行性,并且允许您为单个作业使用的 TPU 核心数量超出单个切片能够容纳的 TPU 核心数量。

TPU 可用于在单个切片或多个切片上运行作业。 如需了解详情,请参阅多切片简介

Cloud TPU ICI 弹性

ICI 弹性有助于提高在立方体之间连接 TPU 的光链路和光学电路开关 (OCS) 的容错能力。(立方体内的 II 连接使用不受影响的铜色链接)。 ICI 弹性允许 ICI 连接围绕 OCS 和光学 ICI 故障进行路由。因此,它提高了 TPU 切片的调度可用性,但以暂时降低 ICI 性能为代价。

与 Cloud TPU v4 类似,对于是一个立方体或更大的 v5p 切片,系统会默认启用 ICI 弹性:

  • 指定加速器类型时为 v5p-128
  • 指定加速器配置时为 4x4x4

已加入队列的资源

TPU 资源的表示形式,用于将针对单切片或多切片 TPU 环境的请求加入队列和管理请求。如需了解详情,请参阅已加入队列的资源用户指南

正在处理

服务是将经过训练的机器学习模型部署到生产环境的过程,在该环境中,模型可用于进行预测或决策。延迟时间和服务级可用性对于服务很重要。

单个主机和多主机

TPU 主机是在连接到 TPU 硬件的实体计算机上运行的虚拟机。TPU 工作负载可以使用一个或多个主机。

单主机工作负载仅限一个 TPU 虚拟机。多主机工作负载会在多个 TPU 虚拟机上分布训练。

切片

Pod 切片是位于同一 TPU Pod 内的一系列芯片,这些芯片由高速芯片间互连 (ICI) 连接。切片根据芯片或 TensorCore 进行描述,具体取决于 TPU 版本。

条状标签形状条状标签拓扑也是指切片形状。

SparseCore

v5p 在每个芯片中包含四个 SparseCore,这些核心是 Dataflow 处理器,能够为依赖于推荐模型中的嵌入的模型加速模型。

TPU Pod

TPU Pod 是通过专用网络分组的一组连续的 TPU。TPU Pod 中的 TPU 芯片的数量取决于 TPU 版本。

TPU 虚拟机或工作器

运行 Linux 的虚拟机,能够访问底层 TPU。TPU 虚拟机也称为“工作器”。

TensorCores

TPU 芯片具有一个或两个 TensorCore,用于运行矩阵乘法。如需详细了解 TensorCore,请参阅这篇 ACM 文章

工作器

请参阅 TPU 虚拟机

TPU 版本

TPU 芯片的确切架构取决于您使用的 TPU 版本。每个 TPU 版本还支持不同的切片大小和配置。如需详细了解系统架构和支持的配置,请参阅以下页面:

Cloud TPU 虚拟机架构

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

TPU 虚拟机架构

借助 TPU 虚拟机架构,您可以使用 SSH 直接连接到物理连接到 TPU 设备的虚拟机。您拥有虚拟机的 root 访问权限,因此可以运行任意代码。您可以访问编译器和运行时调试日志以及错误消息。

图片

TPU 节点架构

TPU 节点架构由通过 gRPC 与 TPU 主机通信的用户虚拟机组成。使用此架构时,您无法直接访问 TPU 主机,因此很难调试训练和 TPU 错误。

图片

后续步骤