Entraînement résilient avec Pathways

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

  • Suspension-reprise : tolérance face aux interruptions planifiées telles que les avis de préemption, sans que l'utilisateur ait besoin d'écrire un code de gestion de préemption personnalisé.
  • Entraînement élastique : tolérance face aux pannes matérielles non planifiées sans provoquer de plantage du client, mais nécessitant que les utilisateurs écrivent un code de récupération spécifique au modèle.

Avant de commencer

Vérifiez que vous disposez bien des éléments suivants :

Suspendre/Reprendre

En règle générale, GKE envoie 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 si la charge de travail actuelle est restaurable, c'est-à-dire si Pathways 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'évince 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, reposent sur des composants Pathways avec état, qui ne sont pas restaurables.

Entraînement élastique

L'entraînement Elastic permet à votre job d'entraînement de se poursuivre même en cas de défaillance matérielle. Pour ce faire, il combine les fonctionnalités du système Pathways et la logique de récupération de modèle définie par l'utilisateur :

  • Détection d'une défaillance : lorsqu'une défaillance matérielle se produit (par exemple, un plantage d'un nœud de calcul TPU), le système Pathways la détecte et en informe le job d'entraînement de l'utilisateur via une exception lors du prochain accès aux données situées sur ce matériel. Cette notification n'entraîne pas le plantage de votre charge de travail. Elle permet à votre code de gérer la notification et de reconfigurer vos ressources pour continuer le traitement ou s'arrêter correctement.
  • Gestionnaire d'élasticité défini par l'utilisateur : le code du modèle de l'utilisateur doit pouvoir gérer cette exception. C'est ce qui en fait une "récupération spécifique au modèle".
    • Instantanés : l'approche la plus courante consiste à enregistrer régulièrement des instantanés de l'état de votre modèle. En cas d'échec, vous pouvez charger l'instantané le plus récent pour reprendre l'entraînement.
    • Reconfiguration : vous devrez probablement reconfigurer votre job d'entraînement pour l'ajuster au nombre de tranches disponibles. Par exemple, si une tranche cesse de fonctionner, vous pouvez réduire d'une unité le nombre de tranches actives jusqu'à ce qu'un remplacement soit disponible. Pour en savoir plus, consultez Gestionnaire Elastic.
    • Mises à jour du graphique de données/calcul : votre code doit gérer toute modification du nombre d'appareils disponibles pour votre calcul en recréant le graphique de calcul si nécessaire. Cela peut impliquer de repartitionner les données ou de recompiler votre modèle.
  • Rôle de Pathways dans la récupération : Pathways fournit les primitives permettant de prendre en charge la reconfiguration définie par l'utilisateur :
    • Remplacement de tranche : si une tranche défaillante est remplacée, le client peut être informé une fois que la nouvelle tranche est disponible. Votre code peut ensuite se reconfigurer pour utiliser cette nouvelle tranche.
    • Récupération transparente : Pathways gère les détails de bas niveau de la récupération, comme le rétablissement des connexions aux parties saines du cluster.
  • Utilitaires dans pathwaysutils : ensemble d'utilitaires Pathways définis dans pathways-utils.

Implémenter un gestionnaire élastique

La majeure partie du code que vous devrez écrire se trouvera dans un gestionnaire élastique défini par l'utilisateur. Ce gestionnaire réagit aux événements élastiques (comme une tranche de TPU qui devient indisponible) en recréant le maillage et en réinitialisant la boucle d'entraînement.

Chaque charge de travail est unique. La complexité du gestionnaire élastique peut évoluer en fonction de la complexité de la charge de travail. Les entrées et les sorties du gestionnaire doivent être les arguments et les valeurs de retour minimaux nécessaires pour réinitialiser la boucle d'entraînement.

def elastic_handler(elastic_utils, *args, **kwargs):
  mesh = initialize_mesh(**kwargs["mesh_kwargs"])
  initial_state, initial_step, jitted_train_step, other_variables =
      initialize_training_loop(mesh, **kwargs["initialize_training_loop_kwargs"])

  step, snapshot = elastic_utils.get_next_snapshot()
  state = initial_state.replace(**snapshot)

  return state, step, mesh, jitted_train_step, other_variables

Mettre à jour votre boucle d'entraînement

Vous devez apporter les modifications suivantes à votre boucle d'entraînement :

  1. Créer un gestionnaire élastique
  2. Encapsulez votre boucle d'entraînement dans un bloc try-except qui gère les jax.errors.JaxRuntimeError.
  3. Dans votre gestionnaire jax.errors.JaxRuntimeError, appelez maybe_reshard_down. Le gestionnaire élastique effectuera un re-sharding si l'erreur est liée à un événement élastique, ou la relancera dans le cas contraire.
  4. Appelez maybe_snapshot et maybe_reshard_up à la fin de la boucle d'entraînement.
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()

def initialize_mesh(**kwargs):
  ...


def initialize_training_loop(**kwargs):
  ...


def train_loop(
    final_step,
    elastic_manager,
    mesh_kwargs,
    initialize_training_loop_kwargs,
):
  mesh = initialize_mesh(**mesh_kwargs)
  initial_state, initial_step, jitted_train_step, other_variables =
      initialize_training_loop(mesh, **initialize_training_loop_kwargs)

  step = initial_step
  while step < final_step:
    try:
      state = jitted_train_step(state)

      elastic_manager.maybe_snapshot(step=step, snapshot=state)
      handler_returns = elastic_manager.maybe_reshard_up(
          step=step,
          snapshot=state,
          elastic_handler=elastic_handler,
          handler_args=(),
          handler_kwargs=dict(
              mesh_kwargs=mesh_kwargs,
              initialize_training_loop_kwargs=initialize_training_loop_kwargs,
          ),
      )
      if handler_returns:
        state, step, mesh, jitted_train_step, other_variables = handler_returns
      step += 1
    except jax.errors.JaxRuntimeError as error:
      handler_returns = elastic_manager.maybe_reshard_down(
          error=error,
          elastic_handler=elastic_handler,
          handler_args=(),
          handler_kwargs=dict(
              mesh_kwargs=mesh_kwargs,
              initialize_training_loop_kwargs=initialize_training_loop_kwargs,
          ),
      )
      if handler_returns:
        state, step, mesh, jitted_train_step, other_variables = handler_returns

  return state


def main():
  elastic_manager = manager.Manager(
      devices=jax.devices(),
      snapshot_period=10,
      snapshot_buffer_size=1,
      reshard_check_period=5,
      max_elastic_down_event_count=10,
      max_reshard_retry_count=3,
  )

  train_loop(100, elastic_manager, {}, {})

Configurer le gestionnaire Elastic

Le gestionnaire élastique peut être configuré de différentes manières. La fréquence de création d'instantanés est déterminée par la période d'instantané. La période d'instantané affecte le nombre moyen d'étapes perdues en raison d'un événement élastique. La période de vérification du repartitionnement détermine la fréquence à laquelle votre boucle d'entraînement interroge la disponibilité des tranches. max_elastic_down_event_count vous permet de définir le nombre d'événements élastiques dus à la perte de tranche que votre boucle d'entraînement acceptera. max_reshard_retry_count spécifie le nombre de fois où le gestionnaire élastique doit réessayer de resegmenter. Le gestionnaire est un objet singleton et ne doit être créé qu'une seule fois.

Instantanés

En fonction de la configuration du gestionnaire élastique, la fonction peut créer un instantané des données dans la mémoire de l'hôte, qui sera disponible pour être utilisé par votre gestionnaire élastique lors d'un événement élastique.

Réduire le partitionnement

Après avoir intercepté un jax.errors.JaxRuntimeError, Pathways vérifie si l'erreur est due à un événement élastique causé par une tranche perdue. Si tel est le cas, il appelle le gestionnaire élastique dans une boucle jusqu'à ce que l'opération réussisse ou que le nombre maximal de nouvelles tentatives soit atteint. Si l'erreur n'est pas due à un événement élastique, elle sera de nouveau générée. Les valeurs renvoyées du gestionnaire élastique sont transmises à l'appelant.

Augmenter le partitionnement

En fonction de la configuration du gestionnaire élastique et si des tranches ne sont pas disponibles, Pathways vérifie si d'autres tranches sont devenues disponibles. Si tel est le cas, il enregistre immédiatement un instantané (si un instantané préexistant pour l'étape actuelle n'a pas déjà été pris) et appelle le gestionnaire élastique dans une boucle jusqu'à ce que l'opération réussisse ou que le nombre maximal de tentatives soit atteint. Si un re-sharding se produit, les valeurs de retour du gestionnaire élastique sont transmises à l'appelant. Sinon, None est renvoyé.

Remplacement à chaud

Le remplacement à chaud fait référence à une fonctionnalité de l'API GKE JobSet qui permet à un job de priorité plus élevée de prendre rapidement le contrôle des 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é. Lors de la reprogrammation de ce JobSet, si vous avez choisi de conserver une tranche de réserve dans votre cluster GKE qui pourrait être utilisée pour un job de priorité inférieure, le système JobSet remappera la charge de travail de la tranche défaillante du job de priorité supérieure sur la tranche de réserve utilisée par le job de priorité inférieure dans le même cluster GKE. Ce remappage prend généralement moins d'une minute.

Lors du redémarrage de JobSet, un remplacement à chaud peut se produire dans les situations suivantes :

  1. Mode par défaut : si des tranches de TPU 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. Cela permet une récupération plus rapide.
  2. Charges de travail hétérogènes : dans les clusters exécutant plusieurs charges de travail avec une PriorityClass Kubernetes 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, ce qui permet au job de priorité supérieure de démarrer immédiatement. Par exemple, vous pouvez configurer vos pods de nœuds de calcul Pathways avec différentes priorités à l'aide de PriorityClass.

Pour utiliser des priorités dans votre cluster, définissez une classe de priorité, par exemple :

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