TPU 硬件支持

Optimum-TPU 支持并针对 v5e 和 v6e TPU 进行了优化。

TPU 命名约定

TPU 命名遵循以下格式：<tpu_version>-<number_of_tpus>

TPU 版本

• v5litepod (v5e)
• v6e

例如，v5litepod-8 是一个具有 8 个 tpu 的 v5e TPU。

TPU 上的内存

v5e、v5p 的每个芯片的 HBM（高带宽内存）容量为 16GB，v6e 的为 32GB。因此，v5e-8 (v5litepod-8) 具有 16GB*8=128GB 的 HBM 内存

TPU 的推荐运行时

在 TPU VM 创建期间，使用以下 TPU VM 基本镜像用于 optimum-tpu

• v2-alpha-tpuv6e (TPU v6e)（推荐）
• v2-alpha-tpuv5 (TPU v5p)（推荐）
• v2-alpha-tpuv5-lite (TPU v5e)（推荐）
• tpu-ubuntu2204-base (默认)

有关安装说明，请参阅我们的TPU 设置教程。我们建议您将alpha 版本与 optimum-tpu 一起使用，因为 optimum-tpu 经过测试并针对这些版本进行了优化。

更多信息请访问 https://cloud.google.com/tpu/docs/runtimes#pytorch_and_jax

下一步

有关不同 TPU 硬件的更多信息，您可以查看： https://cloud.google.com/tpu/docs/v6e https://cloud.google.com/tpu/docs/v5p https://cloud.google.com/tpu/docs/v5e

定价信息可以在这里找到 https://cloud.google.com/tpu/pricing

TPU 可用性可以在这里找到 https://cloud.google.com/tpu/docs/regions-zones