Google Cloud Platform TPU 是 Google 定制设计的 AI 加速器，针对大型 AI 模型的训练和使用进行了优化。它们旨在针对各种 AI 工作负载进行经济高效的扩缩，并提供多种功能，以加速 AI 框架（包括 PyTorch、JAX 和 TensorFlow）上的推理工作负载。如需详细了解 TPU，请参阅 Google Cloud Platform TPU 简介。
在 Dataflow 中使用 TPU 的前提条件
- 您的 Google Cloud 项目必须获得批准才能使用此正式版产品。
限制
此服务存在以下限制：
- 仅支持单主机 TPU 加速器：Dataflow TPU 产品仅支持单主机 TPU 配置，其中每个 Dataflow 工作器管理一个或多个 TPU 设备，这些设备不会与其他工作器管理的 TPU 互连。
- 仅支持同构 TPU 工作器池：Dataflow 适配和 Dataflow Prime 等功能不支持 TPU 工作负载。
价格
使用 TPU 的 Dataflow 作业按所消耗的工作器 TPU 芯片小时数计费，而不按工作器 CPU 和内存计费。如需了解详情，请参阅 Dataflow 价格页面。
可用性
以下 TPU 加速器和处理区域可供使用。
支持的 TPU 加速器
受支持的 TPU 加速器组合通过元组（TPU 类型、TPU 拓扑）来标识。
- TPU 类型是指 TPU 设备的型号。
- TPU 拓扑是指切片中 TPU 芯片的数量和物理排列方式。
如需为 Dataflow 工作器配置 TPU 的类型和拓扑，请使用格式为 type:TPU_TYPE;topology:TPU_TOPOLOGY 的 worker_accelerator 流水线选项。
Dataflow 支持以下 TPU 配置：
| TPU 类型 | 拓扑 | worker_machine_type 为必需项 |
|---|---|---|
| tpu-v5-lite-podslice | 1x1 | ct5lp-hightpu-1t |
| tpu-v5-lite-podslice | 2x2 | ct5lp-hightpu-4t |
| tpu-v5-lite-podslice | 2x4 | ct5lp-hightpu-8t |
| tpu-v6e-slice | 1x1 | ct6e-standard-1t |
| tpu-v6e-slice | 2x2 | ct6e-standard-4t |
| tpu-v6e-slice | 2x4 | ct6e-standard-8t |
| tpu-v5p-slice | 2x2x1 | ct5p-hightpu-4t |
区域
如需了解 TPU 的可用区域和可用区，请参阅 Cloud TPU 文档中的 TPU 区域和可用区。