Cambio di versione software su Cloud TPU

Panoramica

La versione software del framework in esecuzione sulla tua TPU deve corrispondere a quella in esecuzione sulla VM locale. Ora questa versione del software può essere trasferita su una Cloud TPU in esecuzione, senza eliminare e creare nuovamente la TPU. Ciò consente inoltre di configurare Cloud TPU con versioni notturne specifiche dei framework software. Ti consigliamo comunque di selezionare una versione supportata di questi framework.

Utilizzo

Il modo consigliato per cambiare versione è utilizzare la libreria Python cloud-tpu-client.

Esempio di utilizzo per TensorFlow.

from cloud_tpu_client import Client
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--tpu-name',
                    type=str,
                    required=True,
                    help='Name of the TPU Instance')
parser.add_argument('--target-version',
                    type=str,
                    required=True,
                    help='target TPU Runtime version')

args = parser.parse_args()

c = Client(args.tpu_name)
c.configure_tpu_version(args.target_version, restart_type='ifNeeded')
c.wait_for_healthy()

La Cloud TPU viene configurata in modo che corrisponda alla versione di TensorFlow in esecuzione sulla VM locale. Sono incluse sia le release ufficiali sia le build notturne con data.

Il parametro restart_type dell'API configure_tpu_version definisce il comportamento di riavvio della TPU quando si cambia versione. Le opzioni sono "sempre" (valore predefinito) e "senecessario".

• L'opzione 'sempre' può essere utilizzata per correggere una TPU con, ad esempio, stato UNHEALTHY_TENSORFLOW o che restituisce errori di esaurimento memoria (OOM) a causa di perdite di risorse da un'esecuzione precedente. Se questa opzione è impostata, la TPU viene riavviata anche se non è installata una nuova versione del framework.

• "ifneeded" può essere utile perché non riavvia il runtime se è già nella versione corretta, quindi non aggiungerà tempi di avvio significativi a uno script di addestramento. Se questa opzione è impostata, la TPU viene riavviata solo se non è installata la versione del framework corretta.

Poiché la libreria comunica direttamente con Cloud TPU, questo codice deve essere eseguito in una VM nella stessa rete. Ti consigliamo di eseguire questa operazione all'interno del codice per il resto del modello.

Opzioni software aggiuntive

TensorFlow include una stringa tf.__version__ che è il modo più semplice per configurare la versione corretta. Altre opzioni software includono:

• PyTorch: pytorch-1.13, pytorch-nightly-dev20220930, pytorch-nightly
• Jax - tpu_driver, tpu_driver0.1-dev20200320, tpu_driver_nightly

Ad esempio, configurare una TPU da eseguire con l'ultima build notturna di PyTorch.

from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version('pytorch-nightly', restart_type='ifNeeded')
c.wait_for_healthy()