Immagini VM TPU

Quando crei risorse TPU, passi il token --version o --runtime-version che specifica un'immagine VM TPU. Le immagini VM TPU contengono il sistema operativo (Ubuntu) e facoltativamente è necessario altro software per eseguire il codice sulle TPU. Questo documento fornisce indicazioni sulla selezione dell'immagine VM TPU appropriata quando crei 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. La 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 in base all'API del dispositivo (PJRT o esecutore di streaming) che stai utilizzando.

Addestramento su v5p e v5e

TPU v5e e v5p supportano TensorFlow 2.15.0 e versioni successive. Specifica la VM TPU immagine utilizzando il formato: tpu-vm-tf-x.y.z-{pod}-pjrt, dove x è la principale La versione 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 la principale e la patch di TensorFlow in uso. Se utilizzi una singola TPU host, ometti pod.

Pubblicazione su v5e

Sono presenti immagini Docker che contengono tutti i requisiti software necessari con TensorFlow, PyTorch e JAX. Per ulteriori informazioni, vedi 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, utilizza 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 che stai utilizzando. Ad esempio, se utilizzi TensorFlow 2.14.1, utilizza 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 utilizzano. Se utilizzi un pod di aggiunta TPU multi-host alla fine della TPU dell'immagine, 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 la TPU tpu-vm-tf-2.16.1-se dell'immagine. 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 on TPU. PJRT è dotato di 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 Funzionalità 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/D 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 libreria libtpu corrispondente. Se crei una tua immagine VM, utilizza le seguenti versioni del software TensorFlow TPU e i relativi Versioni libtpu:

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