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 tua VM TPU. Ciò include il sistema operativo Ubuntu, Docker e altri software necessari per eseguire il codice sulle TPU.

Se utilizzi Google Cloud CLI, specifica la versione del software TPU utilizzando il parametro --version o --runtime-version. Se utilizzi la console Google Cloud , seleziona una versione software TPU dall'elenco Versione 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 e 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 principale, secondaria e di patch di TensorFlow.
  • -pod (facoltativo): indica che si sta utilizzando una slice 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 ulteriori informazioni sulle versioni delle patch di TensorFlow, consultare Versioni delle patch di TensorFlow supportate.

TPU v6e, v5p e v5e

Le TPU v6e, v5e e v5p supportano TensorFlow 2.15.0 e versioni successive. È possibile specificare la versione del software TPU utilizzando il seguente modulo: tpu-vm-tf-x.y.z-{pod}-pjrt Dovex è la principale versione di TensorFlow,y è la versione minore ez è la versione patch di TensorFlow. Aggiungi pod dopo la versione di TensorFlow se stai utilizzando una TPU multi-host. Ad esempio, se si utilizza TensorFlow 2.16.0 su una TPU multi-host, utilizzare la versione software tpu-vm-tf-2.16.0-pod-pjrt TPU. Per altre versioni di TensorFlow, sostituire 2.16.0 con la versione principale e la versione patch di TensorFlow in uso. Se si utilizza una TPU a singolo host, omettere pod.

TPU v4

Se stai utilizzando TPU v4 e TensorFlow 2.10.1 o versioni successive, segui le istruzioni per TPU v2 e v3. Se si utilizza TensorFlow 2.10.0 o versioni precedenti, utilizzare una versione del software TPU specifica per la 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 si utilizza TPU v2 o v3, utilizzare la versione del software TPU corrispondente alla versione di TensorFlow in uso. Ad esempio, se stai utilizzando TensorFlow 2.14.1, usa la versione software tpu-vm-tf-2.14.1 TPU. Per altre versioni di TensorFlow, sostituire 2.14.1 con la versione di TensorFlow in uso. Se si utilizza una TPU multi-host, aggiungere "pod" alla fine della versione del software della 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 si utilizza TensorFlow 2.16.1 con l'API PJRT, utilizzare la versione software TPU tpu-vm-tf-2.16.1-pjrt. Se stai utilizzando l'API stream executor con la stessa versione di TensorFlow, usa la versione software tpu-vm-tf-2.16.1-se TPU. Le versioni di TensorFlow precedenti alla 2.15.0 supportano solo l'executor 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: semplificare l'integrazione di hardware e framework ML.

Acceleratore Funzionalità Supporto PJRT Supporto per l'esecutore di flussi
TPU v2 - v4 Calcolo denso (nessuna API di incorporamento TPU) Sì Sì
TPU v2 - v4 API Dense Compute + API TPU Embedding No Sì
TPU v2 - v4 tf.summary/tf.print con posizionamento del dispositivo software 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