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, Docker e altro software necessario per eseguire il codice sulle TPU. Questo documento fornisce indicazioni sulla selezione dell'immagine VM TPU appropriata quando devi creare le Cloud TPU.

PyTorch e JAX

Utilizza 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)
  • v2-alpha-tpuv5-lite (TPU v5e)

Per le istruzioni di installazione, consulta la documentazione di avvio rapido 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.17.0
  • 2.16.2
  • 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 delle patch di TensorFlow, consulta Versioni delle 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 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 patch di TensorFlow in uso. Se utilizzi una singola TPU host, ometti pod.

Pubblicazione su v5e

Esistono immagini Docker di servizio che contengono tutti i requisiti software necessari per il servizio 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 per la versione 4:

Versione TensorFlow Versione dell'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 corrispondente alla versione di TensorFlow in uso. Ad esempio, se utilizzi TensorFlow 2.14.1, utilizza l'immagine TPU tpu-vm-tf-2.14.1. Per 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 come parte del nome dell'immagine. Ad esempio, se utilizzi TensorFlow 2.16.1 con l'API PJRT, utilizza l'immagine TPU tpu-vm-tf-2.16.1-pjrt. Se utilizzi l'API stream executor con la stessa versione di TensorFlow, utilizza tpu-vm-tf-2.16.1-se immagine TPU. Versioni di TensorFlow precedenti alla 2.15.0 supportano solo l'esecutore per stream.

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 di framework e hardware per l'IA.

Acceleratore Funzionalità Assistenza PJRT Supporto degli esecutori di streaming
TPU v2 - v4 Computing denso (senza API di incorporamento 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 Computing denso (senza API di incorporamento TPU) Sì No
TPU v5e API di incorporamento TPU N/D No
TPU v5p Calcolo denso (nessuna API di embedding TPU) Sì No
TPU v5p API di embedding 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.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

Passaggi successivi