TPU v6e
本文档介绍了 Cloud TPU v6e (Trillium) 的架构和支持的配置。
Trillium 是 Cloud TPU 的最新一代 AI 加速器。在所有技术平台（例如 API 和日志）以及本文档中，Trillium 都将被称为 v6e。
v6e 每个 Pod 的芯片占用空间为 256 个，与 v5e 有许多相似之处。此系统经过优化，可成为用于转换器、文本到图像和卷积神经网络 (CNN) 训练、微调和服务的最高价值产品。
系统架构
每个 v6e 芯片包含一个 TensorCore。每个 TensorCore 都有 4 个矩阵乘法单元 (MXU)、一个矢量单元和一个标量单元。下表比较了 TPU v6e 与 TPU v5e 的主要规格及其值。
| 规范 | v5e | v6e |
|---|---|---|
| 性能/总拥有成本 (TCO)（预计值） | 0.65x | 1 |
| 每片芯片的峰值计算能力 (bf16) | 197 TFLOPS | 918 TFLOPS |
| 每片芯片的峰值计算能力 (Int8) | 393 TOPs | 1836 TOPs |
| 每个芯片的 HBM 容量 | 16 GB | 32 GB |
| 每个芯片的 HBM 带宽 | 819 GBps | 1640 GBps |
| 芯片间互连 (ICI) 带宽 | 1600 Gbps | 3584 Gbps |
| 每个芯片的 ICI 端口数 | 4 | 4 |
| 每个主机的 DRAM | 512 GiB | 1536 GiB |
| 每个主机的条状标签 | 8 | 8 |
| TPU Pod 大小 | 256 个条状标签 | 256 个条状标签 |
| 互连拓扑 | 2D 环面 | 2D 环面 |
| 每个 Pod 的 BF16 峰值计算 | 50.63 PFLOPs | 234.9 PFLOPs |
| 每个 Pod 的 All-reduce 带宽 | 51.2 TB/s | 102.4 TB/s |
| 每个 Pod 的对分带宽 | 1.6 TB/s | 3.2 TB/s |
| 每个主机的 NIC 配置 | 2 x 100 Gbps NIC | 4 x 200 Gbps NIC |
| 每个 Pod 的数据中心网络带宽 | 6.4 Tbps | 25.6 Tbps |
| 特殊功能 | - | SparseCore |
受支持的配置
TPU v6e 支持最多 256 个 v6e 芯片的训练，以及最多 8 个芯片的单主机推理。
下表显示了 v6e 支持的 2D 切片形状：
| 拓扑 | TPU 芯片 | 主机 | 虚拟机 | 加速器类型 (TPU API) | 机器类型 (GKE API) | 范围 | 是否支持推理？ |
|---|---|---|---|---|---|---|---|
| 1x1 | 1 | 1/8 | 1 | v6e-1 |
ct6e-standard-1t |
子主机 | 是 |
| 2x2 | 4 | 1/2 | 1 | v6e-4 |
ct6e-standard-4t |
子主机 | 是 |
| 2x4 | 8 | 1 | 1 | v6e-8 |
ct6e-standard-8t |
单个主机 | 是 |
| 2x4 | 8 | 1 | 2 | - | ct6e-standard-4t |
单个主机 | 否 |
| 4x4 | 16 | 2 | 4 | v6e-16 |
ct6e-standard-4t |
多主机 | 否 |
| 4x8 | 32 | 4 | 8 | v6e-32 |
ct6e-standard-4t |
多主机 | 否 |
| 8x8 | 64 | 8 | 16 | v6e-64 |
ct6e-standard-4t |
多主机 | 否 |
| 8x16 | 128 | 16 | 32 | v6e-128 |
ct6e-standard-4t |
多主机 | 否 |
| 16x16 | 256 | 32 | 64 | v6e-256 |
ct6e-standard-4t |
多主机 | 否 |
附加到单个虚拟机的包含 8 个芯片的 Slice (v6e-8) 经过优化，可用于推理，从而允许在单个服务工作负载中使用所有 8 个芯片。
如需了解每种拓扑的虚拟机数量，请参阅虚拟机类型。
虚拟机类型
每个 TPU v6e 虚拟机都可能包含 1 个、4 个或 8 个芯片。4 芯片及更小的 slice 具有相同的非统一内存访问 (NUMA) 节点。如需详细了解 NUMA 节点，请参阅维基百科上的非统一内存访问。
v6e 切片是使用半主机虚拟机创建的，每个虚拟机包含 4 个 TPU 芯片。这条规则有两个例外情况：
v6e-1：仅包含单个芯片的虚拟机，主要用于测试v6e-8：针对推理用例进行了优化的全主机虚拟机，其中所有 8 个芯片都连接到单个虚拟机。
下表显示了 TPU v6e 虚拟机类型的比较：
| 虚拟机类型 | 每个虚拟机的 vCPU 数量 | 每个虚拟机的 RAM (GB) | 每个虚拟机的 NUMA 节点数 |
|---|---|---|---|
| 1 芯片虚拟机 | 44 | 176 | 1 |
| 4 芯片虚拟机 | 180 | 720 | 1 |
| 8 芯片虚拟机 | 180 | 1440 | 2 |
指定 v6e 配置
使用 TPU API 分配 TPU v6e slice 时，您可以使用 AcceleratorType 或 AcceleratorConfig 参数指定其大小和形状。
如果您使用的是 GKE，请使用 --machine-type 标志指定支持您要使用的 TPU 的机器类型。如需了解详情，请参阅 GKE 文档中的在 GKE Standard 中部署 TPU 工作负载。
使用 AcceleratorType
分配 TPU 资源时，您可以使用 AcceleratorType 指定切片中的 TensorCore 数量。您为 AcceleratorType 指定的值是一个格式为 v$VERSION-$TENSORCORE_COUNT 的字符串。例如，v6e-8 指定了具有 8 个 TensorCore 的 v6e TPU 切片。
以下示例展示了如何使用 AcceleratorType 创建具有 32 个 TensorCore 的 TPU v6e slice：
gcloud
$ gcloud compute tpus tpu-vm create tpu-name \ --zone=zone \ --accelerator-type=v6e-32 \ --version=v2-alpha-tpuv6e
控制台
使用 AcceleratorConfig
您还可以使用 AcceleratorConfig 指定要使用的 TensorCore 数量。不过，由于 TPU v6e 没有自定义 2D 拓扑变体，因此使用 AcceleratorConfig 和 AcceleratorType 没有区别。
如需使用 AcceleratorConfig 配置 TPU v6e，请使用 --version 和 --topology 标志。将 --version 设置为要使用的 TPU 版本，并将 --topology 设置为切片中 TPU 芯片的物理排列方式。您为 AcceleratorConfig 指定的值是一个格式为 AxB 的字符串，其中 A 和 B 是每个方向上的条状标签数量。
以下示例展示了如何使用按 4x8 拓扑排列的 AcceleratorType 创建 TPU v6e 切片 32 个 TensorCore：
$ gcloud compute tpus tpu-vm create tpu-name \ --zone=zone \ --type=v6e \ --topology=4x8 \ --version=v2-alpha-tpuv6e