Immagini VM TPU

Quando crei risorse TPU, passi il parametro --version o --runtime-version che specifica un'immagine VM TPU. Le immagini VM TPU contengono il sistema operativo (Ubuntu) e, facoltativamente, altro software necessario per eseguire il codice sulle TPU. Questo documento fornisce indicazioni sulla selezione dell'immagine VM TPU appropriata quando crei Cloud TPU.

PyTorch e JAX

Utilizza le seguenti immagini di base delle VM TPU comuni per PyTorch e JAX, quindi installa il framework che vuoi utilizzare.

  • tpu-ubuntu2204-base (predefinito)
  • v2-alpha-tpuv5 (TPU v5p)

Per le istruzioni di installazione, consulta i documenti della guida rapida per PyTorch/XLA e JAX.

TensorFlow

Esistono immagini VM TPU specifiche per ogni versione di TensorFlow. Le seguenti versioni di TensorFlow sono supportate sulle Cloud TPU:

  • 2.16.1
  • 2.15.1
  • 2.15.0
  • 2.14.1
  • 2.14.0
  • 2.13.1
  • 2.13.0
  • 2.12.1
  • 2.12.0
  • 2.11.1
  • 2.11.0
  • 2.10.1
  • 2.10.0
  • 2.9.3
  • 2.9.1
  • 2.8.4
  • 2.8.3
  • 2.8.0
  • 2.7.4
  • 2.7.3

Per ulteriori informazioni sulle versioni della patch TensorFlow, consulta Versioni della patch di TensorFlow supportate.

Per TensorFlow versione 2.15.0 e successive, sono disponibili varianti di immagini VM TPU basate sull'API del dispositivo (PJRT o esecutore flusso) che stai utilizzando.

Addestramento su v5p e v5e

TPU v5e e v5p supportano TensorFlow 2.15.0 e versioni successive. Puoi specificare l'immagine VM 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 l'immagine VM TPU tpu-vm-tf-2.16.0-pod-pjrt. Per altre versioni di TensorFlow, sostituisci 2.16.0 con le versioni principali e con patch di TensorFlow che stai utilizzando. Se utilizzi una singola TPU host, ometti pod.

Pubblicazione su v5e

Esistono immagini Docker che contengono tutti i requisiti software necessari per la pubblicazione con TensorFlow, PyTorch e JAX. Per ulteriori informazioni, consulta introduzione all'inferenza v5e di Cloud TPU.

TPU v4

Se utilizzi TPU v4 e TensorFlow 2.15.0 o versioni successive, segui le istruzioni per l'addestramento su v5p e v5e. Se utilizzi TensorFlow 2.10.0 o versioni precedenti, utilizza un'immagine VM TPU specifica per la versione v4:

Versione TensorFlow Versione immagine VM 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 l'immagine VM TPU che corrisponde alla versione di TensorFlow che stai utilizzando. Ad esempio, se usi TensorFlow 2.14.1, usa l'immagine TPU tpu-vm-tf-2.14.1. Per le altre versioni di TensorFlow, sostituisci 2.14.1 con la versione di TensorFlow che stai utilizzando. Se utilizzi un pod di aggiunta TPU multi-host alla fine dell'immagine TPU, ad esempio tpu-vm-tf-2.14.1-pod.

A partire da TensorFlow 2.15.0, devi anche specificare un'API del dispositivo. Ad esempio, se utilizzi TensorFlow 2.16.1 con l'API PJRT, usa l'immagine TPU tpu-vm-tf-2.16.1-pjrt. Se utilizzi l'API Stream executor con la stessa versione di TensorFlow, usa l'immagine TPU tpu-vm-tf-2.16.1-se. Le versioni di TensorFlow precedenti alla 2.15.0 supportano solo l'esecuzione dei flussi.

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 Selezione delle Supporto di PJRT Supporto per esecutori di flussi
TPU v2 - v4 Dense computing (nessuna API di incorporamento TPU) Sì Sì
TPU v2 - v4 API Dense Compute + API Embedding TPU No Sì
TPU v2 - v4 tf.summary/tf.print con posizionamento soft device No Sì
TPU v5e Dense computing (nessuna API di incorporamento TPU) Sì No
TPU v5e API TPU Embedding N/A No
TPU v5p Dense computing (nessuna API di incorporamento TPU) Sì No
TPU v5p API TPU Embedding Sì No

Versioni Libtpu

Le immagini TensorFlow delle VM TPU contengono una versione specifica di TensorFlow e la corrispondente libreria libtpu. Se stai creando una tua immagine VM, utilizza le seguenti versioni software TensorFlow TPU e le versioni libtpu corrispondenti:

Versione TensorFlow Versione libtpu.so
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

Passaggi successivi