Préserver la progression de l'entraînement à l'aide d'un point de contrôle automatique

Historiquement, lorsqu'une VM TPU nécessite une maintenance, la procédure est lancée immédiatement, sans laisser aux utilisateurs le temps d'effectuer des actions de conservation de la progression, telles que l'enregistrement d'un point de contrôle. C'est ce que montre la figure 1(a).

Fig. 1. Illustration de la fonctionnalité de point de contrôle automatique : (a) Sans point de contrôle automatique, la progression de l'entraînement à partir du dernier point de contrôle est perdue en cas d'événement de maintenance à venir. (b) Avec Autocheckpoint, la progression de l'entraînement depuis le dernier point de contrôle peut être préservée en cas d'événement de maintenance à venir.

Vous pouvez utiliser le point de contrôle automatique (figure 1b) pour préserver la progression de l'entraînement en configurant votre code pour enregistrer un point de contrôle non planifié lorsqu'un événement de maintenance se produit. Lorsqu'un événement de maintenance se produit, la progression depuis le dernier point de contrôle est automatiquement enregistrée. Cette fonctionnalité fonctionne à la fois sur les tranches uniques et sur Multislice.

La fonctionnalité Autocheckpoint fonctionne avec les frameworks pouvant capturer des signaux SIGTERM et enregistrer ensuite un point de contrôle. Les frameworks compatibles incluent:

• MaxText,
• Pax,
• JAX avec Orbax

Utiliser le point de contrôle automatique

La fonctionnalité de point de contrôle automatique est désactivée par défaut. Lorsque vous créez un TPU ou demandez une ressource en file d'attente, vous pouvez activer le point de contrôle automatique en ajoutant l'indicateur --autocheckpoint-enabled lors du provisionnement du TPU. Lorsque cette fonctionnalité est activée, Cloud TPU effectue les étapes suivantes lorsqu'il reçoit une notification d'événement de maintenance:

1. Capturer le signal SIGTERM envoyé au processus à l'aide de l'appareil TPU
2. Attendez la fin du processus ou cinq minutes, selon la première échéance atteinte.
3. Effectuer la maintenance des segments concernés

L'infrastructure utilisée par Autocheckpoint est indépendante du framework de ML. N'importe quel framework de ML peut prendre en charge la création automatique de points de contrôle s'il peut capturer le signal SIGTERM et lancer un processus de création de points de contrôle.

Dans le code de l'application, vous devez activer les fonctionnalités de point de contrôle automatique fournies par le framework de ML. Dans Pax, par exemple, cela signifie activer les indicateurs de ligne de commande lors du lancement de l'entraînement. Pour en savoir plus, consultez le guide de démarrage rapide Autocheckpoint avec Pax. En coulisses, les frameworks enregistrent un point de contrôle non planifié lorsqu'un signal SIGTERM est reçu, et la VM TPU concernée est soumise à une maintenance lorsque le TPU n'est plus utilisé.

Guide de démarrage rapide: Point de contrôle automatique avec MaxText

MaxText est un LLM Open Source hautes performances, évolutif de manière arbitraire et bien testé, écrit en Python/JAX pur et ciblant les Cloud TPU. MaxText contient toute la configuration nécessaire pour utiliser la fonctionnalité de point de contrôle automatique.

Le fichier README MaxText décrit deux façons d'exécuter MaxText à grande échelle:

• Utilisation de multihost_runner.py, recommandé pour les tests
• Utilisation de multihost_job.py, recommandée pour la production

Lorsque vous utilisez multihost_runner.py, activez le point de contrôle automatique en définissant l'indicateur autocheckpoint-enabled lors du provisionnement de la ressource mise en file d'attente.

Lorsque vous utilisez multihost_job.py, activez le point de contrôle automatique en spécifiant l'indicateur de ligne de commande ENABLE_AUTOCHECKPOINT=true lors du lancement de la tâche.

Démarrage rapide: point de contrôle automatique avec Pax sur une seule tranche

Cette section fournit un exemple de configuration et d'utilisation d'Autocheckpoint avec Pax sur une seule tranche. Avec la configuration appropriée:

• Un point de contrôle est enregistré lorsqu'un événement de maintenance se produit.
• Cloud TPU effectuera la maintenance des VM TPU concernées une fois le point de contrôle enregistré.
• Une fois la maintenance terminée, vous pouvez utiliser la VM TPU comme d'habitude.

1. Utilisez l'option autocheckpoint-enabled lorsque vous créez la VM TPU ou demandez une ressource mise en file d'attente.

   Exemple :

   export PROJECT=your-gcp-project-name
   export ZONE=zone-you-want-to-use
   export NODE_ID=your-node-id
   export ACCELERATOR_TYPE=your-accelerator-type
   gcloud config set project $PROJECT
   gcloud config set compute/zone $ZONE

   gcloud alpha compute tpus tpu-vm create $NODE_ID \
       --accelerator-type $ACCELERATOR_TYPE \
       --version tpu-ubuntu2204-base \
       --autocheckpoint-enabled

2. Connectez-vous au TPU à l'aide de SSH:

   gcloud compute tpus tpu-vm ssh $NODE_ID 
   

3. Installer Pax sur une seule tranche

   La fonctionnalité de point de contrôle automatique fonctionne avec les versions Pax 1.1.0 et ultérieures. Sur la VM TPU, installez jax[tpu] et la dernière version de paxml:

   pip install paxml && pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

4. Lancez l'entraînement avec la configuration appropriée.

   L'exemple suivant montre comment configurer le modèle LmCloudSpmd2B pour enregistrer les points de contrôle déclenchés par Autocheckpoint dans un bucket Cloud Storage. Remplacez your-storage-bucket par le nom d'un bucket existant ou créez un bucket.

   export JOB_LOG_DIR=gs://your-storage-bucket
   
   { python3 .local/lib/python3.10/site-packages/paxml/main.py \
       --jax_fully_async_checkpoint=1 \
       --exit_after_ondemand_checkpoint=1 \
       --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2B \
       --job_log_dir=$JOB_LOG_DIR; } 2>&1 | tee pax_logs.txt

   Notez les deux options transmises à la commande:

   • jax_fully_async_checkpoint : lorsque cette option est activée, orbax.checkpoint.AsyncCheckpointer est utilisé. La classe AsyncCheckpointer enregistre automatiquement un point de contrôle lorsque le script d'entraînement reçoit un signal SIGTERM.
   • exit_after_ondemand_checkpoint : lorsque cet indicateur est activé, le processus TPU se termine une fois le point de contrôle automatique enregistré, ce qui déclenche immédiatement la maintenance. Si vous n'utilisez pas cet indicateur, l'entraînement se poursuit une fois le point de contrôle enregistré, et Cloud TPU attend un délai avant expiration (cinq minutes) avant d'effectuer la maintenance requise.

Démarrage rapide: point de contrôle automatique avec Pax sur Multislice

Le point de contrôle automatique fonctionne non seulement pour les tranches uniques, mais aussi pour les tranches multiples. Cette section décrit les étapes à suivre pour utiliser Autocheckpoint avec Multislice.

1. Spécifiez le point de contrôle automatique lors de la création de ressources en file d'attente.

   Un environnement multislice ne peut être provisionné que via une requête de ressources mise en file d'attente. Comme pour le cas d'une seule tranche, utilisez l'indicateur autocheckpoint-enabled dans l'appel pour créer une ressource mise en file d'attente.

   export QR_ID=your-qr-id
   export NODE_COUNT=your-node-count
   export ACCELERATOR_TYPE=your-accelerator-type
   
   gcloud alpha compute tpus queued-resources create $QR_ID \
       --node-count $NODE_COUNT \
       --accelerator-type $ACCELERATOR_TYPE \
       --runtime-version tpu-ubuntu2204-base \
       --autocheckpoint-enabled

   Pour en savoir plus sur toutes les options disponibles, consultez le guide de l'utilisateur de Multislice. Lorsque la requête de ressource mise en file d'attente est créée et qu'elle est à l'état ACTIVE, suivez les étapes suivantes pour exécuter Pax avec Autocheckpoint.

2. Installez jax[tpu] et la dernière version de paxml sur toutes les VM TPU de votre environnement multislice.

   gcloud compute tpus queued-resources ssh $QR_ID \
       --node=all \
       --worker=all \
       --batch-size=your-batch-size \
       --command="pip install paxml && pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"

   Définissez l'indicateur --batch-size sur le nombre de connexions simultanées à établir avec les nœuds de calcul TPU. Pour en savoir plus sur le choix d'une taille de lot pour les charges de travail Multislice, consultez la section Optimiser l'entraînement.

3. Configurez le modèle LmCloudSpmd2B pour le point de contrôle automatique lors de l'entraînement dans un environnement multicouche. Avant d'exécuter le script d'entraînement, définissez DCN_MESH_SHAPE sur [2, 1, 1], comme illustré dans l'exemple suivant:

   @experiment_registry.register
   class LmCloudSpmd2B(LmCloudSpmd):
   
       """SPMD model with 2B params.
   
       Global batch size = 2 * 2 * 1 * 32 = 128
       """
       PERCORE_BATCH_SIZE = 8
   
       NUM_LAYERS = 18
       MODEL_DIMS = 3072
       HIDDEN_DIMS = MODEL_DIMS * 4
   
       CHECKPOINT_POLICY = layers.AutodiffCheckpointType.SAVE_NOTHING
       ICI_MESH_SHAPE = [1, 4, 1]
       DCN_MESH_SHAPE = [2, 1, 1]

4. Pour définir des points de contrôle plus fréquemment, définissez task_p.train.save_interval_steps et task_p.train.save_max_to_keep, comme illustré dans l'exemple suivant:

   @experiment_registry.register
   class LmCloudSpmd2BLimitSteps(LmCloudSpmd2B):
   """SPMD model with 2B params and limited steps.
   
   Global batch size = 2 * 2 * 1 * 32 = 128
   """
   
   def task(self) -> pax_fiddle.Config[tasks_lib.SingleTask]:
       task_p = super().task()
       task_p.train.save_interval_steps = 50
       task_p.train.save_max_to_keep = 5
       return task_p
   

5. Lancez l'entraînement en exécutant la commande suivante pour chaque hôte. Remplacez your-storage-bucket par le nom d'un bucket existant ou créez un bucket.

   export TF_CPP_MIN_LOG_LEVEL=0
   export JOB_LOG_DIR=gs://your-storage-bucket
   
   { python3 .local/lib/python3.10/site-packages/paxml/main.py \
       --jax_fully_async_checkpoint=1 \
       --num_hosts=2 \
       --host_idx=host-index \
       --server_addr=worker0-node0-ip-address \
       --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2B \
       --job_log_dir=$JOB_LOG_DIR; } 2>&1 | tee pax_logs.txt

   Lors du lancement de l'entraînement, en plus des indicateurs de ligne de commande abordés dans le cas d'une seule tranche, trois autres sont requis:

   • num_hosts: nombre total d'hôtes. Dans ce cas, il s'agit de 2.
   • host_idx: index de l'hôte qui lance l'entraînement. Il varie de 0 à N-1, où N est le nombre total d'hôtes.
   • server_addr: adresse IP du worker 0 du nœud 0, avec un port inutilisé (par exemple, 8476). Pour obtenir ces informations, utilisez hostname -i sur le nœud de calcul 0.

Point de contrôle automatique avec Orbax

La fonctionnalité de contrôle automatique n'est pas limitée à MaxText ou Pax. Tout framework capable de capturer le signal SIGTERM et d'initier un processus de point de contrôle fonctionne avec l'infrastructure fournie par Autocheckpoint. Orbax, un espace de noms qui fournit des bibliothèques d'utilitaires courantes pour les utilisateurs de JAX, fournit ces fonctionnalités.

Comme expliqué dans la documentation Orbax, ces fonctionnalités sont activées par défaut pour les utilisateurs de orbax.checkpoint.CheckpointManager. La méthode save appelée après chaque étape vérifie automatiquement si un événement de maintenance est imminent et, le cas échéant, enregistre un point de contrôle, même si le numéro d'étape n'est pas un multiple de save_interval_steps. La documentation GitHub explique également comment arrêter l'entraînement après avoir enregistré un point de contrôle automatique, avec une modification dans le code utilisateur.