Administra eventos de mantenimiento con pods de Cloud TPU

Descripción general

Los nodos TPU y las VM TPU son instancias de VM de Compute Engine con hardware de TPU conectado. 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.

En este documento, se analizan varios enfoques a fin de controlar eventos de mantenimiento para 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 breves de eventos de mantenimiento y deben guardarse con frecuencia: una buena regla general es guardar puntos de control aproximadamente cada una hora. Si no controlas lo suficiente, podrías correr el riesgo de perder mucho progreso del entrenamiento debido a eventos de mantenimiento o interrupciones de otro tipo.

Por lo general, los puntos de control hacen referencia a todos los parámetros guardados que se usan en el entrenamiento (como las ponderaciones del modelo). El tiempo que se tarda en guardar un punto de control puede variar entre el orden de los segundos y el de 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 TPU, y reiniciar el trabajo de entrenamiento desde un punto de control guardado. Para obtener más información sobre cómo detectar fallas de recuperación automáticas y recuperarse de ellas, consulta Detecta y recupera fallas de TPU a continuación.

Los mecanismos usados para guardar y cargar puntos de control son diferentes en cada framework de AA. Por lo general, los modelos de Cloud TPU compatibles tienen puntos de control integrados. Para obtener más información sobre los puntos de control, consulta TensorFlow 2.x, PyTorch o JAX/flax.

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 el entrenamiento. 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 GKE o algún otro sistema de administración de recursos. De lo contrario, puedes implementar una secuencia de comandos bash para volver a ejecutar de forma continua la secuencia de comandos de entrenamiento hasta que se complete. Por ejemplo:

    Con el nodo TPU:

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

    Con la VM de TPU: 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 VM de TPU desde GCP o desde una terminal, no desde la VM de TPU).

  3. Detecta y recupera fallas de 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 borrar y volver a crear la TPU. Puedes encontrar un ejemplo de esta secuencia de comandos aquí. Consulta Administra TPU para obtener detalles sobre cómo borrar y volver a crear TPU de forma manual.