Versioni software TPU

Questo documento fornisce indicazioni per selezionare la versione software TPU appropriata quando crei Cloud TPU.

Quando crei risorse TPU, specifichi la versione software, chiamata anche versione di 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 del software TPU utilizzando il parametro --version o --runtime-version. Se utilizzi la Google Cloud console, 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
versione 4 e precedenti tpu-ubuntu2204-base

Per ulteriori informazioni 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 patch di TensorFlow.
  • -pod (facoltativo): indica che stai utilizzando un slice TPU multi-host.
  • -{device_api} (facoltativo): specifica l'API del dispositivo, ad esempio -pjrt (se utilizzi l'API PJRT).

Per ulteriori informazioni su come specificare una versione software della 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, 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 della 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 patch di TensorFlow in uso. Se utilizzi una TPU con un solo host, omettipod.

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 la versione 4:

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 corrispondente alla versione di TensorFlow in uso. 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 in uso. Se utilizzi una TPU multi-host, appendi pod alla fine della versione software della TPU, ad esempiotpu-vm-tf-2.14.1-pod.

A partire da TensorFlow 2.15.0, devi anche specificare un'API di dispositivo come parte del nome della versione del 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 StreamExecutor con la stessa versione di TensorFlow, utilizza la versione del software TPU tpu-vm-tf-2.16.1-se. Le versioni di TensorFlow precedenti alla 2.15.0 supportano solo l'executor di stream.

Supporto PJRT di TensorFlow

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 ulteriori informazioni su PJRT, consulta PJRT: semplificare l'integrazione di hardware e framework di ML.

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

Passaggi successivi