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

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

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

配置

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

训练作业针对吞吐量和可用性进行了优化,而服务作业针对延迟时间进行了优化。在预配用于服务的 TPU 上运行的训练作业的可用性可能会较低;同样,在预配用于训练的 TPU 上执行的服务作业的延迟时间可能会较长。

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

您还可以使用 AcceleratorConfig 来指定 资源。但是,由于 TPU 没有自定义 2D 拓扑变体, v5e 中,使用 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,请使用以下加速器类型之一 添加以下代码:

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 芯片虚拟机。