Soluciona problemas de JAX en TPU
En esta guía, se proporcionan punteros a información para solucionar problemas de JAX que te ayudarán a identificar y resolver problemas que podrías encontrar mientras entrenas modelos de JAX en Cloud TPU.
Si deseas obtener una guía más general para empezar a usar Cloud TPU, consulta la guía de inicio rápido de JAX.
Problemas generales con JAX
Si tienes problemas mientras desarrollas tu modelo de entrenamiento o entrenas con JAX, consulta las Preguntas frecuentes sobre JAX.
Para obtener más información sobre los errores de programación más generales que puedes encontrar cuando escribes una aplicación de entrenamiento con JAX, consulta Errores de JAX.
Perfil del rendimiento de JAX
Puedes comprender cómo se usan tus recursos TPU con las herramientas que se describen en Generar perfiles de rendimiento de JAX.
Soluciona problemas de la memoria
Puedes supervisar cómo se usa la memoria con el generador de perfiles de memoria del dispositivo JAX, pero no puedes administrar directamente cómo se usa.
El generador de perfiles de memoria del dispositivo JAX se puede usar para lo siguiente:
- Descubrir qué arrays y ejecutables se encuentran en la memoria de la TPU en un momento determinado
- Realizar un seguimiento de las fugas de memoria
No puedes especificar cómo se asigna la memoria de la TPU para operaciones específicas. Para obtener más información sobre los problemas de rendimiento de las TPU específicos de JAX, consulta las notas de rendimiento para usar TPU con JAX.
Soluciona problemas de la TPU
En las siguientes secciones, se describe cómo resolver algunos problemas habituales que puedes encontrar cuando ejecutas un programa de JAX en una TPU.
¿Cómo puedo verificar que la TPU esté en ejecución?
Todo se ejecutará en la TPU, siempre y cuando JAX no imprima el mensaje “No se encontró ninguna GPU o TPU, se recurrió a la CPU”.
Para verificar que la TPU esté activa, puedes mirar jax.devices(), en la que
deberías ver varios dispositivos de TPU o verificarlo
de forma programática con assert jax.devices()[0].platform == 'tpu'.
RuntimeError: No se pudo inicializar el backend “tpu”: UNAVAILABLE: No hay disponible ninguna plataforma de TPU.
Este mensaje de error de entorno de ejecución o encontrar lo siguiente en /tmp/tpu_logs/tpu_driver.WARNING en la VM de TPU:
W1118 17:40:20.985243 23901 tpu_version_flag.cc:57] No hardware is found. Using default TPU version:xxxxxx
puede indicar que está en ejecución la versión incorrecta de la VM de TPU.
Verifica que estés en ejecución la versión actual del entorno de ejecución de JAX y vuelve a intentarlo.
Soluciona problemas de la TPU y GKE
Para ayudar con la solución de problemas, habilita el registro detallado en el manifiesto de tu carga de trabajo de GKE y, luego, 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 la configuración y las resoluciones de TPU y GKE.
No hay extremos disponibles para el servicio “jobset-webhook-service”
Este error significa que el conjunto de trabajos no se instaló correctamente. Comprueba si los Pods de Kubernetes de la implementación de jobset-controller-manager están en ejecución. Para obtener más información, consulta la documentación sobre la solución de problemas de JobSet.
No se pudo inicializar la TPU: No se pudo establecer la conexión
Asegúrate de que la versión de tu nodo de GKE sea 1.30.4-gke.1348000 o posterior (GKE 1.31 no es compatible).