Gérer les événements de maintenance avec les pods Cloud TPU

Présentation

Les nœuds TPU et les VM TPU sont des instances de VM Compute Engine associées à du matériel TPU. Les VM Compute Engine sont soumises aux événements de maintenance des VM Compute Engine. Chaque TPU est connecté à une VM Compute Engine. Par conséquent, l'utilisation de plusieurs TPU (par exemple, dans un pod TPU) augmente la probabilité que l'une de vos VM rencontre un événement de maintenance.

Ce document présente différentes approches permettant de gérer les événements de maintenance pour les tâches d'entraînement de longue durée sur les Cloud TPU.

Utiliser des points de contrôle pour une récupération rapide après des événements de maintenance

Les points de contrôle sont essentiels pour récupérer rapidement les données après un événement de maintenance. Il est recommandé d'y recourir fréquemment: une bonne pratique consiste à enregistrer les points de contrôle environ toutes les heures. Ne définissez pas des points de contrôle assez souvent, car vous risquez de perdre beaucoup de progression de l'entraînement en raison d'événements de maintenance ou d'autres interruptions d'entraînement.

Les points de contrôle font généralement référence à tous les paramètres enregistrés utilisés pendant l'entraînement (tels que les pondérations du modèle). Le temps nécessaire à l'enregistrement d'un point de contrôle peut être compris entre quelques secondes et quelques minutes.

Bien que la plupart des événements de maintenance soient automatiquement récupérés et que les tâches d'entraînement se poursuivent sans intervention manuelle, il peut arriver que la tâche ne redémarre pas et continue automatiquement. Lorsque cela se produit, vous devez supprimer et recréer les ressources TPU, puis redémarrer la tâche d'entraînement à partir d'un point de contrôle enregistré. Pour en savoir plus sur la détection et la récupération des défaillances de récupération automatique, consultez la section Détecter et récupérer les défaillances du TPU ci-dessous.

Les mécanismes permettant de sauvegarder et de charger les points de contrôle sont différents pour chaque framework de ML. Les modèles Cloud TPU compatibles intègrent généralement des points de contrôle. Pour en savoir plus sur les points de contrôle, consultez la page TensorFlow 2.x, PyTorch ou JAX/flax.

Gérer les événements de maintenance

Il existe plusieurs façons de limiter les interruptions liées aux événements de maintenance.

  1. Enregistrer régulièrement les points de contrôle

    Dans l'idéal, lorsqu'un "événement d'interruption" se produit, l'entraînement reprend simplement à partir du dernier point de contrôle.

  2. Nouvelles tentatives d'exécution du script d'entraînement

    Le script d'entraînement peut s'arrêter à la suite d'un "événement d'interruption". Vous pouvez utiliser un script bash pour relancer le script d'entraînement en continu jusqu'à la fin de l'entraînement. Chaque nouvelle tentative doit poursuivre à partir du dernier point de contrôle. Par conséquent, les scripts de nouvelle tentative doivent toujours être utilisés conjointement avec les points de contrôle.

    Les pipelines d'entraînement en production doivent utiliser GKE ou d'autres systèmes de gestion des ressources. Sinon, vous pouvez mettre en œuvre un script bash pour relancer le script d'entraînement en continu jusqu'à la fin de l'opération. Exemple :

    Avec un nœud TPU:

    (Depuis votre VM) while ! python3 [training command]; do sleep 1; done

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

    Notez que vous devez exécuter la commande de VM TPU à partir de GCP ou d'un terminal, et non à partir de la VM TPU.

  3. Détecter et corriger les défaillances du TPU

    Lorsqu'un TPU ne récupère pas d'événement de maintenance, vous pouvez utiliser un script de récupération pour détecter son état, puis le supprimer et le recréer. Pour consulter un exemple de ce script, cliquez ici. Pour en savoir plus sur la suppression et la recréation manuelles de TPU, consultez la page Gérer les TPU.