Résoudre les problèmes liés à JAX – TPU

Ce guide fournit des conseils et des informations de dépannage sur JAX pour vous aider à identifier et résoudre les problèmes que vous pouvez rencontrer lors de l'entraînement de modèles JAX sur Cloud TPU.

Pour obtenir des conseils plus généraux pour faire vos premiers pas avec Cloud TPU, consultez le guide de démarrage rapide de JAX.

Problèmes d'ordre général avec JAX

Si vous rencontrez des problèmes lors du développement de votre modèle d'entraînement ou de votre entraînement avec JAX, consultez les questions fréquentes sur JAX.

Pour connaître les erreurs de programmation plus générales que vous pouvez rencontrer lors de l'écriture d'une application d'entraînement avec JAX, consultez la section Erreurs JAX.

Profiler les performances JAX

Vous pouvez comprendre l'utilisation de vos ressources TPU à l'aide des outils décrits dans la section Profiler les performances JAX.

Résoudre les problèmes liés à la mémoire

Vous pouvez surveiller l'utilisation de la mémoire avec le profileur de mémoire de l'appareil JAX, mais vous ne pouvez pas gérer directement son utilisation.

Le profileur de mémoire de l'appareil JAX peut être utilisé pour :

Vous ne pouvez pas spécifier la façon dont la mémoire TPU est allouée pour des opérations spécifiques. Pour en savoir plus sur les problèmes de performances de TPU spécifiques à JAX, consultez Remarques concernant les performances lors de l'utilisation de TPU avec JAX.

Résoudre les problèmes liés aux TPU

Les sections suivantes décrivent comment résoudre certains problèmes courants que vous pouvez rencontrer lorsque vous exécutez un programme JAX sur un TPU.

Comment vérifier que le TPU est en cours d'exécution ?

Toutes les opérations seront exécutées sur le TPU tant que JAX n'affiche pas "Aucun GPU/TPU trouvé, retour au processeur".

Vous pouvez vérifier que le TPU est actif en consultant jax.devices(), qui devrait afficher plusieurs appareils TPU, ou en effectuant une vérification automatisée avec : assert jax.devices()[0].platform == 'tpu'.

RuntimeError : Impossible d'initialiser le backend "tpu" : INDISPONIBLE : Aucune plate-forme TPU disponible.

Ce message d'erreur d'exécution, ou la présence des éléments suivants dans /tmp/tpu_logs/tpu_driver.WARNING sur la VM TPU : W1118 17:40:20.985243 23901 tpu_version_flag.cc:57] No hardware is found. Using default TPU version:xxxxxx, peuvent indiquer que vous exécutez la mauvaise version de VM TPU.

Vérifiez que vous exécutez la version actuelle de l'environnement d'exécution JAX, puis réessayez.

Résoudre les problèmes liés aux TPU et à GKE

Pour faciliter le dépannage, activez la journalisation détaillée dans votre fichier manifeste de charge de travail GKE, puis fournissez les journaux à l'assistance GKE.

TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0

Les sections suivantes décrivent les messages d'erreur liés aux configurations de TPU et de GKE, ainsi que les solutions correspondantes.

Aucun point de terminaison disponible pour le service "jobset-webhook-service"

Cette erreur signifie que le JobSet n'a pas été installé correctement. Vérifiez si les pods Kubernetes de déploiement jobset-controller-manager sont en cours d'exécution. Pour en savoir plus, consultez la documentation sur le dépannage de JobSet.

Échec d'initialisation du TPU : Échec de la connexion

Assurez-vous que la version de votre nœud GKE est 1.30.4-gke.1348000 ou une version ultérieure (GKE 1.31 n'est pas compatible).