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 durante la creazione di Cloud TPU.

PyTorch e JAX

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

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

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

TensorFlow

Esistono immagini VM TPU specifiche per ciascuna versione di TensorFlow. Sulle Cloud TPU sono supportate le seguenti versioni di TensorFlow:

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

Per TensorFlow 2.15.0 e versioni successive sono disponibili varianti di immagini VM TPU basate sull'API del dispositivo (PJRT o esecutore stream) che utilizzi.

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 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 patch di TensorFlow in uso. Se utilizzi una singola TPU host, ometti pod.

Pubblicazione su v5e

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

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, usa un'immagine VM TPU specifica 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, usa l'immagine VM TPU che corrisponde alla versione di TensorFlow in uso. Ad esempio, se utilizzi TensorFlow 2.14.1, usa l'immagine 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 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, utilizza l'immagine TPU tpu-vm-tf-2.16.1-se. Le versioni di TensorFlow precedenti alla 2.15.0 supportano solo l'esecutore di flussi di dati.

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

Acceleratore Selezione delle Supporto di PJRT Supporto degli esecutori di streaming
TPU v2 - v4 Computing denso (senza API di incorporamento TPU) Sì Sì
TPU v2 - v4 API Dense Compute + API di incorporamento TPU No Sì
TPU v2 - v4 tf.summary/tf.print con posizionamento del dispositivo software No Sì
TPU v5e Computing denso (senza API di incorporamento TPU) Sì No
TPU v5e API di incorporamento TPU N/A No
TPU v5p Computing denso (senza API di incorporamento TPU) Sì No
TPU v5p API di incorporamento TPU 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 di 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