Solucionar problemas de JAX - TPU

En esta guía se proporcionan enlaces a información para solucionar problemas de JAX que te ayudarán a identificar y resolver los problemas que puedan surgir al entrenar modelos de JAX en Cloud TPU.

Para obtener una guía más general sobre cómo empezar a usar Cloud TPU, consulta la guía de inicio rápido de JAX.

Problemas generales de JAX

Si tienes problemas al desarrollar tu modelo de entrenamiento o al entrenar con JAX, consulta las preguntas frecuentes sobre JAX.

Para obtener más información sobre los errores de programación generales que pueden producirse al escribir una aplicación de entrenamiento con JAX, consulta Errores de JAX.

Crear perfiles de rendimiento de JAX

Puedes saber cómo se utilizan tus recursos de TPU con las herramientas descritas en Perfil de rendimiento de JAX.

Solucionar problemas de memoria

Puedes monitorizar cómo se usa la memoria con el profiler de memoria de dispositivos JAX, pero no puedes gestionar directamente cómo se usa.

El perfilador de memoria de dispositivos de JAX se puede usar para lo siguiente:

No puedes especificar cómo se asigna la memoria de la TPU a operaciones concretas. Para obtener más información sobre los problemas de rendimiento de las TPU específicos de JAX, consulta las notas sobre el rendimiento al usar las TPU con JAX.

Solucionar problemas de TPU

En las siguientes secciones se describe cómo solucionar algunos problemas habituales que pueden surgir al ejecutar un programa JAX en una TPU.

¿Cómo puedo verificar que la TPU está en funcionamiento?

Todo se ejecutará en la TPU siempre que JAX no imprima "No se ha encontrado ninguna GPU o TPU. Se va a usar la CPU".

Para comprobar que la TPU está activa, puedes consultar jax.devices(), donde deberían aparecer varios dispositivos TPU, o bien verificarlo mediante programación con assert jax.devices()[0].platform == 'tpu'.

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

Este mensaje de error de tiempo de ejecución o la siguiente información en /tmp/tpu_logs/tpu_driver.WARNING de la VM de TPU: W1118 17:40:20.985243 23901 tpu_version_flag.cc:57] No hardware is found. Using default TPU version:xxxxxx pueden indicar que estás ejecutando la versión incorrecta de la VM de TPU.

Comprueba que estés usando la versión actual del tiempo de ejecución de JAX y vuelve a intentarlo.

Solucionar problemas de TPU y GKE

Para facilitar la solución de problemas, habilita el registro detallado en el manifiesto de tu carga de trabajo de GKE y, a continuación, proporciona los registros al equipo de Asistencia de GKE.

TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0

En las siguientes secciones se describen los mensajes de error relacionados con las configuraciones de TPU y GKE, así como las soluciones.

No hay endpoints disponibles para el servicio "jobset-webhook-service"

Este error significa que el conjunto de trabajos no se ha instalado correctamente. Comprueba si los pods de Kubernetes de la implementación de jobset-controller-manager se están ejecutando. Para obtener más información, consulta la documentación sobre solución de problemas de JobSet.

No se ha podido inicializar la TPU: no se ha podido conectar

Asegúrate de que la versión de tu nodo de GKE sea la 1.30.4-gke.1348000 o una posterior (GKE 1.31 no es compatible).