TPU 软件版本
本文档介绍了如何在创建 Cloud TPU 时选择适当的 TPU 软件版本。
创建 TPU 资源时，您需要指定软件版本（也称为运行时版本），它是指在 TPU 虚拟机上预安装的软件环境。这包括 Ubuntu 操作系统、Docker 以及在 TPU 上运行代码所需的其他软件。
如果您使用的是 Google Cloud CLI，请使用 --version 或 --runtime-version 参数指定 TPU 软件版本。如果您使用的是 Google Cloud 控制台，请从 TPU 软件版本列表中选择 TPU 软件版本。
PyTorch 和 JAX
请使用以下适用于 PyTorch 和 JAX 的常见 TPU 软件版本，然后安装您要使用的框架。
| TPU 版本 | TPU 软件版本 |
|---|---|
| Trillium (v6e) | v2-alpha-tpuv6e |
| v5p | v2-alpha-tpuv5 |
| v5e | v2-alpha-tpuv5-lite |
| v4 及更低版本 | tpu-ubuntu2204-base |
如需详细了解安装以及如何开始使用 PyTorch 或 JAX，请参阅使用 PyTorch 在 Cloud TPU 虚拟机上运行计算以及使用 JAX 在 Cloud TPU 虚拟机上运行计算。
TensorFlow
TPU 软件版本的名称由以下部分组成：
tpu-vm-tf- TensorFlow 版本
-pod（如果您使用的是多主机 TPU 切片）-pjrt（如果您使用的是 PJRT API）
如需详细了解如何指定 TPU 软件版本，请参阅以下部分。
每个 TensorFlow 版本都有特定的 TPU 软件版本。下表显示了支持的 TensorFlow 版本和关联的 libtpu 版本：
| TensorFlow 版本 | libtpu.so 版本 |
|---|---|
| 2.18.0 | 1.12.0 |
| 2.17.1 | 1.11.1 |
| 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 |
如需详细了解 TensorFlow 补丁版本，请参阅支持的 TensorFlow 补丁版本。
TPU v6e、v5p 和 v5e
TPU v6e、v5e 和 v5p 支持 TensorFlow 2.15.0 及更高版本。您可以使用以下格式指定 TPU 软件版本：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。
TPU v4
如果您使用的是 TPU v4 和 TensorFlow 2.10.1 或更高版本，请按照适用于 TPU v2 和 v3 的说明操作。如果您使用的是 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 嵌入 API | 不适用 | 否 |
| TPU v5p | 密集计算（无 TPU 嵌入 API） | 是 | 否 |
| TPU v5p | TPU 嵌入 API | 是 | 否 |