TPU v5e

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

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

系统架构

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

下图展示了 TPU v5e 芯片。

v5e 芯片示意图

下表显示了 v5e 的关键芯片规范及其值。

密钥芯片规范 v5e 值
每个芯片的计算峰值 (bf16) 197 TFLOPS
HBM2 容量和带宽 16 GB,819 GBps
芯片间互连带宽 1600 Gbps

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

主要 Pod 规范 v5e 值
TPU Pod 大小 256 个条状标签
互连拓扑 二维环面
每个 Pod 的峰值计算能力 100 PetaOps(Int8)
每个 Pod 完全减少带宽 51.2 TB/s
每个 Pod 的对分带宽 1.6 TB/s
每个 Pod 的数据中心网络带宽 6.4 Tbps

配置

Cloud TPU v5e 是一款结合了训练和推理(服务)的产品。如需区分训练环境和推理环境,请在创建 GKE 节点池时将 AcceleratorTypeAcceleratorConfig 标志与 TPU API 或 --machine-type 标志搭配使用。

训练作业针对吞吐量和可用性进行了优化,同时处理作业 已针对延迟时间进行了优化预配的 TPU 上的训练作业 可能具有较低的可用性,同样,在 TPU 上执行服务作业 因此延迟时间可能会更长

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

您还可以使用 AcceleratorConfig 指定要使用的 TensorCore 数量。不过,由于 TPU v5e 没有自定义 2D 拓扑变体,因此使用 AcceleratorConfigAcceleratorType 没有区别。

如需使用 AcceleratorConfig 配置 TPU v5e,请使用 --version--topology 标志。将 --version 设置为您要使用的 TPU 版本。 --topology 应用于切片中 TPU 芯片的物理排列。通过 您为 AcceleratorConfig 指定的值是一个格式为 AxB 的字符串, 其中 AB 是每个方向上的条状标签计数。

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

拓扑 TPU 芯片数量 主机数量
1x1 1 1/8
2x2 4 1/2
2x4 8 1
4x4 16 2
4x8 32 4
8x8 64 8
8x16 128 16
16x16 256 32

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

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

NUMA 节点通信

用于传送的 Cloud TPU v5e 类型

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

支持分发的 TPU v5e 配置:1x1、2x2 和 2x4。

如需为分发作业预配 TPU,请在 CLI 或 API TPU 创建请求中使用以下加速器类型之一:

AcceleratorType(TPU API) 机器类型 (GKE API)
v5litepod-1 ct5lp-hightpu-1t
v5litepod-4 ct5lp-hightpu-4t
v5litepod-8 ct5lp-hightpu-8t

支持在超过 8 个 v5e 条状标签上投放(也称为多主机投放) 使用 Sax。如需了解详情,请参阅大型语言模型服务

用于训练的 Cloud TPU v5e 类型

最多支持 256 个条状标签进行训练。

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

AcceleratorType(TPU API) 机器类型 (GKE API) 拓扑
v5litepod-16 ct5lp-hightpu-4t 4x4
v5litepod-32 ct5lp-hightpu-4t 4x8
v5litepod-64 ct5lp-hightpu-4t 8x8
v5litepod-128 ct5lp-hightpu-4t 8x16
v5litepod-256 ct5lp-hightpu-4t 16x16

v5e TPU 虚拟机类型比较:

虚拟机类型 n2d-48-24-v5lite-tpu n2d-192-112-v5lite-tpu n2d-384-224-v5lite-tpu
v5e 芯片数量 1 4 8
vCPU 数量 24 112 224
RAM (GB) 48 192 384
NUMA 节点数量 1 1 2
应用对象 v5litepod-1 v5litepod-4 v5litepod-8
中断

为了为需要更多芯片的工作负载腾出空间,调度程序可能会抢占芯片较少的虚拟机。因此,8 芯片虚拟机可能会抢占 1 芯片和 4 芯片虚拟机。