TPU v5p
本文档介绍了 Cloud TPU v5p 的架构和受支持的配置。
系统架构
本部分介绍了 v5p 版本专有的系统架构。每个 TensorCore 都有四个矩阵乘法单元 (MXU)、一个矢量单元和一个标量单元。
单个 v5p Pod 中包含 8960 个芯片。可调度的最大作业是 96 个 cube（6144 个芯片）作业。
下表显示了 v5p 的主要规范。
| 主要规格 | v5p 值 |
|---|---|
| 每片芯片的峰值计算能力 (bf16) | 459 TFLOPS |
| HBM2e 容量和带宽 | 95GB、2765 GBps |
| TPU Pod 大小 | 8960 个条状标签 |
| 互连拓扑 | 3D 环形* |
| 芯片间互连带宽 | 4800 Gbps |
配置
TPU v5p Pod 由 8960 个芯片组成，这些芯片通过可重配置的高速链路互连。借助 TPU v5p 的灵活网络功能，您可以通过多种方式连接同一大小的切片中的芯片。使用 gcloud compute tpus tpu-vm create 命令创建 TPU 切片时，您可以使用 AcceleratorType 或 AcceleratorConfig 参数指定其类型和形状。
下表显示了 v5p 支持的最常见的单 Slice 形状，以及大于 1 个 Cube 的大多数（但不是全部）完整 Cube 形状。最大 v5p 形状为 16x16x24（6144 个条状标签，96 个立方体）。
| Slice 形状 | 虚拟机大小 | # 核心 | 条状标签 | 机器数量 | 搁架数量 | 是否支持 Twisted？ |
| 2x2x1 | 完整主机 | 8 | 4 | 1 | 无 | 无 |
| 2x2x2 | 完整主机 | 16 | 8 | 2 | 无 | 无 |
| 2x4x4 | 完整主机 | 64 | 32 | 8 | 无 | 无 |
| 4x4x4 | 完整主机 | 128 | 64 | 16 | 1 | 无 |
| 4x4x8 | 完整主机 | 256 | 128 | 32 | 2 | 是 |
| 4x8x8 | 完整主机 | 512 | 256 | 64 | 4 | 是 |
| 8x8x8 | 完整主机 | 1024 | 512 | 128 | 8 | 无 |
| 8x8x16 | 完整主机 | 2048 | 1024 | 256 | 16 | 是 |
| 8x16x16 | 完整主机 | 4096 | 2048 | 512 | 32 | 是 |
| 16x16x16 | 完整主机 | 8192 | 4096 | 1024 | 64 | 无 |
| 16x16x24 | 完整主机 | 12288 | 6144 | 1536 | 96 | 无 |
单 slice 训练最多支持 6144 个芯片。它可以使用多切片扩展到 18432 个芯片。如需了解 Multislice 的详细信息，请参阅 Cloud TPU Multislice 概览。
使用 AcceleratorType 参数
分配 TPU 资源时，您可以使用 --accelerator-type 参数指定切片中的 TensorCore 数量。--accelerator-type 是一个格式化字符串“v$VERSION_NUMBERp-$CORES_COUNT”。例如，v5p-32 指定了具有 32 个 TensorCore（16 个芯片）的 v5p TPU 切片。
如需为 v5p 训练作业预配 TPU，请在 CLI 或 TPU API 创建请求中使用以下加速器类型之一：
- v5p-8
- v5p-16
- v5p-32
- v5p-64
- v5p-128（一个完整立方体/机架）
- v5p-256（2 个立方体）
- v5p-512
- v5p-1024 ... v5p-12288
使用 AcceleratorConfig 参数
对于 v5p 及更高版本的 Cloud TPU，AcceleratorConfig 的使用方式与 Cloud TPU v4 大同小异。不同之处在于，您将 TPU 类型指定为您使用的 TPU 版本（例如，对于 v5p 版本，指定为 --type=v5p），而不是指定为 --type=v4。
Cloud TPU ICI 弹性
ICI 弹性有助于提高光纤链路和光电路开关 (OCS) 的容错性，这些光电路开关用于在各个 Cube 之间连接 TPU。（立方体内的 ICI 连接使用不受影响的铜线）。ICI 弹性功能可绕过 OCS 和光学 ICI 故障来路由 ICI 连接。因此，它可以提高 TPU slice 的调度可用性，但代价是 ICI 性能会暂时下降。
与 Cloud TPU v4 类似，对于大小等于或大于一个立方体的 v5p 切片，系统默认会启用 ICI 弹性功能：
- 指定加速器类型时使用 v5p-128
- 指定加速器配置时为 4x4x4
虚拟机、主机和 slice 属性
| 属性 | TPU 的价值 |
| v5p 芯片数量 | 4 |
| vCPU 的数量 | 208（如果使用 NUMA 绑定以避免跨 NUMA 性能惩罚，则只能使用一半） |
| RAM (GB) | 448（如果使用 NUMA 绑定以避免跨 NUMA 性能惩罚，则只能使用一半） |
| NUMA 节点数 | 2 |
| NIC 吞吐量 (Gbps) | 200 |
一个 Pod 中的 TensorCore、芯片、主机/虚拟机和 Cube 数量之间的关系：
| 核心 | 条状标签 | 主机/虚拟机 | Cubes | |
|---|---|---|---|---|
| 主机 | 8 | 4 | 1 | |
| Cube（也称为机架） | 128 | 64 | 16 | 1 |
| 支持的最大 slice | 12288 | 6144 | 1536 | 96 |
| v5p 完整 Pod | 17920 | 8960 | 2240 | 140 |