Mantieni tutto organizzato con le raccolte Salva e classifica i contenuti in base alle tue preferenze.

Cambio di versione software su Cloud TPU

Panoramica

La versione software del framework in esecuzione sulla TPU deve corrispondere alla versione in esecuzione sulla VM locale. Questa versione software può ora essere attivata su una Cloud TPU in esecuzione, senza eliminare e ricreare la TPU. Ciò consente anche di configurare Cloud TPU con versioni notturne specifiche di framework software. È comunque consigliabile selezionare una versione supportata di questi framework.

Utilizzo

Il metodo 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()

Questo configura Cloud TPU in modo che corrisponda alla versione di TensorFlow in esecuzione sulla tua VM locale. Sono incluse le release ufficiali e le build notturne con data.

Il parametro restart_type dell'API configure_tpu_version definisce il comportamento di riavvio di TPU quando cambi versione. Le opzioni sono 'always' (impostazione predefinita) e 'ifNeeded'.

  • 'always' può essere utilizzato per correggere una TPU con, ad esempio, lo stato UNHEALTHY_TENSORFLOW o che restituisce errori di memoria insufficiente (OOM) a causa di risorse trafugate da un'esecuzione precedente. Quando questa opzione è impostata, la TPU viene riavviata anche quando 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. Quando questa opzione è impostata, la TPU viene riavviata solo se non è installata la versione del framework corretta.

La libreria comunica direttamente con Cloud TPU, pertanto questo codice deve essere eseguito in una VM nella stessa rete. Ti consigliamo di eseguire questa operazione nel codice per il resto del modello.

Opzioni software aggiuntive

TensorFlow include una stringa tf.__version__, 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 per l'esecuzione 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()