TPU v5e
本文档介绍了 Cloud TPU v5e。
TPU v5e 支持单主机和多主机训练以及单主机推理。支持使用 Sax 进行多主机推理。如需了解详情,请参阅大语言模型服务。
系统架构
每个 v5e 芯片包含一个 TensorCore。每个 TensorCore 都有四个矩阵乘法 单位 (MXU)、矢量单位和标量单位。
下图展示了 TPU 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 节点池时将 AcceleratorType
或 AcceleratorConfig
标志与 TPU API 或 --machine-type
标志搭配使用。
训练作业针对吞吐量和可用性进行了优化,而服务作业针对延迟时间进行了优化。在预配用于服务的 TPU 上运行的训练作业的可用性可能会较低;同样,在预配用于训练的 TPU 上执行的服务作业的延迟时间可能会较长。
您可以使用 AcceleratorType
指定要使用的 TensorCore 数量。使用AcceleratorType
gcloud CLI 或 Google Cloud 控制台。您指定的价值
为 AcceleratorType
指定的值是一个字符串,格式如下:
v$VERSION_NUMBER-$CHIP_COUNT
。
您还可以使用 AcceleratorConfig
来指定
资源。但是,由于 TPU 没有自定义 2D 拓扑变体,
v5e 中,使用 AcceleratorConfig
和
AcceleratorType
。
如需使用 AcceleratorConfig
配置 TPU v5e,请使用 --version
和 --topology
标志。将 --version
设置为要使用的 TPU 版本,将 --topology
设置为切片中 TPU 芯片的物理排列方式。通过
您为 AcceleratorConfig
指定的值是一个格式为 AxB
的字符串,
其中 A
和 B
是每个方向上的条状标签计数。
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
通信快。
用于分发的 Cloud TPU v5e 类型
最多 8 个 v5e 芯片支持单个主机传送。支持以下配置:1x1、2x2 和 2x4 切片。每个 slice 分别包含 1、4 和 8 个芯片。
如需为服务作业预配 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 芯片虚拟机。