TPU v5e

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

TPU v5e 支持单主机和多主机训练以及单主机推理。支持使用 Sax 进行多主机推理。如需了解详情，请参阅大语言模型服务。

系统架构

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

下图展示了 TPU v5e 芯片。

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

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

配置

Cloud TPU v5e 是一款结合了训练和推理（服务）功能的产品。如需区分训练环境和推理环境，请在创建 GKE 节点池时将 AcceleratorType 参数与 TPU API 或 --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 通信更快。

用于分发的 Cloud TPU v5e 类型

单主机传送最多支持 8 个 v5e 芯片。支持以下配置：1x1、2x2 和 2x4 切片。每个 slice 分别包含 1、4 和 8 个芯片。

如需为分发作业预配 TPU，请在 CLI 或 API TPU 创建请求中使用以下某种加速器类型：

以下命令会创建一个包含 8 个 v5e 芯片的 v5e TPU Slice，以供分发：

  $ 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 v5e 类型

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

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

以下命令会创建一个包含 256 个 v5e 芯片的 v5e TPU Slice，用于训练：

  $ 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 的系统架构，请参阅系统架构。