张量处理单元 (TPU) 是 Google 定制开发的应用专用集成电路 (ASIC)，旨在加速机器学习 (ML) 和人工智能 (AI) 工作负载。无论您是要花费数周时间训练复杂的基础模型，还是要运行大规模推理，TPU 都能提供可伸缩的专用计算资源，这些资源针对 JAX 和 PyTorch 等框架进行了优化。
Cloud TPU 专为应对要求最严苛的 AI 工作负载而设计。主要优势包括：
针对矩阵计算进行了优化：TPU 专门设计有矩阵乘法单元 (MXU)，可高效执行机器学习算法基础的大规模矩阵运算。
高带宽内存 (HBM)：片上高带宽内存可让您训练和提供更大的模型，并有效利用更大的批次大小。
通过切片实现大规模可伸缩性：TPU 芯片可以分组连接，称为切片。这些切片可让您的工作负载扩缩到多达数千个 TPU 芯片，以执行大规模训练作业。
何时使用 TPU
TPU 针对特定工作负载进行了优化，例如：
- 由矩阵计算主导的模型
- 在主训练循环内没有自定义 PyTorch/JAX 操作的模型
- 需要训练数周或数月的模型
- 有效批量大小较大的大型模型
- 高级排名和推荐工作负载中常见的具有超大嵌入的模型
TPU 不适合以下工作负载：
- 需要频繁分支或包含许多元素级代数运算的线性代数程序
- 需要高精度算法的工作负载
- 主训练循环中包含自定义操作的神经网络工作负载
Google Cloud上的预配选项
您可以根据运营需求，使用以下 Google Cloud 产品访问和预配 TPU。
Compute Engine
借助 Compute Engine，您可以创建和管理单个 TPU 虚拟机或切片，从而能够对 TPU 虚拟机进行完整的生命周期管理。Google 建议您使用 Compute Engine（而非旧版 Cloud TPU API）来预配 TPU 资源。
如需了解详情，请参阅 Compute Engine 中的 Cloud TPU 资源。
Google Kubernetes Engine
Google Kubernetes Engine (GKE) 提供全代管式式多租户 Kubernetes 环境，用于编排大规模 AI 工作负载。GKE 支持 TPU 节点和节点池生命周期管理，包括创建、配置和删除 TPU 虚拟机。
如需了解详情，请参阅 GKE 中的 TPU 简介。
批量
Batch 是一项全代管式服务，可让您安排、排队和执行在 Compute Engine TPU 虚拟机上运行的容器化或基于脚本的工作负载。
如需了解详情，请参阅创建和运行使用 TPU 的批量作业。
Cloud TPU
Cloud TPU API（包括 Google Cloud CLI 和 Cloud TPU 的 Cloud 客户端库）已不再开发。对于预配和管理 TPU 资源，Google 建议您根据编排和工作负载需求使用 Compute Engine 或 GKE。
如需了解详情，请参阅从 Cloud TPU API 迁移。
Compute Engine 支持的 TPU 版本
Compute Engine 支持以下 TPU 版本：
- TPU7x (Ironwood)
- TPU v6e (Trillium)
- TPU v5p
如需详细了解每个 TPU 版本，请参阅 TPU 机器。