加速器优化机器家族中的 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 主机上运行的 Linux 计算实例，可直接访问底层 TPU 硬件，并提供对高性能库、SSH 连接和运行时调试日志的访问权限。

  • TPU 切片：一组互连的 TPU 芯片，可通过一个或多个 TPU 实例进行访问。切片的范围可以是以下其中之一：

    • 单宿主机切片：包含一台宿主机。一般来说，这会映射到一个 TPU 实例。
    • 多主机切片：由多个 TPU 实例组成的切片，这些实例通过高速芯片间互连 (ICI) 进行互连。

  • TPU 立方体：一种由互连的 TPU 芯片组成的 4x4x4 拓扑。这仅适用于 3D 拓扑。

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

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

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

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

选择合适的 TPU 型号取决于工作负载的计算特征、内存要求和可伸缩性需求。选择以下任一标签页即可查看建议。

AI/机器学习训练

  • 预训练大型模型：从头开始训练大型模型。 这些工作负载受计算和通信的限制。

    • 所需的 TPU 功能：高峰值 FLOPS、低精度支持和高互连带宽。
    • 建议使用的 TPU 版本：
      • TPU7x。此版本可提供出色的 FP8 性能。
      • TPU v5p。此版本支持大规模集群配置。

  • 微调和提炼模型：使用参数高效的方法或训练较小的模型变体来调整现有模型。

    • 所需的 TPU 功能：中等计算能力和充足的内存容量。
    • 建议使用的 TPU 版本：
      • TPU v6e。此版本可经济高效地完成标准训练任务。
      • TPU7x。此版本可为大规模微调提供高性能。

AI/机器学习推理

  • 在线提供模型：部署模型以实现实时互动，其中低延迟至关重要。在按顺序生成令牌期间，这些工作负载会受到内存带宽的限制。

    • 所需的 TPU 功能：高 HBM 带宽和大型片上 SRAM，以最大限度地减少数据检索延迟。
    • 建议使用的 TPU 版本：
      • TPU v6e。此版本经过优化，可经济高效地大规模提供服务。
      • TPU7x。此版本提供高 HBM 带宽和大 SRAM，以加速令牌生成。

  • 运行批次推理：离线处理大型数据集，主要目标是实现高吞吐量。在提示预填充阶段，这些工作负载受计算限制。

    • 所需的 TPU 功能：高计算密度，可最大限度地提高每美元的吞吐量。
    • 建议使用的 TPU 版本：
      • TPU v6e。此版本可实现高吞吐量批处理，经济实惠。
      • TPU7x。此版本提供高计算密度，可快速处理大批量数据。

推荐系统

  • 处理大型嵌入：训练和提供使用大型嵌入表的推荐模型。这些工作负载受容量和通信限制。
    • 所需的 TPU 功能：专用 SparseCore 硬件，用于并行处理稀疏操作；大 HBM 容量；支持卸载到主机内存。
    • 建议使用的 TPU 版本：
      • TPU7x。此版本包含每个芯片四个 SparseCore，并提供最大的 HBM 容量。
      • TPU v6e。此版本包含每个芯片两个 SparseCore，可实现经济高效的投放。
      • TPU v5p。此版本包含每个芯片四个 SparseCore，可用于大规模训练。

强化学习

  • 运行强化学习循环：管理需要紧密循环生成样本和更新政策权重的混合工作负载。

    • 所需的 TPU 功能：低延迟 ICI，用于在芯片之间快速传输权重和激活，以及高带宽主机连接。

    • 推荐的 TPU 版本：

      这些 TPU 版本使用高速 3D 环面互连。

使用选项

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

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

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

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

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

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

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

如需了解支持此使用模式的底层预配模型，请参阅虚拟机预配模型简介。

各 TPU 版本的消费选项可用性

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

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

1 使用“所有容量”模式时，您无法在日历模式下创建未来预留请求。

2 Spot、灵活启动和日历模式下的未来预留仅适用于列入许可名单的 TPU7x。如需申请访问权限，请与您的客户支持团队或销售团队联系。

全容量模式

默认情况下，TPU 预留包含自动故障恢复功能。Google Cloud 通过自动将有故障的 TPU 机器替换为预留容量中的正常机器来管理此流程，以确保有可用的资源来重启 TPU 切片。

或者，对于 TPU v6e 和 TPU7x，您可以在全容量模式下请求 TPU 容量。此模式可授予对所有预留容量的访问权限，并提供完整的拓扑可见性。在此模式下，您还可以定位到特定块或子块，以实现精确的拓扑感知型工作负载布置。不过，您需要负责管理硬件故障和维护。

如需详细了解这些模式，请参阅 Cloud TPU 文档中的 TPU 容量模式。

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 带宽 (GBps) 7380 1638 2765
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 虚拟机包含 4 个 TPU 芯片。所有 TPU7x 切片均采用整机配置，由 4 芯片虚拟机组成。

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

借助 Ironwood 编程模型，您可以调用两个 TPU 设备，而不再受限于前几代产品所采用的单逻辑核心架构。如需了解详情，请参阅 Cloud TPU 文档中的双芯片组架构。

机器类型 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. 计算方法是用拓扑产品除以四。↩

后续步骤