Versioni software TPU

Questo documento fornisce indicazioni sulla selezione della versione software TPU appropriata quando crei Cloud TPU.

Quando crei risorse TPU, specifichi la versione software, chiamata anche versione runtime, che si riferisce all'ambiente software preinstallato sulla VM TPU. Sono inclusi il sistema operativo Ubuntu, Docker e altro software necessario per eseguire il codice sulle TPU.

Se utilizzi Google Cloud CLI, specifica la versione software TPU utilizzando il parametro --version o --runtime-version. Se utilizzi la Google Cloud console, seleziona una versione software TPU dall'elenco Versione software software TPU.

PyTorch e JAX

Utilizza le seguenti versioni software TPU comuni per PyTorch e JAX, quindi installa il framework che vuoi utilizzare.

Versione TPU Versione software TPU
Trillium (v6e) v2-alpha-tpuv6e
v5p v2-alpha-tpuv5
v5e v2-alpha-tpuv5-lite
v4 e versioni precedenti tpu-ubuntu2204-base

Per saperne di più sull'installazione e su come iniziare a utilizzare PyTorch o JAX, consulta Eseguire un calcolo su una VM Cloud TPU utilizzando PyTorch ed Eseguire un calcolo su una VM Cloud TPU utilizzando JAX.

TensorFlow

Le versioni software TPU per TensorFlow seguono una convenzione di denominazione specifica:

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

  • x.y.z: rappresenta le versioni principali, secondarie e delle patch di TensorFlow.
  • -pod (facoltativo): indica che stai utilizzando una sezione TPU multi-host.
  • -{device_api} (facoltativo): specifica l'API del dispositivo, ad esempio -pjrt (se utilizzi l'API PJRT).

Per saperne di più su come specificare una versione software TPU, consulta le sezioni seguenti.

Esistono versioni software TPU specifiche per ogni versione di TensorFlow. La tabella seguente mostra le versioni di TensorFlow supportate e le versioni di libtpu associate:

Versione TensorFlow Versione 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

Per saperne di più sulle versioni delle patch di TensorFlow, consulta Versioni delle patch di TensorFlow supportate.

TPU v6e, v5p e v5e

TPU v6e, v5e e v5p supportano TensorFlow 2.15.0 e versioni successive. Specifica la versione software TPU utilizzando il formato: tpu-vm-tf-x.y.z-{pod}-pjrt dove x è la versione principale di TensorFlow, y è la versione secondaria e z è la versione della patch di TensorFlow. Aggiungi pod dopo la versione di TensorFlow se utilizzi una TPU multi-host. Ad esempio, se utilizzi TensorFlow 2.16.0 su una TPU multi-host, utilizza la versione software TPU tpu-vm-tf-2.16.0-pod-pjrt. Per altre versioni di TensorFlow, sostituisci 2.16.0 con le versioni principali e delle patch di TensorFlow che stai utilizzando. Se utilizzi una TPU single-host, ometti pod.

TPU v4

Se utilizzi TPU v4 e TensorFlow 2.10.1 o versioni successive, segui le istruzioni per TPU v2 e v3. Se utilizzi TensorFlow 2.10.0 o versioni precedenti, utilizza una versione software TPU specifica per v4:

Versione TensorFlow Versione software 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 e v3

Se utilizzi TPU v2 o v3, utilizza la versione software TPU che corrisponde alla versione di TensorFlow che stai utilizzando. Ad esempio, se utilizzi TensorFlow 2.14.1, utilizza la versione software TPU tpu-vm-tf-2.14.1. Per altre versioni di TensorFlow, sostituisci 2.14.1 con la versione di TensorFlow che stai utilizzando. Se utilizzi una TPU multi-host, aggiungi pod alla fine della versione software TPU, ad esempio tpu-vm-tf-2.14.1-pod.

A partire da TensorFlow 2.15.0, devi specificare anche un'API del dispositivo come parte del nome della versione software. Ad esempio, se utilizzi TensorFlow 2.16.1 con l'API PJRT, utilizza la versione software TPU tpu-vm-tf-2.16.1-pjrt. Se utilizzi l'API dell'esecutore di stream con la stessa versione di TensorFlow, utilizza la versione software TPU tpu-vm-tf-2.16.1-se. Le versioni di TensorFlow precedenti alla 2.15.0 supportano solo l'esecutore di stream.

Supporto di TensorFlow PJRT

A partire da TensorFlow 2.15.0, puoi utilizzare l'interfaccia PJRT per TensorFlow su TPU. PJRT offre la deframmentazione automatica della memoria del dispositivo e semplifica l'integrazione dell'hardware con i framework. Per saperne di più su PJRT, consulta PJRT: Simplifying ML Hardware and Framework Integration.

Acceleratore Funzionalità Supporto PJRT Supporto dell'esecutore di stream
TPU v2 - v4 Calcolo denso (nessuna API di incorporamento TPU) Sì Sì
TPU v2 - v4 API di calcolo denso + API di incorporamento TPU No Sì
TPU v2 - v4 tf.summary/tf.print con posizionamento soft del dispositivo No Sì
TPU v5e Calcolo denso (nessuna API di incorporamento TPU) Sì No
TPU v5e API di incorporamento TPU N/D No
TPU v5p Calcolo denso (nessuna API di incorporamento TPU) Sì No
TPU v5p API di incorporamento TPU Sì No

Passaggi successivi