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

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

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