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_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

以下命令会创建一个具有 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