加速器优化机器家族中的 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 计算实例，可提供对高性能库、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/机器学习推理

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

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

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

    • 所需的 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 2

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

2 TPU7x 的 Spot、Flex-start、日历模式下的未来预留以及容量块的未来预留均受许可清单限制。如需申请访问权限，请与您的客户支持团队或销售团队联系。

全容量模式

默认情况下，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 由 8,960 个 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 芯片数量：1024
  • 主机数量：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. 计算方法是用拓扑产品除以四。↩

后续步骤