加速器优化机器家族中的 TPU 机器

本文档介绍了加速器优化机器家族中具有张量处理单元 (TPU) 的 Compute Engine 实例。TPU 是 Google 定制开发的应用专用集成电路 (ASIC),专门针对人工智能 (AI) 和机器学习 (ML) 工作负载进行了优化。

Compute Engine 支持以下 TPU 版本:

  • TPU7x
  • TPU v6e
  • TPU v5p

每个版本中的每种机器类型都有特定的拓扑和所连接的 TPU 芯片数量。

TPU 架构基础知识

了解 TPU 架构的基础知识有助于您为工作负载选择合适的 TPU 版本和机器类型。

  • TPU 芯片:TPU 芯片是 Google 专为机器学习设计的专用加速器。每个 TPU 芯片都包含一个或多个 TensorCore,用于处理大规模矩阵运算。每个 TensorCore 都包含一个或多个矩阵乘法单元 (MXU),这些单元使用脉动阵列架构,可在每个周期内执行数千次乘法累加运算,而无需持续访问内存。虽然 TPU 芯片主要用于高速矩阵处理,但它还包含向量单元和标量单元,用于执行一般计算和控制流操作。

  • TPU Pod:TPU Pod 是指通过专用网络组合在一起的一组连续的 TPU。TPU Pod 中的 TPU 芯片数量取决于 TPU 版本。

  • TPU 虚拟机:TPU 虚拟机是一种在 TPU 主机上运行的 Linux 虚拟机,可访问底层 TPU。您可以使用 SSH 直接连接到 TPU 虚拟机。您拥有虚拟机的根访问权限,因此可以运行任意代码。您可以访问编译器和运行时调试日志及错误消息。

  • TPU 切片:一组互连的 TPU 芯片,可通过一个或多个 TPU 虚拟机进行访问。切片的范围可以是以下其中之一:

    • 单宿主机切片:由一台宿主机组成的切片。一般来说,这会映射到 1 个 TPU 虚拟机。
    • 多主机切片:由多个 TPU 虚拟机组成的切片,这些虚拟机通过高速芯片间互连 (ICI) 进行互连。
  • TPU 立方体:一种由互连的 TPU 芯片组成的 4x4x4 拓扑。此参数仅适用于 3D 拓扑。

  • SparseCore:SparseCore 是一种数据流处理器,可利用稀疏运算来加速模型。主要应用场景是加速高度依赖嵌入的推荐模型。

  • TPU 版本:TPU 芯片的确切架构取决于您使用的 TPU 版本。每个 TPU 版本还支持不同的切片大小和配置。

如需了解 TPU 的工作方式,请参阅 Cloud TPU 文档中的 TPU 架构文档。

按工作负载类型推荐的 TPU 版本

TPU 版本 主要工作负载类型
TPU7x (Ironwood)
  • 大规模密集模型和混合专家模型 (MoE)
  • 针对大规模基础模型进行密集预训练
  • 抽样和解码密集型推理
TPU v6e (Trillium)
  • 训练和微调(Transformer、CNN)
  • 大规模推理(Gemma 2、Llama、Diffusion 模型)
  • 推荐引擎和个性化(使用 SparseCore)
TPU v5p
  • 可实现大规模基础模型训练的最高性能
  • 大规模多模态 AI 训练
  • 嵌入密集型工作负载,例如大型推荐系统

使用选项

为了在平衡工作负载性能的同时优化资源利用率和费用,Compute Engine 支持以下 TPU 使用情况选项:

  • 按需:无需提前安排容量即可使用 TPU。在请求资源之前,您必须拥有足够的按需配额,以用于特定类型和数量的 TPU 虚拟机。按需是最灵活的使用选项;不过,我们无法保证有足够的按需资源可用于满足您的请求。

  • Spot 虚拟机:如需预配 Spot 虚拟机,您可以享受大幅折扣,但 Spot 虚拟机随时可能被抢占,并且会提前 30 秒发出警告。如需了解详情,请参阅 Spot 虚拟机简介

  • 灵活启动:预配最多可运行 7 天的灵活启动虚拟机,Compute Engine 会根据可用性尽力自动分配硬件。如需了解详情,请参阅灵活启动虚拟机简介

  • 未来预留:请求一年或更长时间的未来预留。如需了解详情,请参阅 Cloud TPU 文档中的请求一年或更长时间的未来预留

  • 日历模式下的未来预留:在指定时间段内预配 TPU 资源,预留时间最长为 90 天。如需了解详情,请参阅日历模式下的未来预留请求简介

如果您未指定其他选项,按需是 TPU 的默认使用模式。

如需了解支持此消费选项的底层预配模型,请参阅虚拟机预配模型简介

使用方案可用性(按 TPU 版本)

下表总结了每种使用选项的可用性(按 TPU 版本)。

TPU 版本 按需 Spot 灵活启动 按需预留 未来预留 日历模式下的未来预留
1 1 1

1 Spot、灵活启动和日历模式下的未来预留仅限 TPU7x 允许名单中的用户使用。如需申请访问权限,请与您的客户支持团队销售团队联系。

TPU 版本比较

比较不同 TPU 版本的特征。 您可以在选择要比较的属性字段中选择特定属性,以比较下表中所有 TPU 版本的这些属性。

加速器优化 加速器优化 加速器优化
虚拟机 虚拟机 虚拟机
Intel Emerald Rapids AMD EPYC Genoa Intel Sapphire Rapids
x86 x86 x86
224 44 - 180 208
消息串 消息串 消息串
960 GB 176 到 1440 GB 448 GB
NUMA NUMA NUMA
NVMe NVMe NVMe
gVNIC gVNIC gVNIC
400 Gbps 50 到 400 Gbps 200 Gbps
4 8 4
折扣 折扣 折扣
折扣 折扣 折扣

TPU 架构规范

下表列出了每个 TPU 版本的主要规范。

规范 TPU7x TPU v6e TPU v5p
每个 pod 的芯片数量 9216 256 8960
每个芯片的峰值计算能力 (BF16) (TFLOPs) 2307 918 459
每个芯片的峰值计算能力 (FP8) (TFLOPs) 4614 918 459
每个芯片的 HBM 容量 (GiB) 192 32 95
每个芯片的 HBM 带宽 (GiBps) 7380 1638 2575
vCPU 数量(4 芯片虚拟机) 224 180 208
RAM (GiB)(4 芯片虚拟机) 960 720 448
每个芯片的 TensorCore 数量 2 1 2
每个芯片的 SparseCore 数量 4 2 4
每个芯片的双向芯片间互连 (ICI) 带宽 (GBps) 1200 800 1200
每个芯片的数据中心网络 (DCN) 带宽 (Gbps) 100 100 50

TPU 机器类型

以下部分介绍了每种 TPU 版本可用的机器类型。

TPU7x (Ironwood)

每个 TPU7x 虚拟机 (VM) 包含 4 个 TPU 芯片。所有 TPU7x 切片都使用全主机、4 芯片虚拟机。

每个 TPU7x 芯片包含两个 TensorCore 和四个 SparseCore。

借助 Ironwood 编程模型,您可以访问两个 TPU 设备,而不是之前世代中使用的单个逻辑核心架构。如需了解详情,请参阅 Cloud TPU 文档中的双 chiplet 架构

机器类型 vCPU 的数量 实例内存 (GiB) 物理 NIC 数量 最大网络带宽 (Gbps) 每个虚拟机的 TPU 芯片数量 NUMA 节点的数量 TPU 总内存(GiB HBM)
tpu7x-standard-4t 224 960 2 400 4 2 768

如需详细了解 TPU7x 架构,请参阅 Cloud TPU 文档中的 TPU7x (Ironwood)

TPU v6e (Trillium)

每个 TPU v6e 虚拟机可以包含 1、4 或 8 个 TPU 芯片。芯片数量不超过 4 块的切片使用相同的非统一内存访问 (NUMA) 节点。

v6e 切片是使用半主机虚拟机创建的,每个虚拟机有 4 个 TPU 芯片,但以下情况除外:

  • 仅包含单个 TPU 芯片的 ct6e-standard-1t 主要用于测试。
  • ct6e-standard-8t 是一种针对推理应用场景进行了优化的全主机虚拟机,可在一个部署工作负载中使用连接到单个虚拟机的全部 8 个 TPU 芯片。
机器类型 vCPU 的数量 实例内存 (GB) 物理 NIC 数量 最大网络带宽 (Gbps) 每个虚拟机的 TPU 芯片数量 NUMA 节点的数量 TPU 总内存(GiB HBM)
ct6e-standard-1t 44 176 1/4 50 1 1 32
ct6e-standard-4t 180 720 2 400 4 1 128
ct6e-standard-8t 360 1440 1 200 8 2 256

如需详细了解 TPU v6e 架构,请参阅 Cloud TPU 文档中的 TPU v6e

TPU v5p

TPU v5p Pod 由 8960 个 TPU 芯片组成,这些芯片通过可重配置的高速链路相互连接。TPU v5p 的灵活联网功能可让您以多种方式连接相同大小的切片中的 TPU 芯片。单切片训练最多支持 6,144 个 TPU 芯片。

机器类型 vCPU 的数量 实例内存 (GB) 物理 NIC 数量 最大网络带宽 (Gbps) 每个虚拟机的 TPU 芯片数量 NUMA 节点的数量 TPU 总内存(GiB HBM)
ct5p-hightpu-4t 208 448 1 200 4 2 380

如需详细了解 TPU v5p 架构,请参阅 Cloud TPU 文档中的 TPU v5p

TPU 拓扑

拓扑定义了 TPU 切片内 TPU 的物理排列方式。拓扑是二维或三维拓扑,具体取决于 TPU 版本。您可以通过计算拓扑中各个尺寸的乘积来确定切片中的 TPU 芯片数量。例如:

  • 具有 2x2x2 拓扑的 tpu7x-standard-4t 机器类型是 8 芯片多主机 TPU7x 切片。

下表列出了每个 TPU 版本可用的拓扑。

TPU 版本 机器类型 范围 技术规范
TPU7x (Ironwood) tpu7x-standard-4t 单个主机
  • 拓扑:2x2x1
  • 拓扑的 TPU 芯片数量:4
  • 主机数量:1
  • 虚拟机数量:1
  • 立方体数量:1/16
TPU7x (Ironwood) tpu7x-standard-4t 多主机
  • 拓扑:2x2x2
  • 拓扑的 TPU 芯片数量:8
  • 主机数量:2
  • 虚拟机数量:2
  • 立方体数量:1/8
TPU7x (Ironwood) tpu7x-standard-4t 多主机
  • 拓扑:2x2x4
  • 拓扑的 TPU 芯片数量:16
  • 主机数量:4
  • 虚拟机数量:4
  • 立方体数量:1/4
TPU7x (Ironwood) tpu7x-standard-4t 多主机
  • 拓扑:2x4x4
  • 拓扑的 TPU 芯片数量:32
  • 主机数量:8
  • 虚拟机数量:8
  • 立方体数量:1/2
TPU7x (Ironwood) tpu7x-standard-4t 多主机
  • 拓扑:4x4x4
  • 拓扑的 TPU 芯片数量:64
  • 主机数量:16
  • 虚拟机数量:16
  • 立方体数量:1
TPU7x (Ironwood) tpu7x-standard-4t 多主机
  • 拓扑:4x4x8
  • 拓扑的 TPU 芯片数量:128
  • 主机数量:32
  • 虚拟机数量:32
  • 立方体数量:2
TPU7x (Ironwood) tpu7x-standard-4t 多主机
  • 拓扑:4x8x8
  • 拓扑的 TPU 芯片数量:256
  • 主机数量:64
  • 虚拟机数量:64
  • 立方体数量:4
TPU7x (Ironwood) tpu7x-standard-4t 多主机
  • 拓扑:8x8x8
  • 拓扑的 TPU 芯片数量:512
  • 主机数量:128
  • 虚拟机数量:128
  • 立方体数量:8
TPU7x (Ironwood) tpu7x-standard-4t 多主机
  • 拓扑:8x8x16
  • 拓扑的 TPU 芯片数量:1,024
  • 主机数量:256
  • 虚拟机数量:256
  • 立方体数量:16
TPU7x (Ironwood) tpu7x-standard-4t 多主机
  • 拓扑:{A}x{B}x{C}(其中 A、B 和 C 是 2 的倍数)
  • 相应拓扑的 TPU 芯片数量:A*B*C
  • 主机数量:(A*B*C)/4
  • 虚拟机数量:(A*B*C/4)
  • 立方体数量:(A*B*C/64)
TPU v6e (Trillium) ct6e-standard-1t 单个主机
  • 拓扑:1x1
  • 拓扑的 TPU 芯片数量:1
  • 虚拟机数量:1
TPU v6e (Trillium) ct6e-standard-8t 单个主机
  • 拓扑:2x4
  • 拓扑的 TPU 芯片数量:8
  • 虚拟机数量:1
TPU v6e (Trillium) ct6e-standard-4t 单个主机
  • 拓扑:2x2
  • 拓扑的 TPU 芯片数量:4
  • 虚拟机数量:1
TPU v6e (Trillium) ct6e-standard-4t 多主机
  • 拓扑:2x4
  • 拓扑的 TPU 芯片数量:8
  • 虚拟机数量:2
TPU v6e (Trillium) ct6e-standard-4t 多主机
  • 拓扑:4x4
  • 拓扑的 TPU 芯片数量:16
  • 虚拟机数量:4
TPU v6e (Trillium) ct6e-standard-4t 多主机
  • 拓扑:4x8
  • 拓扑的 TPU 芯片数量:32
  • 虚拟机数量:8
TPU v6e (Trillium) ct6e-standard-4t 多主机
  • 拓扑:8x8
  • 拓扑的 TPU 芯片数量:64
  • 虚拟机数量:16
TPU v6e (Trillium) ct6e-standard-4t 多主机
  • 拓扑:8x16
  • 拓扑的 TPU 芯片数量:128
  • 虚拟机数量:32
TPU v6e (Trillium) ct6e-standard-4t 多主机
  • 拓扑:16x16
  • 拓扑的 TPU 芯片数量:256
  • 虚拟机数量:64
TPU v5p ct5p-hightpu-4t 单个主机
  • 拓扑:2x2x1
  • 拓扑的 TPU 芯片数量:4
  • 虚拟机数量:1
TPU v5p ct5p-hightpu-4t 多主机
  • 拓扑:2x2x2
  • 拓扑的 TPU 芯片数量:8
  • 虚拟机数量:2
TPU v5p ct5p-hightpu-4t 多主机
  • 拓扑:2x2x4
  • 拓扑的 TPU 芯片数量:16
  • 虚拟机数量:4
TPU v5p ct5p-hightpu-4t 多主机
  • 拓扑:2x4x4
  • 拓扑的 TPU 芯片数量:32
  • 虚拟机数量:8
TPU v5p ct5p-hightpu-4t 多主机
  • 拓扑:{A}x{B}x{C}(其中 A、B 和 C 是 2 的倍数)
  • 相应拓扑的 TPU 芯片数量:A*B*C
  • 虚拟机数量:(A*B*C/4)1
  1. 计算方法是用拓扑产品除以四。

后续步骤