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. 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 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 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 di 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 patch di TensorFlow, vedi Versioni patch di TensorFlow supportate.

TPU v6e, v5p e v5e

TPU v6e, v5e e v5p supportano TensorFlow 2.15.0 e versioni successive. Specifichi la versione del software TPU utilizzando il modulo: 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 che stai utilizzando. Se utilizzi una TPU con un solo 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 del 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, aggiungi pod alla fine della versione 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 del software. Ad esempio, se utilizzi TensorFlow 2.16.1 con l'API PJRT, utilizza la versione del software TPU tpu-vm-tf-2.16.1-pjrt. Se utilizzi l'API Stream Executor 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 lo stream executor.

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 dell'esecutore di stream
TPU v2 - v4 Calcolo denso (nessuna API TPU Embedding) Sì Sì
TPU v2 - v4 API Dense Compute + API TPU Embedding No Sì
TPU v2 - v4 tf.summary/tf.print con posizionamento soft dei dispositivi No Sì
TPU v5e Calcolo denso (nessuna API TPU Embedding) Sì No
TPU v5e API TPU embedding N/D No
TPU v5p Calcolo denso (nessuna API TPU Embedding) Sì No
TPU v5p API TPU embedding Sì No

Passaggi successivi