TPU v5e

本文档介绍了 Cloud TPU v5e 的架构和支持的配置。

TPU v5e 支持单主机和多主机训练，以及单主机推理。使用 Sax 支持多主机推理。 如需了解详情，请参阅 Cloud TPU 推理。

系统架构

每个 v5e 芯片包含一个 TensorCore。每个 TensorCore 都有四个矩阵乘法单元 (MXU)、一个向量单元和一个标量单元。

下图展示了 TPU v5e 芯片。

下表展示了 v5e 的主要芯片规格及其值。

下表展示了 v5e 的 Pod 规范及其值。

配置

Cloud TPU v5e 是一款集训练和推理（服务）于一体的产品。如需区分训练环境和推理环境，请在创建 Google Kubernetes Engine (GKE) 节点池时，使用 TPU API 的 AcceleratorType 参数或 --machine-type 标志。

训练作业针对吞吐量和可用性进行了优化，而服务作业则针对延迟时间进行了优化。在为部署而配置的 TPU 上运行的训练作业的可用性可能较低，同样，在为训练而配置的 TPU 上执行的部署作业的延迟时间可能较长。

您可以使用 AcceleratorType 指定要使用的 TensorCore 数量。 使用 gcloud CLI 或 Google Cloud 控制台创建 TPU 时，您需要指定 AcceleratorType。您为 AcceleratorType 指定的值是一个字符串，格式为：v$VERSION_NUMBER-$CHIP_COUNT。

v5e 支持以下 2D 切片形状：

虚拟机类型

v5e TPU 切片中的每个 TPU 虚拟机都包含 1 个、4 个或 8 个芯片。在 4 芯片及更小的切片中，所有 TPU 芯片共享同一非统一内存访问 (NUMA) 节点。

对于 8 芯片 v5e TPU 虚拟机，CPU-TPU 通信在 NUMA 分区内会更高效。例如，在下图中，CPU0-Chip0 通信将比 CPU0-Chip4 通信更快。

下表比较了 TPU v5e 虚拟机类型：

用于服务的 Cloud TPU v5e 类型

单主机部署最多支持 8 个 v5e 芯片。支持以下配置：1x1、2x2 和 2x4 切片。每个切片分别有 1 个、4 个和 8 个芯片。

如需为服务作业预配 TPU，请在 CLI 或 API TPU 创建请求中使用以下加速器类型之一：

以下命令会创建一个具有 8 个 v5e 芯片的 v5e TPU 切片，用于提供服务：

$ gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=us-central1-a \
    --accelerator-type=v5litepod-8 \
    --version=v2-alpha-tpuv5-lite

如需详细了解如何管理 TPU，请参阅管理 TPU。 如需详细了解 Cloud TPU 的系统架构，请参阅系统架构。

使用 Sax 可在超过 8 个 v5e 芯片上进行部署，也称为多主机部署。如需了解详情，请参阅 Cloud TPU 推理。

用于训练的 Cloud TPU v5e 类型

支持最多 256 个芯片的训练。

如需为 v5e 训练作业预配 TPU，请在 CLI 或 API TPU 创建请求中使用以下加速器类型之一：

以下命令会创建一个具有 256 个 v5e 芯片的 v5e TPU 切片，用于训练：

$ gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=us-east5-a \
    --accelerator-type=v5litepod-256 \
    --version=v2-alpha-tpuv5-lite

如需详细了解如何管理 TPU，请参阅管理 TPU。 如需详细了解 Cloud TPU 的系统架构，请参阅系统架构。