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 个条状标签
互连拓扑 2D 环面
每个 Pod 的峰值计算能力 100 PetaOps(Int8)
每个 Pod 的 All-reduce 带宽 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

使用 Sax 支持在 8 个以上的 v5e 芯片上提供服务(也称为多主机服务)。如需了解详情,请参阅大型语言模型服务

用于训练的 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 芯片虚拟机。