TPU v5p
本文档介绍了 Cloud TPU v5p 的架构和支持的配置。
系统架构
本部分介绍了 v5p 版本特有的系统架构。每个 TensorCore 都有四个矩阵乘法单元 (MXU)、一个向量单元和一个标量单元。
单个 v5p 切片中有 8960 个芯片。可安排的最大作业是 96 个立方体(6144 个芯片)作业。
下表展示了 v5p 的主要规范。
主要规范 | v5p 值 |
---|---|
每个芯片的峰值计算能力 (bf16) | 459 TFLOPs |
HBM2e 容量和带宽 | 95 GB、2765 GBps |
TPU Pod 大小 | 8960 个芯片 |
互连拓扑 | 3D 环面 * |
芯片间互联带宽 | 4800 Gbps |
配置
TPU v5p Pod 由 8960 个芯片组成,这些芯片通过可重配置的高速链路相互连接。TPU v5p 的灵活网络可让您通过多种方式连接相同大小的切片中的芯片。使用 gcloud compute tpus tpu-vm create
命令创建 TPU 切片时,您可以使用 AcceleratorType
参数指定其类型和形状。
下表展示了 v5p 支持的最常见单切片形状,以及大于 1 个立方体的大多数(但不是全部)完整立方体形状。v5p 的最大形状为 16x16x24(6144 个芯片、96 个立方体)。
拓扑 | 核心 | 芯片 | 主机 | 立方体 | 是否支持扭曲? |
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 | 不适用 |
单切片训练最多支持 6144 个芯片。您可以使用多切面扩容到最多 18432 个芯片。如需详细了解多切片,请参阅 Cloud TPU 多切片概览。
使用 AcceleratorType 参数
分配 TPU 资源时,您可以使用 --accelerator-type
参数指定切片中的 TensorCore 数量。--accelerator-type
是格式化字符串“v$VERSION_NUMBER
p-$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
以下命令会创建一个具有 256 个 v5p 芯片的 v5p TPU 切片,用于训练:
$ gcloud compute tpus tpu-vm create your-tpu-name \ --zone=us-east5-a \ --accelerator-type=v5p-256 \ --version=v2-alpha-tpuv5
如需详细了解如何管理 TPU,请参阅管理 TPU。如需详细了解 Cloud TPU 的系统架构,请参阅系统架构。
Cloud TPU ICI 弹性
ICI 弹性有助于提高光学链路和光学电路开关 (OCS) 的容错能力,用于在立方体之间连接 TPU。(立方体内的 ICI 连接使用不受影响的铜链路)。借助 ICI 弹性,ICI 连接可绕过 OCS 和光学 ICI 故障。因此,它会提高 TPU 切片的调度可用性,但会导致 ICI 性能暂时下降。
与 Cloud TPU v4 类似,对于大小为一个立方体或更大的 v5p 切片,ICI 弹性默认处于启用状态:
- 指定加速器类型时为 v5p-128
- 指定加速器配置时为 4x4x4
虚拟机、主机和切片属性
属性 | TPU 中的值 |
v5p 芯片数量 | 4 |
vCPU 的数量 | 208(如果使用 NUMA 绑定来避免跨 NUMA 性能损失,则只有一半可用) |
RAM (GB) | 448(如果使用 NUMA 绑定以避免跨 NUMA 性能损失,则只有一半可用) |
NUMA 节点数 | 2 |
NIC 吞吐量 (Gbps) | 200 |
一个 Pod 中的 TensorCore、芯片、主机/虚拟机和立方体数量之间的关系:
核心数 | 芯片 | 主机/虚拟机 | 立方体 | |
---|---|---|---|---|
主机 | 8 | 4 | 1 | |
立方体(又名机架) | 128 | 64 | 16 | 1 |
支持的最大切片 | 12288 | 6144 | 1536 | 96 |
v5p 完整 Pod | 17920 | 8960 | 2240 | 140 |