Risoluzione dei problemi relativi a JAX - TPU

Questa guida fornisce indicazioni per le informazioni sulla risoluzione dei problemi di JAX per aiutarti a identificare e risolvere i problemi che potresti riscontrare durante l'addestramento dei modelli JAX su Cloud TPU.

Per una guida più generale su come iniziare a utilizzare Cloud TPU, consulta Guida rapida: crea un'istanza TPU.

Problemi generali di JAX

Se riscontri problemi durante lo sviluppo del modello di addestramento o l'addestramento con JAX, consulta le domande frequenti su JAX.

Per errori di programmazione più generali che potresti riscontrare durante la scrittura di un'applicazione di addestramento con JAX, consulta Errori JAX.

Profilare le prestazioni di JAX

Puoi capire come vengono utilizzate le risorse TPU utilizzando gli strumenti descritti in Profilazione delle prestazioni di JAX.

Risolvere i problemi di memoria

Puoi monitorare l'utilizzo della memoria con il profilatore della memoria del dispositivo JAX, ma non puoi gestirne direttamente l'utilizzo.

Il profiler della memoria del dispositivo JAX può essere utilizzato per:

Non puoi specificare come viene allocata la memoria TPU per operazioni specifiche. Per ulteriori informazioni sui problemi di rendimento delle TPU specifici di JAX, consulta Note sul rendimento per l'utilizzo delle TPU con JAX.

Risolvere i problemi relativi alla TPU

Le sezioni seguenti descrivono come risolvere alcuni problemi comuni che potresti riscontrare quando esegui un programma JAX su una TPU.

Come faccio a verificare che la TPU sia in esecuzione?

Tutto verrà eseguito sulla TPU, a condizione che JAX non stampi "No GPU/TPU found, falling back to CPU." (Nessuna GPU/TPU trovata, ripristino della CPU).

Puoi verificare che la TPU sia attiva esaminando jax.devices(), dove dovresti vedere diversi dispositivi TPU visualizzati, oppure verificare in modo programmatico con: assert jax.devices()[0].platform == 'tpu'.

RuntimeError: Unable to initialize backend 'tpu': UNAVAILABLE: No TPU Platform available.

Visualizzi il messaggio di errore RuntimeError: Unable to initialize backend 'tpu': UNAVAILABLE: No TPU Platform available o trovi il seguente avviso /tmp/tpu_logs/tpu_driver.WARNING sulla VM TPU:

W1118 17:40:20.985243   23901 tpu_version_flag.cc:57] No hardware is found. Using default TPU version:VERSION

Questo errore può indicare che stai eseguendo la versione errata della VM TPU. Verifica di utilizzare la versione attuale di JAX Runtime e riprova.

Risolvere i problemi relativi a TPU e GKE

Per facilitare la risoluzione dei problemi, attiva il logging dettagliato nel manifest del carico di lavoro GKE e poi fornisci i log all'assistenza GKE.

TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0

Le sezioni seguenti descrivono i messaggi di errore relativi alle configurazioni e alle risoluzioni di TPU e GKE.

nessun endpoint disponibile per il servizio "jobset-webhook-service"

Questo errore indica che il jobset non è stato installato correttamente. Controlla se i pod Kubernetes di deployment di jobset-controller-manager sono in esecuzione. Per saperne di più, consulta la documentazione sulla risoluzione dei problemi relativi a JobSet.

Inizializzazione della TPU non riuscita: impossibile connettersi

Assicurati che la versione del nodo GKE sia 1.30.4-gke.1348000 o successiva (GKE 1.31 non è supportato).