系统架构
张量处理单元 (TPU) 是 Google 设计的机器学习加速器。Cloud TPU 使 TPU 可用作可扩缩的 Google Cloud 资源。您可以使用机器学习框架（如 TensorFlow、Pytorch 和 JAX）在 Cloud TPU 上运行机器学习工作负载。
TPU 芯片上的每个 TensorCore 都由一个或多个矩阵乘法单元 (MXU)、向量和标量单元组成。
MXU 由脉动阵列中的 128 x 128 乘法/累加器组成。MXU 可在 TensorCore 中提供大部分计算能力。每个 MXU 能够在每个周期中执行 16K 乘法累加运算。所有乘法都接受 bfloat16 输入，但所有累积都以 FP32 数字格式执行。
矢量单位用于激活和 softmax 等常规计算。标量单元用于控制流、计算内存地址和其他维护操作。
TPU 版本
TPU 的确切布局取决于您使用的 TPU 版本。如需了解 TPU v2 和 v3 的架构详情和性能特性，请参阅用于训练深度神经网络的领域专用超级计算机。
TPU v4
每个 v4 TPU 芯片包含两个 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 芯片。
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 v2
最小的 TPU v2 配置包含四个 TPU 芯片和 16 GiB 的 HBM。每个 TPU 芯片包含两个 TensorCore。每个 TensorCore 都有一个 MXU、矢量单位和标量单位。下图展示了 TPU v2 芯片。
Cloud TPU 提供以下 TPU 配置：
- 单个 TPU 设备
- TPU Pod - 通过高速互连连接的一组 TPU 设备
- TPU Slice - TPU Pod 的一个细分
与 v3 相比，TPU v4 的性能优势
内存系统：
- v4 TPU 芯片在整个芯片上具有统一的 32 GiB HBM 内存空间，因此可以更好地协调两个芯片上 TensorCore。
- 使用最新的内存标准和速度提升了 HBM 性能。
- 通过对 512B 粒度的高性能划分，内置了对 DMA 性能配置文件的改进。
TensorCores：
- 实现 275 个最大 TFLOPS 后，MXU 的数量翻了一番，时钟频率更高了。
- 转置和排列带宽的 2 倍。
- 通用内存 (Cmem) 的加载存储区内存访问模型。
- 更快的 MXU 权重加载带宽和 8 位模式支持，允许较低的批次大小和缩短的推断延迟时间。
芯片间互连：
每个芯片有六个互连链接，以便启用网络直径较小的网络拓扑。
其他：
- 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 可让您通过 topology 字段指定条状标签大小。
TPU v4 拓扑使用 3 元组指定，该元组描述 TPU 芯片之间的数字、形状和互连。以下插图显示了一些常见的 TPU v4 拓扑。
更大的切片可以从一个或多个 4x4x4“立方体”的条状标签构建。
可以采用不同方式配置给定数量的芯片的 TPU 切片。例如，现在可将 AcceleratorType 为 v4-1024 的 TPU 切片配置为 4x4x32、4x8x16 或 8x8x8。
TPU v2 和 v3 配置
TPU 支持以下配置：
- 单个 TPU 板
- TPU Pod
- 一个 TPU Pod 切片
单个 TPU 板
单板 TPU 配置是具有四个 TPU 芯片（八个 TensorCore）的独立板，未与其他 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 设备的虚拟机。您拥有虚拟机的根访问权限，因此可以运行任意代码。您可以访问编译器和运行时调试日志以及错误消息。
支持的框架
JAX、PyTorch 和 TensorFlow 等框架使用每个虚拟机上都存在的名为 libtpu 的共享库访问 TPU。libtpu 包括 XLA 编译器、TPU 运行时软件和 TPU 驱动程序。
TensorFlow
TPU 虚拟机
使用 TPU 虚拟机时，Python 代码可以直接在 TPU 主机上运行，而不是在用户虚拟机上运行。
如需详细了解 TensorFlow 和 Cloud TPU 虚拟机，请参阅 Cloud TPU 虚拟机用户指南。
TPU 节点
Cloud TPU Node 系统架构最初是为 TensorFlow 构建的。用户无法访问 TPU 主机，并运行 TensorFlow 服务器的无头副本。它们不运行 Python 或未表示为 TensorFlow 图的任何用户代码。用户代码在单独的用户虚拟机中运行，该虚拟机通过 gRPC 网络与 TPU 主机通信。
PyTorch
TPU 虚拟机
使用 TPU 虚拟机时，您的 PyTorch 代码会直接在 TPU 主机上运行。
如需详细了解 PyTorch 和 Cloud TPU，请参阅 PyTorch/XLA 用户指南。
TPU 节点
PyTorch 使用一个名为 XRT 的库在 Cloud TPU 节点架构上运行。XRT 通过 gRPC 发送 XLA 图和运行时指令，以便在 TensorFlow 服务器上运行。 每个 TPU 主机都需要一个用户虚拟机。
JAX
TPU 虚拟机
使用 TPU 虚拟机时，您的 JAX 代码会直接在 TPU 主机上运行。
如需详细了解如何在 Cloud TPU 上运行 JAX，请参阅 JAX 快速入门。
TPU 节点
JAX 在 Cloud TPU 节点上的运行方式与 PyTorch 类似，每个主机虚拟机都需要一个单独的用户虚拟机。
