了解 TPU

张量处理单元 (TPU) 是 Google 定制开发的应用专用集成电路 (ASIC)，旨在加速机器学习 (ML) 和人工智能 (AI) 工作负载。无论您是训练复杂的基础模型数周，还是运行大规模推理，TPU 都能提供可伸缩的专用计算资源，这些资源针对高级 AI 框架进行了优化。

本文档介绍了 TPU 在 Google Cloud中的作用、每个可用 TPU 版本的技术规范、性能特征、价格注意事项，以及它们与 Compute Engine 机器系列的对应关系。

什么是 TPU？

TPU 是一种定制设计的 ASIC，旨在满足机器学习算法中核心的大规模矩阵运算的需求。

TPU 系统的关键架构组件包括：

  • TensorCore：TPU 芯片的处理单元。每个 TensorCore 都包含一个或多个矩阵乘法单元 (MXU)，这些单元使用脉动阵列架构高效执行矩阵乘法，同时还包含用于控制流和通用计算的标量单元和向量单元。

  • SparseCore：专为加速稀疏运算而设计的数据流处理器。SparseCore 可与 TensorCore 协同工作，高效处理大型嵌入表，因此非常适合加速推荐模型和大规模嵌入。

  • 高带宽内存 (HBM)：连接到 TPU 芯片的大容量物理内存，可提供极高的内存带宽。HBM 支持以更大的批次大小训练和提供更大的模型。

  • TPU 实例：一种在 TPU 主机上运行的 Linux 计算实例，可直接访问底层 TPU 硬件，并提供对高性能库、SSH 连接和运行时调试日志的访问权限。

  • TPU 切片和 Pod：TPU Pod 是一个由物理链接的 TPU 芯片组成的庞大连续集群。TPU 切片是 Pod（包含一个或多个通过高速互连连接的主机）的逻辑分区，用于运行单个工作负载。

  • XLA 编译器：加速线性代数 (XLA) 编译器将机器学习图编译为在 TPU TensorCore 上运行的优化机器指令。

TPU 版本和机器系列

Compute Engine 支持多个代系和版本的 TPU。下表提供了每个受支持的 TPU 版本的关键规范：

TPU 版本 机器系列 单芯片峰值算力 (TFLOPs) TPU 内存和带宽 每个芯片的核心数 单芯片互连 (ICI) 带宽 每个芯片的网络带宽 (DCN) 每个 Pod 的芯片数
TPU7x (Ironwood) tpu7x-standard

BF16：2,307

FP8：4,614

 192 GiB（7,380 GBps）

2 个 TensorCore

4 个 SparseCore

 1,200 GBps 100 Gbps 9,216
TPU v6e (Trillium) ct6e-standard

BF16：918

FP8：918

 32 GiB（1,638 GBps）

1 个 TensorCore

2 个 SparseCore

 800 GBps 100 Gbps 256
TPU v5p ct5p-hightpu

BF16：459

FP8：459

 95 GiB (2,765 GBps)

2 个 TensorCore

4 个 SparseCore

 1,200 GBps 50 Gbps 8,960

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 环面互连。

工作负载注意事项

对于以下工作负载，TPU 可能不是最佳选择：

  • 需要频繁的逻辑分支或包含许多元素级代数运算的线性代数程序。
  • 依赖于在 CPU 上运行的 JAX 或 PyTorch 中的自定义运算符的工作负载。
  • 需要双精度浮点 (FP64) 运算的科学工作负载。

价格注意事项

以下因素决定了 Google Cloud上的 TPU 价格：

  • TPU 版本

  • 位置

  • 使用模式

如需了解完整的价格详情，请参阅 TPU 价格页面。

使用和购买模式

您可以根据可用性要求和费用容忍度，使用以下任一消费选项来预配 TPU 实例：

  • 按需：标准随用随付 (PAYG) 价格，可立即执行。可提供最大的生命周期灵活性，但可用性受物理资源限制。
  • Spot：以极低的价格（最高可节省 70%）使用过剩的 Google Cloud 容量。由于 Google Cloud 可能会在提前 30 秒通知的情况下停止或删除 Spot 虚拟机来回收容量，因此这些计算实例非常适合容错工作负载和设置了检查点的工作负载。
  • 灵活启动：可从专用池请求计算实例，最多可运行 7 天，可用性更高。
  • 未来预留：可保障容量并提供比按需价格更低的折扣价格的预留。预留指定未来日期和时间的 TPU 容量，预留时间最长为 90 天（日历模式），或预留 1 年或更长时间。如果Google Cloud 批准了您的请求，您就可以在预留期内开始使用相应容量。

承诺使用折扣 (CUD)

当您请求预留 1 年或更长时间的未来预留时，可以为预留的 TPU 容量获得 CUD。将 CUD 应用于计算实例后，您只需承诺使用一段时间，即可享受远低于按需费率的价格。

如需了解详情，请参阅 Compute Engine 的 CUD。

后续步骤