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, Docker e altro software necessario per eseguire il codice sulle TPU. Questo documento fornisce indicazioni per selezionare l'immagine VM TPU appropriata quando crei 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 (predefinita)
  • v2-alpha-tpuv6e (TPU v6e)
  • 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 ogni versione di TensorFlow. Sulle Cloud TPU sono supportate le seguenti versioni di TensorFlow:

  • 2.18.0
  • 2.17.1
  • 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

Il nome dell'immagine VM TPU è composto da tpu-vm-tf, la versione di TensorFlow, -pod se utilizzi una sezione TPU multihost e -pjrt se utilizzi l'API PJRT. Non tutte le versioni di TPU supportano PJRT. Consulta le sezioni seguenti per ulteriori informazioni su come specificare un'immagine VM TPU.

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

Per le versioni di TensorFlow 2.15.0 e successive sono disponibili varianti di immagini VM TPU basate sull'API del dispositivo (PJRT o executor di stream) in uso.

Addestramento su v6e, v5p e v5e

TPU v6e, 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 v6e e v5e

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

TPU v4

Se utilizzi TPU v4 e TensorFlow 2.15.0 o versioni successive, segui le istruzioni per l'addestramento su v6e, 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 in uso. Se utilizzi una TPU multi-host, aggiungi il pod 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 di 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 StreamExecutor 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'executor di stream.

Supporto PJRT di TensorFlow

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

Acceleratore Funzionalità Supporto PJRT Supporto dell'eseguitore dello stream
TPU v2 - v4 Calcolo denso (nessuna API di embedding 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 Calcolo denso (nessuna API di embedding TPU) Sì No
TPU v5e API di embedding 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 la tua immagine VM, utilizza le seguenti versioni software di TensorFlow TPU e le versioni corrispondenti di libtpu:

Versione TensorFlow Versione 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

Passaggi successivi