系统架构
张量处理单元 (TPU) 是 Google 设计的应用专用集成电路 (ASIC)，用于加速机器学习工作负载。Cloud TPU 是一项 Google Cloud 服务，用于将 TPU 作为可伸缩资源提供。
TPU 旨在快速执行矩阵操作，使其非常适合用于机器学习工作负载。您可以使用 TensorFlow、Pytorch 和 JAX 等框架在 TPU 上运行机器学习工作负载。
Cloud TPU 术语
如果您刚开始接触 Cloud TPU，请参阅 TPU 文档首页。以下部分介绍了本文档中使用的术语和相关概念。
批量推断
批量或离线推断是指在生产流水线之外进行推断，通常基于大量输入。批量推断可用于离线任务（例如为数据加标签），以及评估经过训练的模型。延迟时间 SLO 不是批量推断的优先级。
推断
推断是使用经过训练的模型根据新数据进行预测的过程。供传送进程使用。
已加入队列的资源
一种 TPU 资源表示法，用于将单切片或多切片 TPU 环境加入队列和管理请求。如需了解详情，请参阅队列资源用户指南。
提供服务
服务是将经过训练的机器学习模型部署到生产环境时，该模型可用于进行预测或决策的过程。延迟时间和服务级可用性对于传送非常重要。
单个主机和多主机
TPU 主机是在连接到 TPU 硬件的物理计算机上运行的虚拟机。TPU 工作负载可以使用一个或多个主机。
单主机工作负载仅限 1 个 TPU 虚拟机，可以使用 1 个、4 个或 8 个 TPU 芯片。多主机 TPU v5e 工作负载可以访问 8 个、12 个、16 个、32、64、128 或 256 个 TPU 芯片，并且每 4 个 TPU 芯片对应一个 TPU 虚拟机。多主机工作负载会将训练分布到多个 TPU 虚拟机。
TPU v5e 支持单主机和多主机训练以及单主机推断。使用 Sax 支持多主机推断。 如需了解详情，请参阅大型语言模型服务。
切片
Pod 切片是全部位于同一个 TPU Pod 内的芯片集合，通过高速芯片间互连 (ICI) 连接。
v5e 切片使用 2D 切片形状进行描述。切片形状中的每个数字对应于某个维度中的 v5e 条状标签数量。例如，4x2 描述了 8 个 v5e 芯片在 4 x 2 网格中的排列。
TPU v4 切片可以用 v4 芯片进行描述。v4 切片以 3D 形状进行描述。切片形状中的每个数字均对应于某个维度中的 v4 条状标签数量。例如，4x4x8 描述了 128 个 v4 芯片在 4 x 4 x 8 的立方体中排列。
v4 切片也可以根据切片中的 TensorCores 数量进行描述。例如，v4-128 描述了一个包含 128 个 TensorCore 的 v4 切片。v4 切片具有 16、32、64、128、256、512、1024、2048 或 4096 个 TensorCore。
TPU v3 切片根据 TensorCore 进行描述，提供 32 个、128 个、512 个、1024 个或 2048 个 TensorCore。TPU v2 切片也是从 TensorCore 的角度进行描述，并且具有 32、128、256 或 512 个 TensorCore。
条状标签形状和条状标签拓扑也指切片形状。
如需查看 v5e 支持的切片形状的列表，请参阅加速器类型部分中的表格。
TPU 芯片
一个 TPU 芯片包含一个或多个 TensorCore。TensorCore 的数量取决于 TPU 芯片的版本。每个 TensorCore 由一个或多个矩阵乘法单元 (MXU)、一个矢量单位和一个标量单位组成。
MXU 由脉动阵列中的 128 x 128 乘法累加器组成。MXU 在 TensorCore 中提供大部分计算能力。每个 MXU 在每个周期能够执行 16K 乘法累加操作。所有乘法都采用 bfloat16 输入，但所有累计都以 FP32 数字格式执行。
该矢量单位用于一般计算，例如激活和 Softmax。 标量单位用于控制流、计算内存地址和其他维护操作。
TensorCores
TPU 芯片具有一个或两个 TensorCore，用于运行矩阵乘法。与 v2 和 v3 Pod 类似，v5e 的每个芯片中都有一个 TensorCore。相比之下，v4 Pod 的每个芯片有 2 个 TensorCore。如需详细了解 TensorCore，请参阅 ACM 文章。
TPU Pod
TPU Pod 是通过专用网络组合在一起的一组连续的 TPU。TPU Pod 中的 TPU 芯片的数量取决于 TPU 版本。
TPU 虚拟机
一个运行 Linux 并可以访问底层 TPU 的虚拟机。对于 v5e TPU，每个 TPU 虚拟机可以直接使用 1、4 或 8 个芯片，具体取决于用户指定的加速器类型。对于 v4 及更低版本，每个 TPU 虚拟机可以使用 4 个 TPU 芯片。TPU 虚拟机也称为“工作器”。
工作器
请参阅 TPU 虚拟机。
TPU 版本
TPU 的确切布局取决于您使用的 TPU 版本。如需了解 TPU v2 和 v3 的架构详情和性能特征，请参阅用于训练深度神经网络的领域专用超级计算机。
TPU v5e
每个 v5e 芯片包含一个 TensorCore。每个 TensorCore 都有 4 个矩阵乘法单位 (MXU)、一个矢量单位和一个标量单位。
下图展示了 TPU v5e 芯片。
下表显示了 v5e 的密钥条状标签规范及其值。
| 密钥条状标签规范 | v5e 值 |
| 每个芯片的峰值计算 (bf16) | 197 个 TFLOP |
| 每个芯片的峰值计算 (Int8) | 393 个 TFLOP |
| HBM2 容量和带宽 | 16 GB，819 GBps |
| 芯片间互连 BW | 1600 Gbps |
下表显示了 v5e 的 Pod 规范及其值。
| 关键 Pod 规范 | v5e 值 |
| TPU Pod 大小 | 256 个条状标签 |
| 互连拓扑 | 2D 环岛 |
| 每个 Pod 的峰值计算 | 100 PetaOps(Int8) |
| 每个 Pod 的全减带宽 | 51.2 TB/秒 |
| 每个 Pod 的对分带宽 | 1.6 TB/秒 |
| 每个 Pod 的数据中心网络带宽 | 6.4 Tbps |
TPU v4
每个 TPU v4 芯片包含两个 TensorCore。每个 TensorCore 都有四个 MXU、一个矢量单位和一个标量单位。下表显示了 v4 TPU Pod 的关键规范。
| 密钥规范 | v4 Pod 值 |
|---|---|
| 每个芯片的峰值计算能力 | 每秒 275 万亿次浮点运算（bf16 或 int8） |
| HBM2 容量和带宽 | 32 GiB，1200 GBps |
| 测量的最低/平均/最大功率 | 90/170/192 瓦 |
| TPU Pod 大小 | 4096 芯片 |
| 互连拓扑 | 3D 网格 |
| 每个 Pod 的峰值计算 | 1.1 百亿亿次浮点运算（bf16 或 int8） |
| 每个 Pod 的全减带宽 | 1.1 PB/秒 |
| 每个 Pod 的对分带宽 | 24 TB/秒 |
下图展示了 TPU v4 芯片。
3D 网格和 3D 环面
v4 TPU 在三维空间内与最近的相邻芯片建立直接连接，从而形成网络连接的 3D 网格。当切片等于或大于单个立方体时，可以将连接配置为 3D 环面。一般来说，3D 配置的性能要优于 3D 网格配置。
TPU v3
每个 v3 TPU 芯片包含两个 TensorCore。每个 TensorCore 有两个 MXU、一个矢量单位和一个标量单位。下表显示了 v3 TPU Pod 的关键规范及其值。
| 密钥规范 | v3 Pod 值 |
|---|---|
| 每个芯片的峰值计算能力 | 123 万亿次浮点运算 (bf16) |
| HBM2 容量和带宽 | 32 GiB、900 GBps |
| 测量的最低/平均/最大功率 | 123/220/262 瓦 |
| TPU Pod 大小 | 1024 芯片 |
| 互连拓扑 | 2D 环面 |
| 每个 Pod 的峰值计算 | 每秒 126 千万亿次浮点运算 (bf16) |
| 每个 Pod 的全减带宽 | 340 TB/秒 |
| 每个 Pod 的对分带宽 | 6.4 TB/秒 |
下图展示了 TPU v3 芯片。
TPU v4 与 v3 相比的性能优势
本部分介绍 TPU v4 的性能优势
内存系统：
非统一内存访问 (NUMA) 是一种适用于具有多个 CPU 的机器的计算机内存架构。每个 CPU 可以直接访问高速内存块。CPU 和内存称为 NUMA 节点。NUMA 节点连接到直接相邻的 NUMA 节点。来自一个 NUMA 节点的 CPU 可以访问另一个 NUMA 节点中的内存，但这种访问比访问 NUMA 节点中的内存慢。
在多 CPU 机器上运行的软件可以将 CPU 所需的数据存储在其 NUMA 节点中，从而提高内存吞吐量。如需详细了解 NUMA，请参阅维基百科上的非统一内存访问。
您可以将训练脚本绑定到 NUMA 节点 0，以利用 NUMA 位置方面的优势。
要启用 NUMA 节点绑定，请执行以下操作：
安装 numactl 命令行工具。
$ sudo apt-get update $ sudo apt-get install numactl
启动训练脚本时使用
numactl --cpunodebind=0。这会将您的脚本代码绑定到 NUMA 节点 0。$ numactl --cpunodebind=0 python3 your-training-script
在以下情况下，请启用 NUMA 节点绑定：
- 如果您的工作负载严重依赖 CPU 工作负载（例如，图片分类、推荐工作负载），而不考虑框架。
- 如果您使用的是不带 -pod 后缀的 TPU 运行时版本（例如
tpu-vm-tf-2.10.0-v4），
内存系统的其他差异：
- v4 TPU 芯片在整个芯片中拥有统一的 32 GiB HBM 内存空间，从而能够提升两个芯片上的 TensorCore 之间的更好地协调。
- 使用最新的内存标准和速度提升了 HBM 性能。
- 通过对 512B 粒度的高性能步进支持，改进了 DMA 性能配置文件。
TensorCore：
- 将 MXU 的数量翻了一番，并提高了时钟频率，最多可提供 275 个 TFLOPS。
- 2 倍转置和排列带宽。
- 通用内存 (Cmem) 的加载存储内存访问模型。
- 更快的 MXU 权重加载带宽和 8 位模式支持，可以降低批次大小并缩短推断延迟时间。
芯片间互连：
每个芯片具有六个互连链接，可实现网络直径较小的网络拓扑。
其他：
- 与主机之间的 x16 PCIE gen3 接口（直接连接）。
- 改进了安全模型。
- 提高了能效。
TPU v3 相对于 v2 的性能优势
TPU v3 配置中每个 TensorCore 的 FLOPS 增加和内存容量可以通过以下方式提高模型的性能：
对于受限于计算的模型，每个 TensorCore 的 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 虚拟机架构：TPU 节点还是 TPU 虚拟机。
TPU 节点架构
TPU 节点架构由一个通过 gRPC 与 TPU 主机通信的用户虚拟机组成。使用此架构时，您无法直接访问 TPU 主机，因此难以调试训练和 TPU 错误。
TPU 虚拟机架构
借助 TPU 虚拟机架构，您可以使用 SSH 直接连接到物理连接到 TPU 设备的虚拟机。您拥有虚拟机的 root 访问权限，因此可以运行任意代码。您可以访问编译器和运行时调试日志以及错误消息。
