Risoluzione dei problemi relativi a JAX - TPU

Questa guida fornisce indicazioni sulle informazioni per la 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 la guida rapida di JAX.

Nota :se non riesci a risolvere il problema con questa guida, consulta la sezione Richiedere assistenza per ulteriore aiuto.

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

Nota: la memoria TPU (memoria ad alta larghezza di banda) è condivisa tra i Tensor Core su ogni chip TPU. Ciò consente un coordinamento efficiente tra i Tensor Core on-chip.

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:

• Scopri quali array ed eseguibili si trovano nella memoria TPU in un determinato momento oppure
• Individua le perdite di memoria.

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, oppure verificarlo a livello di programmazione con: assert jax.devices()[0].platform == 'tpu'.

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

Questo messaggio di errore di runtime o il seguente messaggio in /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:xxxxxx 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, abilita la registrazione dettagliata nel manifest del carico di lavoro GKE, quindi 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 ulteriori informazioni, 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).