Soluciona problemas de JAX: TPU

En esta guía, se proporcionan punteros a información de solución de problemas de JAX para ayudarte a identificar y resolver los problemas que puedes encontrar cuando entrenas modelos de AJAX en Cloud TPU.

Si quieres obtener una guía más general para comenzar a usar Cloud TPU, consulta la guía de inicio rápido de JAX.

Problemas generales de JAX

Si tienes problemas durante el desarrollo de tu modelo de entrenamiento o el entrenamiento con JAX, consulta las Preguntas frecuentes sobre JAX.

Para descubrir los errores de programación más generales que puedes encontrar cuando escribes una aplicación de entrenamiento con JAX, consulta Errores de JavaScript.

Crear perfiles de rendimiento de JAX

Puedes comprender cómo se usan tus recursos de TPU con las herramientas descritas en Cómo generar perfiles de rendimiento de JAX.

Solución de problemas de memoria

Puedes usar el Generador de perfiles de memoria de dispositivos de AJAX para ver cómo se usa la memoria de la TPU. Se puede usar para lo siguiente:

Soluciona problemas de TPU

¿Cómo puedo verificar que la TPU se esté ejecutando?

Detalles

Todo se ejecutará en la TPU, siempre y cuando JAX no imprima "No se encontró ninguna GPU/TPU, y recurre 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 verificarla de manera programática con assert jax.devices()[0].platform == 'tpu'.