Administra eventos de mantenimiento con pods de Cloud TPU

Descripción general

Las VMs de TPU son instancias de VMs de Compute Engine con hardware de TPU adjunto. Las VMs 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 VMs 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.

Importante: Cuando un nodo TPU falla o se cierra debido a un evento de mantenimiento, es posible que se pierdan los datos almacenados en el disco de inicio de la VM. Por lo tanto, te recomendamos que guardes los puntos de control en algún tipo de almacenamiento externo, como un disco persistente o un bucket de Cloud Storage. Para obtener más información, consulta Opciones de almacenamiento de TPU

.

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.

Los mecanismos que se usan para 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.

Cómo detectar eventos de mantenimiento

Puedes detectar si se produce un evento de mantenimiento en tu TPU y cuándo lo hace con el siguiente comando gcloud describe:

$ gcloud compute tpus tpu-vm 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

Registros de eventos de mantenimiento

Puedes ver los registros históricos de los eventos de mantenimiento en tu TPU en los registros de auditoría de eventos del sistema.

En el menú de navegación de la consola de Google Cloud, haz clic en Compute Engine > Instancias de VM y busca, por ejemplo:

"tpu.nodes.terminate" OR "tpu.nodes.restart"

Dentro del período de búsqueda, se muestran las interrupciones y reparaciones de los trabajadores de TPU. Los registros mostrarán la fecha y hora del evento, el tipo de evento y, para los eventos de "terminación", el motivo de la terminación en protoPayload.metadata.terminateReason.

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 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). Para obtener más información sobre el uso de Google Kubernetes Engine con VMs de TPU, consulta Cómo implementar cargas de trabajo de TPU. Puedes implementar una secuencia de comandos bash para reintentar la secuencia de comandos de entrenamiento de forma continua hasta que se complete. Por ejemplo:

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

   Nota: Debes ejecutar la secuencia de comandos desde Cloud Shell 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 TPUs 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.