Entraînement résilient avec Pathways

Cette page explique comment configurer des fonctionnalités de résilience pour les charges de travail Pathways sur Google Kubernetes Engine (GKE) afin que les jobs d'entraînement à grande échelle puissent récupérer après des interruptions planifiées et des défaillances matérielles non planifiées.

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 avis de préemption, sans nécessiter de code personnalisé de gestion de la préemption.
• 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.
• Points de contrôle multiniveaux : permet à votre charge de travail d'enregistrer des points de contrôle locaux et répliqués par des pairs à haute fréquence sur les travailleurs TPU à l'aide de Python colocalisé, avec des sauvegardes Cloud Storage asynchrones pour réduire le temps de récupération.
• Remplacement à chaud : permet à une charge de travail GKE JobSet de priorité supérieure de prendre rapidement le relais des tranches de TPU de rechange ou de priorité inférieure en cas de défaillance matérielle.

Avant de commencer

Avant de commencer, effectuez les tâches suivantes :

• Configurer Cluster Toolkit
• Installer les outils Kubernetes
• Installer Google Cloud CLI
• Activer l'API Kubernetes Engine

Si vous utilisez l'entraînement élastique, installez les packages suivants dans l'image de conteneur de votre charge de travail :

• pathwaysutils
• Orbax (orbax-checkpoint) : >= 0.12.6 (requis pour les instantanés en mémoire avec ocp.v1.training.pathways.Snapshotter)
• Grain (grain) : obligatoire si vous utilisez grain.experimental.ElasticIterator

Si vous utilisez la création de points de contrôle multiniveaux, remplissez également les conditions préalables suivantes :

• Configurez votre cluster GKE pour les points de contrôle multiniveaux en suivant les instructions de la page Entraîner des modèles de machine learning à grande échelle sur GKE avec des points de contrôle multiniveaux :
  • Activez les modules complémentaires Multi-Tier Checkpointing et Cloud Storage FUSE (GcsFuseCsiDriver), ainsi que Workload Identity Federation for GKE, sur votre cluster.
  • Créez un bucket Cloud Storage avec l'espace de noms hiérarchique activé.
  • Attribuez les rôles IAM (Identity and Access Management) roles/storage.objectUser et roles/storage.insightsCollectorService au compte de service Kubernetes de point de contrôle multiniveau.
  • Déployez une ressource personnalisée CheckpointConfiguration en spécifiant votre bucket Cloud Storage, le sélecteur de nœuds TPU et la taille du volume en mémoire (inMemoryVolumeSize).
• Créez et transférez une image de conteneur side-car Python colocalisée, par exemple à l'aide de l'exemple de fichier Dockerfile side-car. Dans vos images de charge de travail et de side-car, utilisez des versions jax et jaxlib correspondantes qui s'alignent sur votre version de Pathways et répondent aux exigences de version minimale suivantes :
  • Version d'Orbax : >= 0.12.4
  • Version de JAX (y compris Python colocalisé) : >= 0.10.1

Suspendre/Reprendre

GKE envoie généralement une notification de préemption à un pod d'accélérateur avant de le préempter. Pathways active par défaut la tolérance à la préemption 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 possible, Pathways tente de suspendre de manière transparente votre charge de travail de ML en écrivant son état actuel dans un espace de 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 charge de travail 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 reposent sur des composants XLA Pathways 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 et 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 : si vous définissez minimum_slice_count sur None ou sur le nombre total de tranches, l'entraînement est mis en pause lorsqu'une tranche échoue et reprend une fois que toutes les tranches sont à nouveau opérationnelles. É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 : si vous définissez minimum_slice_count sur une valeur inférieure au nombre total de tranches, l'entraînement se poursuit sur les tranches actives restantes lorsqu'une tranche échoue 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 saines dans elastic_manager.available_inactive_slices. Votre boucle d'entraînement peut ensuite recompiler le maillage d'appareils 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/cloud-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. En revanche, ElasticIterator suit une seule position d'ensemble de données global que vous pouvez restaurer sur un nombre illimité de segments actifs.

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.
• L'application de @elastic_manager.elastic_retry à run_training_loop configure la fonction pour qu'elle attende des tranches opérationnelles, définisse jax.default_device sur une tranche active, maintienne le thread de surveillance en arrière-plan actif et réessaie sur un jax.errors.JaxRuntimeError.
• Au début d'une tentative de récupération, 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. Laissez les tampons de l'appareil de la tentative ayant échoué sortir du champ d'application afin que JAX puisse libérer leur mémoire.
• À 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 qui s'exécute au début de chaque tentative après la mise à jour de active_slice_indices par le gestionnaire et avant l'exécution de la fonction de boucle d'entraînement encapsulée.
• on_elastic_event_callback : rappel facultatif qui s'exécute après que elastic_retry a intercepté 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). Dans une fonction décorée avec @elastic_manager.elastic_retry, 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 prochaine tentative, elastic_retry appelle elastic.wait_for_slices une fois pour interroger l'état des tranches jusqu'à ce qu'au moins minimum_slice_count tranches soient actives. Il met ensuite à jour elastic_manager.active_slice_indices et exécute votre fonction de boucle d'entraînement sur les tranches restantes.

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é.

Points de contrôle multiniveaux

Les jobs d'entraînement à grande échelle subissent des interruptions fréquentes, et l'écriture ou la restauration de points de contrôle directement vers et depuis Cloud Storage peut ralentir la progression de l'entraînement. La création de points de contrôle multicouches réduit les coûts de création de points de contrôle et le temps de récupération en coordonnant trois niveaux de stockage sur GKE :

• Disque RAM local : les nœuds de calcul TPU écrivent des points de contrôle à haute fréquence sur un volume CSI en mémoire (multitier-checkpoint.csi.storage.gke.io) installé sur chaque nœud de calcul. Ce volume persiste lors des redémarrages de pods sur le même nœud.
• Réplication de pairs dans le cluster : le service de réplication GKE réplique les segments de points de contrôle locaux sur les nœuds de calcul TPU pairs du cluster.
• Sauvegarde Cloud Storage : le service de réplication sauvegarde de manière asynchrone les points de contrôle locaux dans un bucket Cloud Storage avec l'espace de noms hiérarchique activé, sans bloquer les étapes d'entraînement TPU.

Contrairement à JAX multi-contrôleur, où chaque hôte TPU exécute son propre script d'entraînement, Pathways exécute un seul contrôleur sur le pod pathways-head et répartit le travail sur les pods pathways-worker. Pour lire et écrire des segments de point de contrôle directement sur le disque RAM local de chaque hôte TPU sans acheminer les données de point de contrôle via le contrôleur unique, Pathways utilise Python colocalisé. Chaque pod worker exécute un conteneur colocated-python-sidecar avec le conteneur pathways-worker, et les deux conteneurs installent le volume CSI de point de contrôle à plusieurs niveaux et un répertoire de mémoire partagée (/tmp/sidecar). Lors de l'initialisation et de l'entraînement, Orbax distribue la configuration, les sauvegardes et les restaurations des points de contrôle du contrôleur Pathways aux side-cars de nœuds de calcul sur chaque hôte TPU. La création de points de contrôle multiniveau sur Pathways est également compatible avec l'entraînement élastique lors de la réduction ou de l'augmentation du nombre de tranches de TPU actives.

Pathways est compatible avec les points de contrôle multiniveaux avec MaxText. Pour obtenir des instructions de configuration, consultez le guide sur les points de contrôle multiniveaux MaxText.

Configurer votre charge de travail Pathways

Vous pouvez configurer une charge de travail Pathways pour les points de contrôle multiniveaux à l'aide de Cluster Toolkit (gcluster job submit) ou en configurant un fichier manifeste JobSet Pathways personnalisé.

Utiliser Cluster Toolkit

Lorsque vous envoyez un job avec gcluster job submit --pathways, transmettez les options --gke-mtc-enabled, --gke-mtc-ramdisk-dir et --pathways-colocated-python-sidecar-image :

gcluster job submit \
    --pathways \
    --gke-mtc-enabled \
    --gke-mtc-ramdisk-dir="/tmp/mtc_checkpoints" \
    --pathways-colocated-python-sidecar-image="SIDECAR_IMAGE" \
    ...

Remplacez SIDECAR_IMAGE par l'URI de votre image de conteneur side-car Python colocalisée.

Lorsque vous définissez ces options, gcluster configure automatiquement les éléments suivants :

• Transmet --sidecar_name=external au conteneur pathways-proxy.
• Exécute le conteneur colocated-python-sidecar dans chaque pod worker et installe un répertoire /tmp/sidecar partagé (--cloud_pathways_sidecar_shm_directory=/tmp/sidecar) sur pathways-worker et colocated-python-sidecar.
• Installe le volume CSI multitier-checkpoint.csi.storage.gke.io au niveau du chemin d'accès spécifié par --gke-mtc-ramdisk-dir dans pathways-worker et colocated-python-sidecar sur les pods worker, mais pas sur le pod pathways-head.

Utiliser un fichier manifeste JobSet Pathways personnalisé

Si vous gérez votre propre fichier manifeste JobSet Pathways au lieu d'utiliser gcluster, transmettez --sidecar_name=external à pathways-proxy. Dans la spécification de pod worker, installez le volume de mémoire side-car partagé et le volume CSI multitier-checkpoint.csi.storage.gke.io sur les conteneurs pathways-worker et colocated-python-sidecar :

containers:
- name: pathways-worker
  args:
  - --cloud_pathways_sidecar_shm_directory=/tmp/sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
- name: colocated-python-sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
volumes:
- name: shared-tmp
  emptyDir:
    medium: Memory
- name: checkpoint-mount
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

Utiliser la création de points de contrôle multiniveau Orbax dans une charge de travail JAX personnalisée

Dans un script d'entraînement JAX personnalisé exécuté sur Pathways, utilisez le package orbax.checkpoint.experimental.emergency.multi_tier_checkpointing et définissez use_colocated_python=True dans initialize_multi_tier_checkpointing et ReplicatorCheckpointManagerOptions :

from etils import epath
import jax
import orbax.checkpoint as ocp
from orbax.checkpoint.experimental.emergency.multi_tier_checkpointing import (
    initialization,
    replicator_checkpoint_manager,
)
import pathwaysutils

pathwaysutils.initialize()

local_checkpoint_dir = epath.Path("/tmp/mtc_checkpoints")

# Initialize multi-tier checkpointing across worker sidecars.
# When using elastic training, pass the active slice devices to `devices`
# on initial setup and whenever the active slice count changes.
initialization.initialize_multi_tier_checkpointing(
    local_checkpoint_directory=local_checkpoint_dir,
    backup_interval_minutes=10,
    num_slices=num_slices,
    run_name=run_name,
    data_parallelism=data_parallelism,
    use_colocated_python=True,
    devices=jax.devices(),
)

options = replicator_checkpoint_manager.ReplicatorCheckpointManagerOptions(
    save_interval_steps=20,
    use_colocated_python=True,
)
checkpoint_manager = (
    replicator_checkpoint_manager.ReplicatorCheckpointManager(
        local_checkpoint_dir,
        options=options,
        global_mesh=mesh,
    )
)

# Restore from the latest checkpoint if one exists.
start_step = 1
latest_step = checkpoint_manager.latest_step()
if latest_step is not None:
  restore_args = jax.tree.map(
      lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding),
      state,
  )
  restored = checkpoint_manager.restore(
      latest_step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeRestore(
              item=state,
              restore_args=restore_args,
          )
      ),
  )
  state = restored.state
  start_step = latest_step + 1

# Step 0 is reserved by the replicator service to indicate "no checkpoint".
for step in range(start_step, final_step + 1):
  state = jitted_train_step(state)
  checkpoint_manager.save(
      step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeSave(item=state)
      ),
  )

checkpoint_manager.wait_until_finished()
checkpoint_manager.close()

Remplacement à chaud

Le remplacement à chaud est une fonctionnalité de l'API JobSet de GKE qui permet à un job de priorité supérieure de reprendre rapidement les ressources d'un job de priorité inférieure, ce qui réduit les temps d'arrêt et accélère la récupération.

Lorsque vous créez un JobSet, 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, GKE marque les pods concernés comme ayant échoué. Si votre cluster GKE dispose d'une tranche de rechange exécutant une tâche de priorité inférieure, le contrôleur 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 et enregistrez-le sous le nom high-prior-job.yaml :

apiVersion: scheduling.k8s.io/v1
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 charge de travail Pathways. Si vous envoyez votre charge de travail à l'aide de Cluster Toolkit, transmettez --priority=high-prior-job à la commande gcluster job submit, qui définit priorityClassName sur les pods pathways-head et worker. Si vous écrivez votre propre fichier manifeste JobSet, définissez priorityClassName dans le modèle de pod du job répliqué worker, qui exécute le conteneur pathways-worker :

replicatedJobs:
- name: worker
  template:
    spec:
      template:
        spec:
          priorityClassName: high-prior-job

Si vous avez créé votre cluster à l'aide d'un plan TPU Cluster Toolkit avec Pathways activé, le cluster inclut déjà les classes de priorité very-low, low, medium, high et very-high. Vous pouvez donc transmettre l'une d'elles à --priority au lieu de créer la vôtre.

Étapes suivantes

• Créer un cluster GKE avec Pathways
• Exécuter une charge de travail par lot avec Pathways
• Mode interactif de Pathways
• Porter des charges de travail JAX vers Pathways
• Résoudre les problèmes liés à Pathways sur Cloud
• Entraîner des modèles de machine learning à grande échelle sur GKE avec des points de contrôle multiniveaux
• Guide sur les points de contrôle multiniveaux MaxText
• Utiliser des points de contrôle à plusieurs niveaux pour les jobs d'entraînement d'IA volumineux