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.
Para obtener una guía más general a fin de comenzar a usar Cloud TPU, consulta la guía de inicio rápido de JAX.
Problemas generales de 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 generales que puedes encontrar cuando escribes una aplicación de entrenamiento con JAX, consulta Errores de JAX.
Cómo generar perfiles del rendimiento de JAX
Puedes comprender cómo se utilizan tus recursos de TPU con las herramientas que se describen en Cómo crear perfiles del rendimiento de JAX.
Soluciona problemas de 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:
- Descubre qué arrays y ejecutables se encuentran en la memoria de la TPU en un momento determinado.
- Realiza 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 Notas de rendimiento para usar TPU con JAX.
Soluciona problemas de TPU
En las siguientes secciones, se describe cómo resolver algunos problemas comunes que puedes encontrar cuando ejecutas un programa de JAX en una TPU.
¿Cómo puedo verificar que la TPU se esté ejecutando?
Todo se ejecutará en la TPU, siempre y cuando JAX no imprima el mensaje "No se encontró ninguna GPU o TPU. Se recurrirá a la CPU".
Puedes verificar que la TPU esté activa mirando jax.devices(), donde deberías ver varios dispositivos TPU, o bien 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 tiempo 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ás ejecutando la versión incorrecta de la VM de TPU.
Verifica que estés ejecutando la versión actual del tiempo de ejecución de JAX y vuelve a intentarlo.
Soluciona problemas de TPU y GKE
Para ayudar a solucionar 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 se están ejecutando los Pods de Kubernetes de la implementación de jobset-controller-manager. 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).