Cloud de Confiance by S3NS 上的 TPU 简介

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

Cloud TPU 旨在处理要求最严苛的 AI 工作负载。主要优势包括:

  • 针对矩阵计算进行了优化:TPU 专门配备了矩阵乘法单元 (MXU),能够以极高的效率执行对机器学习算法至关重要的大规模矩阵运算。

  • 高带宽内存 (HBM):借助片上高带宽内存,您可以训练和提供更大的模型,并有效利用更大的批次大小。

  • 通过切片实现大规模可伸缩性:TPU 芯片可以连接在名为切片的组中。借助切片,您的工作负载可将规模扩展到数千个 TPU 芯片,以处理大规模训练作业。

何时使用 TPU

TPU 针对特定工作负载进行了优化,例如:

  • 由矩阵计算主导的模型
  • 在主训练循环内没有自定义 PyTorch/JAX 操作的模型
  • 需要训练数周或数月的模型
  • 有效批量大小较大的大型模型
  • 在高级排名和推荐工作负载中常见的具有超大嵌入的模型

TPU 适合以下工作负载:

  • 需要频繁分支或包含许多元素级代数运算的线性代数程序
  • 需要高精度算法的工作负载
  • 主训练循环中包含自定义操作的神经网络工作负载

Cloud de Confiance by S3NS上的预配选项

您可以根据自己的运营需求,使用以下 Cloud de Confiance by S3NS 产品 来访问和预配 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 简介

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 机器

后续步骤