Entraînement résilient avec Pathways

Les chemins d'accès offrent des avantages en termes de résilience de plusieurs façons :

  • Suspendre/Reprendre : tolérance face aux interruptions planifiées telles que les notifications de préemption sans que l'utilisateur ait besoin d'écrire de code personnalisé pour gérer la préemption.
  • Entraînement élastique : tolérance face aux défaillances matérielles imprévues sans provoquer le 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 à la préemption des chemins d'accès est activée par défaut sur tous les déploiements cloud, et les tâches d'accélérateur Pathways écoutent ces notifications.

    Lorsqu'une notification 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 tel 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 relisant son état persistant.

    Si la charge de travail ne peut pas être restaurée, Pathways arrête le job d'accélérateur et transmet l'échec à votre job si l'entraînement élastique est configuré. Si l'entraînement Elastic 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 typiques définies à l'aide de JAX reposent 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. Ceux-ci ne sont pas restaurables.

    Entraînement élastique

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

    • Détection des défaillances : lorsqu'une défaillance matérielle se produit (par exemple, un nœud de calcul TPU plante), le système Pathways la détecte et en informe le job d'entraînement de l'utilisateur par le biais d'une exception la prochaine fois que les données situées sur ce matériel sont consultées. Cette notification ne plante pas votre charge de travail. Elle permet à votre code de gérer la notification et de reconfigurer vos ressources pour continuer le traitement ou quitter correctement.
    • Gestionnaire d'élasticité défini par l'utilisateur : le code du modèle de l'utilisateur doit être en mesure de gérer cette exception. C'est ce qui fait qu'il s'agit d'une "récupération spécifique au modèle".
      • Instantanés : l'approche la plus courante consiste à enregistrer périodiquement 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'adapter au nombre de tranches disponibles. Par exemple, si une tranche cesse de fonctionner, vous pouvez réduire le nombre de tranches actives d'une unité jusqu'à ce qu'un remplacement soit disponible. Pour en savoir plus, consultez Gestionnaire Elastic.
      • Mises à jour du graphique de données/de 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 des pathways dans la récupération : les pathways fournissent les primitives permettant 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 la nouvelle tranche disponible. Votre code peut ensuite être reconfiguré 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 opérationnelles 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 (par exemple, lorsqu'une tranche de TPU 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 sorties du gestionnaire doivent correspondre aux arguments et aux 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 des blocs try-except qui gèrent les jax.errors.JaxRuntimeError.
    3. Dans votre gestionnaire jax.errors.JaxRuntimeError, appelez maybe_reshard_down. Le gestionnaire élastique effectuera un resharding à la baisse si l'erreur est liée à un événement élastique ou la relèvera autrement.
    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 des instantanés est déterminée par la période d'instantané. La période du snapshot a une incidence sur le nombre moyen de pas perdus 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 prendra en charge. max_reshard_retry_count spécifie le nombre de fois où le gestionnaire élastique doit réessayer de repartitionner. 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 votre gestionnaire élastique lors d'un événement élastique.

    Réduire le partitionnement

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

    Augmenter le partitionnement

    En fonction de la configuration du gestionnaire élastique et s'il existe des tranches indisponibles, Pathways vérifie si des tranches supplémentaires sont devenues disponibles. Si c'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 en boucle jusqu'à ce que l'opération réussisse ou que le nombre maximal de tentatives soit atteint. En cas de re-partitionnement, les valeurs de retour du gestionnaire élastique sont transmises à l'appelant. Sinon, None est renvoyé.

    Échange à chaud

    L'échange à 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 relais des ressources d'un job de priorité inférieure, ce qui minimise les temps d'arrêt et assure 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é. Lorsque vous reprogrammez ce JobSet, si vous avez choisi de conserver un emplacement libre dans votre cluster GKE qui pourrait être utilisé pour une tâche de priorité inférieure, le système JobSet remappera la charge de travail de l'emplacement ayant échoué de la tâche de priorité supérieure sur l'emplacement libre utilisé par la tâche de priorité inférieure dans le même cluster GKE. Cette réaffectation prend généralement moins d'une minute.

    Lors du redémarrage de JobSet, l'échange à chaud peut se produire dans les situations suivantes :

    1. Mode par défaut : si des tranches TPU de secours et inactives sont disponibles dans le même cluster, le planificateur Kubernetes donnera la priorité à la planification des jobs redémarrés sur ces tranches plutôt que d'attendre la réparation des tranches ayant échoué. 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 échange à 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 worker Pathways en ajoutant le texte suivant au fichier podspec de votre pod pathways-worker.

    priorityClassName: high-prior-job

    Étapes suivantes