TPU 軟體版本

本文提供指引，協助您在建立 Cloud TPU 時，選取合適的 TPU 軟體版本。

建立 TPU 資源時，您會指定軟體版本 (即執行階段版本)，也就是預先安裝在 TPU VM 的軟體環境。包括 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

如要進一步瞭解如何安裝及開始使用 JAX 或 PyTorch，請參閱「使用 JAX 在 Cloud TPU VM 上執行運算」和「Cloud TPU 上的 PyTorch」。

TensorFlow

TensorFlow 的 TPU 軟體版本遵循特定命名慣例：

tpu-vm-tf-x.y.z[-{pod}][-{device_api}]

  • x.y.z：代表 TensorFlow 的主要、次要和修訂版本。
  • -pod (選用)：表示您使用的是多主機 TPU 配量。
  • -{device_api} (選用)：指定裝置 API，例如 -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，請在 TPU 軟體版本結尾附加 pod，例如 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 開始，您可以在 TPU 上使用 TensorFlow 的 PJRT 介面。PJRT 具備自動重組裝置記憶體的功能，可簡化硬體與框架的整合。如要進一步瞭解 PJRT，請參閱「PJRT：簡化機器學習硬體和架構整合」。

加速器 功能 PJRT 支援 Stream 執行器支援
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 是 否

後續步驟