Entraînement résilient avec Pathways

Pathways offre des avantages en termes de résilience de différentes manières :

  • Suspension/reprise : permet à votre charge de travail de tolérer les interruptions planifiées, telles que les notifications de préemption, sans nécessiter de code de gestion de préemption personnalisé.
  • Entraînement élastique : permet à votre charge de travail de tolérer les défaillances matérielles non planifiées sans planter le client, en utilisant une logique de récupération dans votre code d'entraînement.

Avant de commencer

Avant de commencer, effectuez les tâches suivantes :

Suspendre/Reprendre

Google Kubernetes Engine (GKE) envoie généralement une notification de préemption à un pod d'accélérateur avant que le pod ne soit préempté. La tolérance de préemption Pathways est activée par défaut sur tous les déploiements cloud, et les jobs d'accélérateur Pathways écoutent ces notifications.

Lorsqu'un avis de préemption arrive, Pathways détermine d'abord s'il peut enregistrer et restaurer la charge de travail de manière transparente. Si c'est le cas, il tente de suspendre de manière transparente votre charge de travail de ML en écrivant son état actuel dans un stockage persistant tel que Cloud Storage avant que GKE n'expulse vos jobs d'accélérateur. Lorsque GKE reprogramme vos jobs ultérieurement, Pathways reprend votre charge de travail de ML en lisant son état persistant.

Si la charge de travail ne peut pas être restaurée, Pathways arrête le job de l'accélérateur et transmet l'échec à votre job si l'entraînement élastique est configuré. Si l'entraînement élastique n'est pas configuré, GKE redémarre l'intégralité de la charge de travail en fonction de la règle de redémarrage JobSet.

Les charges de travail de ML standards définies à l'aide de JAX s'appuient sur des composants Pathways XLA sans état qui peuvent être restaurés à l'aide d'un instantané de mémoire à haut débit (HBM). Certaines charges de travail de ML, telles que celles définies à l'aide de l'API Python colocalisée JAX, s'appuient sur des composants Pathways avec état. Ces charges de travail ne peuvent pas être restaurées.

Entraînement élastique

L'entraînement élastique permet à votre job d'entraînement de continuer en cas de défaillance matérielle ou de préemption de tranche, sans redémarrer l'intégralité de la charge de travail. Dans un job d'entraînement multitranche, Pathways maintient le processus client à contrôleur unique et les tranches de TPU opérationnelles lorsqu'une tranche individuelle devient indisponible.

L'entraînement élastique combine les fonctionnalités du système Pathways avec la logique de récupération dans votre code d'entraînement :

  • Détection des défaillances : lorsqu'une défaillance matérielle ou une préemption se produit sur une tranche TPU, Pathways annule l'attribution de la tranche concernée et génère une erreur jax.errors.JaxRuntimeError la prochaine fois que votre charge de travail accède aux données de cette tranche. Étant donné que le processus client reste actif, votre code peut intercepter cette erreur, reconfigurer vos ressources et poursuivre l'entraînement.
  • Modes de récupération : la classe pathwaysutils.elastic.manager.Manager et son décorateur elastic_retry sont compatibles avec deux stratégies de récupération :
    • Mettre en pause et reprendre : lorsque minimum_slice_count est défini sur None ou est égal au nombre total de tranches du job, l'entraînement est mis en pause lorsqu'une tranche est perdue et attend que toutes les tranches soient à nouveau opérationnelles avant de reprendre. Étant donné que le nombre de tranches ne change pas d'une tentative à l'autre, votre charge de travail peut réutiliser son maillage d'appareils et ses exécutables compilés existants.
    • Redimensionnement de l'instance répliquée : lorsque minimum_slice_count est inférieur au nombre total de tranches, l'entraînement se poursuit immédiatement sur les tranches actives restantes en recréant le maillage de l'appareil et en recompilant pour la topologie plus petite. Un thread de surveillance en arrière-plan vérifie régulièrement les tranches inactives et enregistre les tranches nouvellement opérationnelles dans elastic_manager.available_inactive_slices. Votre boucle d'entraînement peut ensuite recompiler le maillage de l'appareil avec les tranches de remplacement, réinitialiser la topologie étendue et placer l'état actuel sur le nouveau maillage.
  • Récupération d'état : pour reprendre l'entraînement après un événement élastique, votre charge de travail restaure les paramètres du modèle, l'état de l'optimiseur et l'état de l'itérateur de données à l'aide d'instantanés en mémoire. La classe Snapshotter dans orbax.checkpoint copie de manière asynchrone l'état répliqué de la mémoire HBM du TPU vers la mémoire du processeur pinned_host sur chaque tranche. Pendant la récupération, Snapshotter reconstruit l'état à partir des instances répliquées de l'hôte survivant et le repartitionne directement sur le maillage TPU actif sans lire à partir du stockage distant.

Activer l'élasticité sur le serveur proxy Pathways

Pour activer les ensembles d'appareils élastiques sur Pathways dans GKE, ajoutez l'option --num_elastic_instances au champ args du conteneur pathways-proxy dans votre fichier manifeste JobSet et spécifiez le nombre de tranches TPU élastiques dans votre charge de travail :

- name: pathways-proxy
  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
  args:
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --server_port=29000
  - --gcs_scratch_location=gs://BUCKET_NAME
  - --num_elastic_instances=NUM_SLICES

Remplacez les éléments suivants :

  • BUCKET_NAME : bucket Cloud Storage utilisé pour les fichiers temporaires Pathways.
  • NUM_SLICES : nombre de tranches de TPU dans votre charge de travail.

Filtrer les appareils actifs et créer le maillage d'appareils

Dans une charge de travail Pathways, jax.devices() renvoie tous les appareils de l'allocation initiale pendant la durée de vie du processus client. Lorsque vous utilisez l'entraînement élastique, utilisez elastic_manager.active_devices afin que votre maillage d'appareils et votre taille de lot n'utilisent que des tranches saines et actives :

import jax
import numpy as np
from pathwaysutils.elastic import manager

elastic_manager = manager.Manager()


def create_mesh(
    elastic_manager: manager.Manager, devices_per_slice: int
) -> jax.sharding.Mesh:
  device_array = np.asarray(elastic_manager.active_devices).reshape(
      (elastic_manager.active_slice_count, devices_per_slice)
  )
  return jax.sharding.Mesh(device_array, axis_names=("replica", "model"))

Lorsque le nombre de tranches actives change en mode de redimensionnement des instances répliquées, utilisez un chargeur de données élastique tel que grain.experimental.ElasticIterator. Les itérateurs de données par hôte standards enregistrent l'état par hôte et ne peuvent pas être restaurés sur un nombre d'hôtes différent, tandis que ElasticIterator suit une seule position d'ensemble de données global qui peut être restaurée sur n'importe quel nombre de tranches actives.

Boucle d'entraînement élastique

Une boucle d'entraînement élastique coordonne la surveillance des tranches, la récupération après défaillance et la création d'instantanés en mémoire, comme illustré dans l'exemple de code suivant :

  • Votre fonction train crée une seule instance manager.Manager() et un Snapshotter Orbax en dehors du wrapper de nouvelle tentative, de sorte que le suivi des tranches et les instantanés de la mémoire hôte persistent entre les tentatives.
  • La décoration de la fonction run_training_loop interne avec @elastic_manager.elastic_retry attend les tranches opérationnelles, définit jax.default_device sur une tranche active, maintient le thread de surveillance en arrière-plan actif entre les étapes et relance la boucle lorsqu'une jax.errors.JaxRuntimeError se produit.
  • Lors de la récupération après un échec de tranche au début d'une tentative, votre code doit recompiler le maillage de l'appareil à partir de elastic_manager.active_devices et appeler snapshotter.load pour restaurer le dernier instantané sur les segmentations de maillage mises à jour, tout en laissant les tampons de l'appareil de la tentative ayant échoué hors du champ d'application.
  • À chaque étape de la boucle, votre code appelle régulièrement snapshotter.save pour sauvegarder l'état de l'entraînement et vérifie elastic_manager.available_inactive_slices pour détecter les tranches de remplacement opérationnelles. Lorsque des tranches de remplacement sont disponibles, votre code les ajoute à elastic_manager.active_slice_indices, recompile le maillage, réinitialise la boucle d'entraînement et place l'état sur le maillage développé.

L'exemple suivant montre comment implémenter une boucle d'entraînement élastique en combinant manager.Manager, elastic_retry et le Snapshotter en mémoire d'Orbax :

import jax
import orbax.checkpoint as ocp
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()


def train(
    final_step: int,
    snapshot_interval: int,  # in steps
    devices_per_slice: int,
):
  elastic_manager = manager.Manager()
  snapshotter = ocp.v1.training.pathways.Snapshotter(replica_axis_index=0)

  @elastic_manager.elastic_retry(
      minimum_slice_count=1,
      poll_interval=10,  # in seconds
      timeout=300,  # in seconds
      retry_policy=manager.ElasticRetryLimit(max_attempts=10),
  )
  def run_training_loop():
    mesh = create_mesh(elastic_manager, devices_per_slice)
    state, abstract_state, jitted_train_step, data_iterator = (
        initialize_training_loop(mesh)
    )
    step = 0

    # Restore from the latest in-memory snapshot after a slice failure.
    if snapshotter.latest is not None:
      snapshotter.join()
      restored = snapshotter.load({
          "state": abstract_state,
          "step": step,
          "iterator_state": data_iterator.get_state(),
      })
      state = restored["state"]
      step = int(restored["step"]) + 1
      data_iterator.set_state(restored["iterator_state"])

    while step < final_step:
      batch = next(data_iterator)
      state = jitted_train_step(state, batch)

      if step % snapshot_interval == 0:
        snapshotter.save(
            step,
            {
                "state": state,
                "step": step,
                "iterator_state": data_iterator.get_state(),
            },
        )

      # If replacement slices are healthy, reinitialize and place state onto the new mesh.
      if elastic_manager.available_inactive_slices:
        elastic_manager.active_slice_indices |= (
            elastic_manager.available_inactive_slices
        )
        mesh = create_mesh(elastic_manager, devices_per_slice)
        iterator_state = data_iterator.get_state()
        _, abstract_state, jitted_train_step, data_iterator = (
            initialize_training_loop(mesh)
        )
        data_iterator.set_state(iterator_state)
        state_shardings = jax.tree.map(lambda x: x.sharding, abstract_state)
        state = jax.device_put(state, state_shardings)

      step += 1

    return state

  state = run_training_loop()
  snapshotter.join()
  return state

Configurer le gestionnaire élastique et le comportement de nouvelle tentative

Vous instanciez manager.Manager(devices=None) une fois par processus. Si vous omettez devices, le gestionnaire utilise jax.devices() et regroupe les appareils en fonction de leur attribut slice_index.

L'instance Manager expose les propriétés suivantes pour inspecter la topologie du cluster pendant l'entraînement :

  • slice_to_devices : mappage de chaque slice_index à la séquence d'objets jax.Device appartenant à cette tranche.
  • all_slice_indices : frozenset de tous les indices de tranche de l'allocation.
  • active_slice_indices : frozenset d'indices de tranche participant au maillage de calcul actif.
  • active_devices : liste d'objets jax.Device appartenant aux tranches actives, classés par slice_index.
  • inactive_slice_indices : frozenset d'indices de tranche dans l'allocation qui ne participent pas au maillage de calcul.
  • available_inactive_slices : frozenset d'indices de tranche inactifs que le thread de surveillance en arrière-plan a détectés comme opérationnels et prêts à être réintégrés.
  • total_slice_count et active_slice_count : nombre total de tranches allouées et nombre de tranches actives.
  • default_device : appareil 0 d'une tranche active, que elastic_retry définit automatiquement comme appareil JAX par défaut pour chaque tentative.

Vous pouvez personnaliser le comportement des nouvelles tentatives et de la mise à l'échelle en transmettant les arguments suivants à @elastic_manager.elastic_retry :

  • minimum_slice_count : nombre minimal de tranches opérationnelles requises avant de lancer l'exécution initiale ou une nouvelle tentative. Si la valeur est définie sur None (valeur par défaut), elastic_retry attend toutes les tranches de l'allocation (total_slice_count), qui fonctionnent en mode pause et reprise. Si la valeur est définie sur un entier plus petit (tel que 1), elastic_retry fonctionne en mode de redimensionnement de l'instance répliquée.
  • poll_interval : intervalle en secondes entre les vérifications de l'état des tranches effectuées par le gestionnaire d'entraînement élastique (Manager) pendant l'entraînement. La valeur par défaut est 10.
  • timeout : nombre maximal de secondes d'attente avant que les tranches minimum_slice_count ne deviennent actives avant de générer une TimeoutError. La valeur par défaut est None (pas de délai avant expiration).
  • pre_callback : rappel facultatif exécuté au début de chaque tentative après la mise à jour de active_slice_indices et avant l'exécution de la fonction de boucle d'entraînement encapsulée.
  • on_elastic_event_callback : rappel facultatif exécuté après la détection d'un jax.errors.JaxRuntimeError de réduction, avant le début de la prochaine tentative. Vous pouvez utiliser ce rappel pour nettoyer les points de contrôle incomplets ou réinitialiser les itérateurs de données.
  • retry_policy : (attempt: int, error: Exception) -> bool appelable qui détermine s'il faut continuer à réessayer après un événement élastique. Utilisez manager.ElasticRetryLimit(max_attempts) pour limiter le nombre total de tentatives. Si aucune valeur n'est spécifiée, elastic_retry effectue de nouvelles tentatives indéfiniment.

Lors d'un événement élastique, JAX libère les tampons de l'appareil de la tentative ayant échoué une fois que le bloc de pile est déroulé et qu'aucune variable ne fait référence à ces tableaux. Évitez de conserver des références aux tableaux d'appareils d'une tentative infructueuse dans des variables globales ou des rappels de longue durée, afin que la mémoire soit libérée avant le début de la tentative suivante.

Instantanés en mémoire avec Orbax

La classe ocp.v1.training.pathways.Snapshotter (disponible dans orbax-checkpoint version 0.12.6 et ultérieures) conserve une sauvegarde en mémoire de votre état d'entraînement dans la mémoire pinned_host sur plusieurs tranches :

  • Snapshotter(replica_axis_index=0) : initialise l'outil d'instantanés et spécifie l'axe de maillage qui représente les instances répliquées de tranche de données parallèles. Tous les tableaux JAX de l'état enregistré doivent être répliqués le long de cet axe de maillage afin que les tranches survivantes conservent une copie complète de l'état du modèle et de l'optimiseur en cas d'échec d'une tranche. Si votre état contient des tableaux à un seul appareil, répliquez-les sur le maillage avant d'appeler save.
  • save(step, state) : transfère de manière asynchrone state vers la mémoire pinned_host et le divise en arbres par instance répliquée le long de replica_axis_index. Si un transfert d'instantané précédent est toujours en cours, save ignore la nouvelle requête pour éviter de bloquer la boucle d'entraînement.
  • load(abstract_state, reset_snapshot_state=True) : sonde les instances répliquées de l'hôte enregistrées pour identifier les tranches survivantes, concatène les instances répliquées opérationnelles et repartitionne l'état sur les partitionnements cibles spécifiés par abstract_state.
  • join() : bloque jusqu'à ce que toutes les opérations d'instantané en file d'attente soient terminées.
  • latest : renvoie les métadonnées (y compris step) de l'instantané le plus récent ou None si aucun instantané n'a été effectué.

Réduire la segmentation en cas d'échec de tranche

Lorsqu'une tranche de TPU échoue ou est préemptée pendant l'exécution, Pathways génère une exception jax.errors.JaxRuntimeError (telle que DATA_LOSS, DEADLINE_EXCEEDED, NOT_FOUND ou INTERNAL). Le décorateur elastic_retry intercepte cette exception et vérifie elastic.is_error_due_to_slice_down(error) :

  • Si l'erreur n'est pas causée par un événement de réduction, elastic_retry lève à nouveau l'exception immédiatement.
  • Si l'erreur est causée par un événement de réduction de tranche, elastic_retry arrête toute trace de profileur JAX active, exécute on_elastic_event_callback et évalue retry_policy(attempt, error).
  • Avant de lancer la tentative suivante, elastic_retry appelle une fois elastic.wait_for_slices, qui interroge l'état des tranches jusqu'à ce qu'au moins minimum_slice_count tranches soient actives, met à jour elastic_manager.active_slice_indices et exécute votre fonction de boucle d'entraînement sur les tranches survivantes.

Augmenter le partitionnement lorsque les segments se rejoignent

Lorsque votre charge de travail s'exécute sur un nombre réduit de tranches en mode de redimensionnement de l'instance répliquée, le thread de surveillance en arrière-plan exécute régulièrement des vérifications d'état sur elastic_manager.inactive_slice_indices. Lorsque les tranches réparées ou reprogrammées réussissent la vérification de l'état, le thread de surveillance ajoute leurs index de tranche à elastic_manager.available_inactive_slices.

Lorsque elastic_manager.available_inactive_slices n'est pas vide à la fin d'une étape d'entraînement, la charge de travail effectue un scaling à la hausse sur place sans perdre la progression de l'entraînement :

  • L'ajout des indices de tranche dans elastic_manager.available_inactive_slices à elastic_manager.active_slice_indices intègre les tranches réunies dans elastic_manager.active_devices et les efface de available_inactive_slices.
  • La recompilation du maillage de l'appareil à partir de elastic_manager.active_devices et la réinitialisation de l'étape d'entraînement et de l'itérateur de données adaptent le calcul à la topologie étendue.
  • L'appel de jax.device_put place l'état d'entraînement actuel sur les nouveaux partitionnements de maillage sans restauration à partir d'un instantané.

Remplacement à chaud

Le remplacement à chaud est une fonctionnalité de l'API GKE JobSet qui permet à un job de priorité plus élevée de reprendre rapidement les ressources d'un job de priorité inférieure, ce qui réduit les temps d'arrêt et garantit une récupération plus rapide.

Lorsqu'un JobSet est créé, GKE planifie la charge de travail sur plusieurs tranches, comme spécifié dans la configuration du JobSet. En cas de défaillance matérielle sur une ou plusieurs tranches, les pods concernés sont marqués comme ayant échoué. Si votre cluster GKE dispose d'une tranche de rechange exécutant une tâche de priorité inférieure, le système JobSet remappe la charge de travail de la tranche défaillante de la tâche de priorité supérieure sur cette tranche de rechange. Ce remappage prend généralement moins d'une minute.

Lorsqu'un JobSet redémarre, un remplacement à chaud peut se produire dans les situations suivantes :

  • Mode par défaut : si des tranches de TPU de rechange inactives sont disponibles dans le même cluster, le programmeur Kubernetes donne la priorité à la planification des jobs redémarrés sur ces tranches plutôt que d'attendre que les tranches défaillantes soient réparées. Ce comportement de planification permet une récupération plus rapide.
  • Charges de travail hétérogènes : dans les clusters exécutant plusieurs charges de travail avec une ressource Kubernetes PriorityClass configurée, un JobSet redémarré peut déclencher un remplacement à chaud. Si l'affinité du job redémarré correspond aux ressources d'un job de priorité inférieure, Kubernetes préempte le job de priorité inférieure afin que le job de priorité supérieure puisse démarrer sans attendre la réparation de la tranche. Par exemple, vous pouvez configurer vos pods de nœuds de calcul Pathways avec différentes priorités à l'aide d'une ressource PriorityClass.

Pour utiliser des priorités dans votre cluster, définissez un fichier manifeste PriorityClass :

kind: PriorityClass
metadata:
  name: high-prior-job
value: 2000
globalDefault: false
description: "This priority class should be used for high priority job."

Appliquez ce fichier YAML à votre cluster GKE :

kubectl apply -f high-prior-job.yaml

Ensuite, associez la nouvelle classe de priorité à votre job de nœud de calcul Pathways en ajoutant le texte suivant à la spécification de pod de votre pod pathways-worker.

priorityClassName: high-prior-job

Étapes suivantes