系统架构

张量处理单元 (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 由 256 x 256(TPU v6e)或 128 x 128(TPU 版本低于 v6e)的乘法累加器组成,这些累加器位于一个脉动阵列中。MXU 在 TensorCore 中提供大部分计算能力。每个 MXU 能够在每个周期中执行 16K 乘法累加运算。所有乘法都采用 bfloat16 输入,但所有累加均采用 FP32 数字格式执行。

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

TPU Cube

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

推断

推理是一种使用经过训练的模型对新数据进行预测的过程。服务进程会使用它。

多切片与单切片

多切片是一组切片,可将 TPU 连接扩展到芯片间互连 (ICI) 连接之外,并利用数据中心网络 (DCN) 在切片之外传输数据。每个 slice 中的数据仍由 ICI 传输。借助这种混合连接,多切片可实现跨切片并行处理,让您可以为单个作业使用比单个切片可容纳的更多 TPU 核心。

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

Cloud TPU ICI 弹性

ICI 弹性有助于提高在立方体之间连接 TPU 的光纤链路和光电路开关 (OCS) 的容错性。(立方体内的 ICI 连接使用不受影响的铜线)。ICI 弹性功能可绕过 OCS 和光学 ICI 故障来路由 ICI 连接。因此,它可以提高 TPU slice 的调度可用性,但代价是 ICI 性能会暂时下降。

与 Cloud TPU v4 类似,对于大小为 1 个立方或更大的 v5p 切片,系统默认启用 ICI 弹性功能:

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

已排队的资源

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

提供服务

服务是指将训练好的机器学习模型部署到生产环境中的过程,以便在该环境中使用该模型进行预测或决策。延迟时间和服务级别可用性对于广告投放至关重要。

单主机、多主机和子主机

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

单主机工作负载仅限于一个 TPU 虚拟机。多主机工作负载会在多个 TPU 虚拟机之间分配训练任务。子主机工作负载不会使用 TPU 虚拟机上的所有芯片。

切片

Pod 切片是一组位于同一 TPU Pod 内的芯片,这些芯片通过高速芯片间互连 (ICI) 连接。切片以芯片或 TensorCore 的形式进行描述,具体取决于 TPU 版本。

条状标签形状条状标签拓扑也指的是 slice 形状。

SparseCore

SparseCore 是一种数据流处理器,可利用推荐模型中的嵌入来加速模型。v5p 每个芯片包含 4 个 SparseCore,v6e 每个芯片包含 2 个 SparseCore。

TPU Pod

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

TPU 虚拟机或工作进程

运行 Linux 且有权访问底层 TPU 的虚拟机。TPU 虚拟机也称为工作器

TensorCores

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

工作器

请参阅 TPU 虚拟机

TPU 版本

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

TPU 架构

有两种 TPU 架构描述了虚拟机如何与 TPU 设备进行物理连接:TPU 节点和 TPU 虚拟机。TPU 节点是 v2 和 v3 TPU 版本的原始 TPU 架构。在 v4 中,TPU 虚拟机成为默认架构,但同时支持这两种架构。TPU 节点架构已废弃,仅支持 TPU 虚拟机。如果您使用的是 TPU 节点,请参阅从 TPU 节点架构迁移到 TPU 虚拟机架构,了解如何从 TPU 节点架构迁移到 TPU 虚拟机架构。

TPU 虚拟机架构

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

图片

TPU 节点架构

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

图片

从 TPU 节点架构迁移到 TPU 虚拟机架构

如果您有使用 TPU 节点架构的 TPU,请按照以下步骤将其识别、删除并重新配置为 TPU 虚拟机。

  1. 前往 TPU 页面:

    前往 TPU

    1. 架构标题下方找到您的 TPU 及其架构。 如果架构为“TPU VM”,则无需执行任何操作。如果架构为“TPU 节点”,您需要删除并重新预配 TPU。
  2. 删除并重新预配 TPU。

    如需了解如何删除重新预配 TPU,请参阅管理 TPU

后续步骤