Administra eventos de mantenimiento con pods de Cloud TPU

Descripción general

Los nodos TPU y las VM de TPU son instancias de VM de Compute Engine con hardware de TPU adjunto. Las VM de Compute Engine están sujetas a los eventos de mantenimiento de VM de Compute Engine. Cada TPU está conectada a una VM de Compute Engine, por lo que usar más TPU (por ejemplo, en un pod de TPU) aumenta la probabilidad de que una de tus VM encuentre un evento de mantenimiento.

Nota: Las diferencias entre las VM de TPU y los nodos TPU se describen en Arquitectura del sistema.

En este documento se analizan varios enfoques para controlar eventos de mantenimiento de trabajos de entrenamiento de larga duración en Cloud TPU.

Usa puntos de control para una recuperación rápida de los eventos de mantenimiento

Los puntos de control son clave para las recuperaciones cortas de los eventos de mantenimiento y deben guardarse con frecuencia: una buena regla general es guardar puntos de control aproximadamente a cada hora. Si no guardas puntos de control a menudo corres el riesgo de perder mucho progreso del entrenamiento debido a eventos de mantenimiento o a otras interrupciones del entrenamiento.

En general, los puntos de control se refieren a todos los parámetros guardados que se usan en el entrenamiento (como los pesos de los modelos). El tiempo que lleva guardar un punto de control puede variar desde segundos hasta minutos.

Aunque la mayoría de los eventos de mantenimiento se recuperan de forma automática y los trabajos de entrenamiento continúan sin intervención manual, puede haber casos extremos en los que el trabajo no se reinicie y continúe de forma automática. Cuando esto sucede, debes borrar y volver a crear los recursos de TPU y reiniciar el trabajo de entrenamiento desde un punto de control guardado. Para obtener información sobre cómo detectar y recuperarse ante fallas automáticas, consulta Detecta y recupera ante fallas de TPU a continuación.

Los mecanismos que se usan a fin de guardar y cargar puntos de control son diferentes para cada framework de AA. En general, los modelos compatibles de Cloud TPU tienen puntos de control integrados. Para obtener más información sobre los puntos de control, consulta: TensorFlow 2.x, PyTorch o JAX/flax.

Detecta eventos de mantenimiento

Puedes detectar si ocurre un evento de mantenimiento en tu TPU y cuándo puede hacerlo con el siguiente comando gcloud describe:

VM de TPU

$ gcloud compute tpus tpu-vm describe tpu-name --zone=zone  | grep 'health'

Nodos TPU

$ gcloud compute tpus describe tpu-name --zone=zone | grep 'health'

El resultado de este comando muestra el estado actual de la TPU y una descripción del evento de mantenimiento más reciente. El resultado debería ser similar al siguiente:

health: HEALTHY
healthDescription: The TPU had a maintenance event at 2022-01-26T03:44:36.265703305Z

Controla eventos de mantenimiento

Existen varias formas de mitigar las interrupciones del evento de mantenimiento.

1. Guarda puntos de control de forma periódica

   En la situación ideal, cuando ocurre un "evento de interrupción", el entrenamiento solo se reanuda desde el último punto de control.

2. Reintentos de la secuencia de comandos de entrenamiento

   La secuencia de comandos de entrenamiento puede detenerse como resultado de un "evento de interrupción". Puedes usar una secuencia de comandos bash para reintentar la secuencia de comandos de entrenamiento de forma continua hasta que se complete. Cada reintento debe continuar desde el último punto de control, por lo que las secuencias de comandos de reintento siempre deben usarse junto con los puntos de control.

   Las canalizaciones de entrenamiento listas para la producción deben usar un sistema de administración de recursos, como Google Kubernetes Engine (GKE). En la actualidad, GKE es compatible con los nodos TPU, pero no con las VM de TPU. De lo contrario, puedes implementar una secuencia de comandos bash para reintentar la secuencia de comandos de entrenamiento de forma continua hasta que se complete. Por ejemplo:

   Con un nodo TPU:

   Desde tu VM while ! python3 [training command]; do sleep 1; done

   Con la VM de TPU, haz lo siguiente:

   while ! gcloud compute tpus tpu-vm ssh ${TPU_NAME} --command "python3 [training command]"; do sleep 1; done
   

   Ten en cuenta que debes ejecutar el comando de la VM de TPU desde una shell de Google Cloud o desde una terminal, no desde la VM de TPU.

3. Detecta y recupera fallas de la TPU

   Cuando una TPU no se recupera de un evento de mantenimiento, puedes usar una secuencia de comandos de recuperación para detectar el estado de la TPU y volver a crearla. Puedes encontrar un ejemplo de esta secuencia de comandos aquí. Consulta Administra TPU para obtener detalles sobre cómo borrar y volver a crear las TPU de forma manual.

   Cuando creas o vuelves a crear una VM de TPU, puedes especificar una secuencia de comandos de inicio con el parámetro --metadata startup-script. Una secuencia de comandos de inicio se ejecuta cada vez que se crea una VM de TPU. Consulta Ejecuta secuencias de comandos de instalación estándar para obtener más información.