TPU 虚拟机映像
创建 TPU 资源时,您需要传递 --version
或 --runtime-version
参数,用于指定 TPU 虚拟机映像。TPU 虚拟机映像包含 Ubuntu 操作系统、Docker 以及在 TPU 上运行代码所需的其他软件。本文档提供了有关在创建 Cloud TPU 时选择适当 TPU 虚拟机映像的指南。
PyTorch 和 JAX
使用以下常见的 TPU VM 基础映像来安装 PyTorch 和 JAX,然后安装要使用的框架。
- tpu-ubuntu2204-base(默认)
- v2-alpha-tpuv6e(TPU v6e)
- v2-alpha-tpuv5(TPU v5p)
- v2-alpha-tpuv5-lite(TPU v5e)
如需了解安装说明,请参阅 PyTorch/XLA 和 JAX 的快速入门文档。
TensorFlow
每个版本的 TensorFlow 都有专用的 TPU VM 映像。Cloud TPU 支持以下 TensorFlow 版本:
- 2.18.0
- 2.17.0
- 2.16.2
- 2.16.1
- 2.15.1
- 2.15.0
- 2.14.1
- 2.14.0
- 2.13.1
- 2.13.0
- 2.12.1
- 2.12.0
- 2.11.1
- 2.11.0
- 2.10.1
- 2.10.0
- 2.9.3
- 2.9.1
- 2.8.4
- 2.8.3
- 2.8.0
- 2.7.4
- 2.7.3
如需详细了解 TensorFlow 补丁版本,请参阅支持的 TensorFlow 补丁版本。
对于 TensorFlow 2.15.0 及更高版本,有基于您使用的设备 API(PJRT 或流执行器)的 TPU VM 映像变体。
在 v6e、v5p 和 v5e 上进行训练
TPU v6e、v5e 和 v5p 支持 TensorFlow 2.15.0 及更高版本。您可以使用以下形式指定 TPU VM 映像:tpu-vm-tf-x.y.z-{pod}-pjrt
,其中 x
是 TensorFlow 的主要版本,y
是次要版本,z
是 TensorFlow 补丁版本。如果您使用的是多主机 TPU,请在 TensorFlow 版本后面添加 pod
。例如,如果您在多主机 TPU 上使用 TensorFlow 2.16.0,请使用 tpu-vm-tf-2.16.0-pod-pjrt
TPU 虚拟机映像。对于其他版本的 TensorFlow,请将 2.16.0
替换为您所使用的 TensorFlow 的主要版本和补丁版本。如果您使用的是单主机 TPU,请省略 pod
。
在 v6e 和 v5e 上提供
有可供部署的 Docker 映像,其中包含使用 TensorFlow、PyTorch 和 JAX 进行部署所需的所有软件要求。如需了解详情,请参阅 Cloud TPU v5e 推理简介。
TPU v4
如果您使用的是 TPU v4 和 TensorFlow 2.15.0 或更高版本,请按照在 v6e、v5p 和 v5e 上训练的说明操作。如果您使用的是 TensorFlow 2.10.0 或更低版本,请使用专门针对 v4 的 TPU 虚拟机映像:
TensorFlow 版本 | TPU 虚拟机映像版本 |
---|---|
2.10.0 | tpu-vm-tf-2.10.0-v4 tpu-vm-tf-2.10.0-pod-v4 |
2.9.3 | tpu-vm-tf-2.9.3-v4 tpu-vm-tf-2.9.3-pod-v4 |
2.9.2 | tpu-vm-tf-2.9.2-v4 tpu-vm-tf-2.9.2-pod-v4 |
2.9.1 | tpu-vm-tf-2.9.1-v4 tpu-vm-tf-2.9.1-pod-v4 |
TPU v2 和 v3
如果您使用的是 TPU v2 或 v3,请使用与您使用的 TensorFlow 版本匹配的 TPU 虚拟机映像。例如,如果您使用的是 TensorFlow 2.14.1,请使用 tpu-vm-tf-2.14.1
TPU 映像。对于其他版本的 TensorFlow,请将 2.14.1
替换为您使用的 TensorFlow 版本。如果您使用的是多主机 TPU,请将 Pod 附加到 TPU 映像的末尾,例如 tpu-vm-tf-2.14.1-pod
。
从 TensorFlow 2.15.0 开始,您还必须在映像名称中指定设备 API。例如,如果您将 TensorFlow 2.16.1 与 PJRT API 搭配使用,请使用 TPU 映像 tpu-vm-tf-2.16.1-pjrt
。如果您将流式执行器 API 与相同版本的 TensorFlow 搭配使用,请使用 tpu-vm-tf-2.16.1-se
TPU 映像。低于 2.15.0 的 TensorFlow 版本仅支持流执行器。
对 TensorFlow PJRT 的支持
从 TensorFlow 2.15.0 开始,您可以将 PJRT 接口用于 TPU 上的 TensorFlow。PJRT 具有自动设备内存碎片整理功能,并简化了硬件与框架的集成。如需详细了解 PJRT,请参阅 PJRT:简化机器学习硬件和框架集成。
加速器 | 功能 | PJRT 支持 | 流执行器支持 |
---|---|---|---|
TPU v2 - v4 | 密集计算(无 TPU 嵌入 API) | 是 | 是 |
TPU v2 - v4 | 密集计算 API + TPU 嵌入 API | 否 | 是 |
TPU v2 - v4 | 使用软设备放置的 tf.summary/tf.print | 否 | 是 |
TPU v5e | 密集计算(无 TPU 嵌入 API) | 是 | 否 |
TPU v5e | TPU Embedding API | 不适用 | 否 |
TPU v5p | 密集计算(无 TPU 嵌入 API) | 是 | 否 |
TPU v5p | TPU Embedding API | 是 | 否 |
libtpu 版本
TPU VM TensorFlow 映像包含特定的 TensorFlow 版本和相应的 libtpu 库。如果您要创建自己的虚拟机映像,请使用以下 TensorFlow TPU 软件版本和相应的 libtpu 版本:
TensorFlow 版本 | libtpu.so 版本 |
---|---|
2.18.0 | 1.12.0 |
2.17.0 | 1.11.0 |
2.16.2 | 1.10.1 |
2.16.1 | 1.10.1 |
2.15.1 | 1.9.0 |
2.15.0 | 1.9.0 |
2.14.1 | 1.8.1 |
2.14.0 | 1.8.0 |
2.13.1 | 1.7.1 |
2.13.0 | 1.7.0 |
2.12.1 | 1.6.1 |
2.12.0 | 1.6.0 |
2.11.1 | 1.5.1 |
2.11.0 | 1.5.0 |
2.10.1 | 1.4.1 |
2.10.0 | 1.4.0 |
2.9.3 | 1.3.2 |
2.9.1 | 1.3.0 |
2.8.3 | 1.2.3 |
2.8.0 | 1.2.0 |
2.7.3 | 1.1.2 |