Resilienztraining mit Pathways

Pathways bietet folgende Vorteile in Bezug auf die Ausfallsicherheit:

Suspend-Resume: Toleranz bei geplanten Unterbrechungen wie Benachrichtigungen über das Vorrecht, ohne dass der Nutzer benutzerdefinierten Code für die Vorrechtsbehandlung schreiben muss.
Elastisches Training: Toleranz bei ungeplanten Hardwarefehlern, ohne dass der Client abstürzt. Nutzer müssen jedoch modellspezifischen Wiederherstellungscode schreiben.

Hinweis

Sie benötigen Folgendes:

Suspend-Resume

Normalerweise sendet GKE einen Hinweis auf vorzeitiges Beenden an einen Accelerator-Pod, bevor der Pod vorzeitig beendet wird. Die Pathways-Unterbrechungstoleranz ist standardmäßig für alle Cloud-Bereitstellungen aktiviert und Pathways-Beschleunigerjobs reagieren auf diese Benachrichtigungen.

Wenn eine Benachrichtigung über die Unterbrechung eingeht, prüft Pathways zuerst, ob die aktuelle Arbeitslast wiederhergestellt werden kann, d. h., ob Pathways die Arbeitslast transparent speichern und wiederherstellen kann. Wenn ja, wird versucht, Ihre ML-Arbeitslast transparent zu pausieren, indem der aktuelle Status in einem persistenten Speicher wie Cloud Storage gespeichert wird, bevor GKE Ihre Beschleunigerjobs entfernt. Wenn GKE Ihre Jobs später neu plant, setzt Pathways Ihre ML-Arbeitslast fort, indem der gespeicherte Status gelesen wird.

Wenn die Arbeitslast nicht wiederhergestellt werden kann, beendet Pathways den Accelerator-Job und leitet den Fehler an Ihren Job weiter, wenn Elastic Training konfiguriert ist. Wenn das elastische Training nicht konfiguriert ist, startet GKE die gesamte Arbeitslast basierend auf der Neustartrichtlinie für JobSet neu.

Typische ML-Arbeitslasten, die mit JAX definiert werden, basieren auf zustandslosen Pathways-XLA-Komponenten, die mit einem HBM-Snapshot (High Bandwidth Memory) wiederhergestellt werden können. Bestimmte ML-Arbeitslasten, z. B. solche, die mit der JAX-Python-API für die gemeinsame Ausführung definiert werden, basieren auf zustandsbehafteten Pathways-Komponenten. Diese können nicht wiederhergestellt werden.

Elastisches Training

Durch elastisches Training kann Ihr Trainingsjob auch bei Hardwarefehlern fortgesetzt werden. Dies wird durch eine Kombination aus den Funktionen des Pathways-Systems und der benutzerdefinierten Logik zur Modellwiederherstellung erreicht:

Erkennung von Fehlern: Wenn ein Hardwarefehler auftritt (z. B. ein TPU-Worker abstürzt), erkennt das Pathways-System dies und benachrichtigt den Trainingsjob des Nutzers durch eine Ausnahme, wenn das nächste Mal auf Daten zugegriffen wird, die sich auf dieser Hardware befinden. Diese Benachrichtigung führt nicht zum Absturz Ihrer Arbeitslast. Ihr Code kann die Benachrichtigung verarbeiten und Ihre Ressourcen so neu konfigurieren, dass die Verarbeitung fortgesetzt oder ordnungsgemäß beendet wird.
Benutzerdefinierte Elastizitätshandler: Der Modellcode des Nutzers muss diese Ausnahme verarbeiten können. Daher spricht man von „modellbezogener Wiederherstellung“.
- Snapshots: Der gängigste Ansatz besteht darin, regelmäßig Snapshots des Status Ihres Modells zu speichern. Wenn ein Fehler auftritt, können Sie das Training mit dem letzten Snapshot fortsetzen.
- Neukonfiguration: Sie müssen Ihren Trainingsjob wahrscheinlich neu konfigurieren, um die Anzahl der verfügbaren Slices anzupassen. Wenn beispielsweise ein Slice nicht mehr funktioniert, können Sie die Anzahl der aktiven Slices um eins reduzieren, bis ein Ersatz verfügbar ist. Weitere Informationen finden Sie unter Elastic Handler.
- Aktualisierungen des Daten-/Berechnungsgraphen: Ihr Code muss alle Änderungen an der Anzahl der für die Berechnung verfügbaren Geräte verarbeiten, indem er den Berechnungsgraphen nach Bedarf neu erstellt. Dazu müssen Sie möglicherweise Daten neu partitionieren oder Ihr Modell neu kompilieren.
Rolle von Pathways bei der Wiederherstellung: Pathways bietet die Primitiven zur Unterstützung der benutzerdefinierten Neukonfiguration:
- Segmentersetzung: Wenn ein fehlgeschlagenes Segment ersetzt wird, kann der Client benachrichtigt werden, sobald das neue Segment verfügbar ist. Ihr Code kann dann neu konfiguriert werden, um diesen neuen Slice zu verwenden.
- Transparente Wiederherstellung: Pathways übernimmt die Details der Wiederherstellung auf niedriger Ebene, z. B. das Wiederherstellen von Verbindungen zu den fehlerfreien Teilen des Clusters.
Hilfsprogramme in pathwaysutils: Eine Reihe von Pathways-Hilfsprogrammen, die in pathways-utils definiert sind.

Elastischen Handler implementieren

Der größte Teil des Codes, den Sie schreiben müssen, befindet sich in einem benutzerdefinierten elastischen Handler. Dieser Handler reagiert auf elastische Ereignisse (z. B. wenn eine TPU-Slice nicht mehr verfügbar ist), indem er das Mesh neu erstellt und die Trainingsschleife neu initialisiert.

Jede Arbeitslast ist einzigartig. Die Komplexität des elastischen Handlers kann mit der Komplexität der Arbeitslast skaliert werden. Die Ein- und Ausgaben des Handlers sollten die Mindestargumente und Rückgabewerte sein, die zum Reinitialisieren der Trainingsschleife erforderlich sind.

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

Trainingsschleife aktualisieren

Sie müssen die folgenden Änderungen an Ihrem Trainingszyklus vornehmen:

Elastischen Manager erstellen
Kapseln Sie Ihre Trainingsschleife in try-except-Blöcken, die jax.errors.JaxRuntimeErrors verarbeiten.
Rufen Sie in Ihrem jax.errors.JaxRuntimeError-Handler maybe_reshard_down auf. Der Elastic Manager führt ein Resharding durch, wenn der Fehler mit einem elastischen Ereignis zusammenhängt. Andernfalls wird er noch einmal ausgelöst.
Rufen Sie maybe_snapshot und maybe_reshard_up am Ende des Trainingsdurchlaufs auf.

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, {}, {})

Elastic Manager konfigurieren

Der elastische Manager kann auf verschiedene Arten konfiguriert werden. Die Häufigkeit der Snapshots wird durch den Snapshot-Zeitraum bestimmt. Der Snapshot-Zeitraum wirkt sich auf die durchschnittliche Anzahl der Schritte aus, die aufgrund eines elastischen Ereignisses verloren gegangen sind. Der Zeitraum für die Reshard-Prüfung bestimmt, wie oft Ihr Trainingslauf nach der Verfügbarkeit von Slices sucht. Mit max_elastic_down_event_count können Sie die Anzahl der elastischen Ereignisse festlegen, die aufgrund von Slice-Verlust in Ihrem Trainingslauf unterstützt werden. Der Parameter max_reshard_retry_count gibt an, wie oft der Elastic Manager die Resharding-Vorgänge wiederholen soll. Der Manager ist ein Singleton-Objekt und sollte nur einmal erstellt werden.

Snapshots

Je nach Konfiguration des Elastic Manager kann die Funktion Daten im Hostspeicher erfassen, die während eines elastischen Ereignisses von Ihrem elastischen Handler verwendet werden können.

Sharding reduzieren

Nachdem ein jax.errors.JaxRuntimeError abgefangen wurde, prüft Pathways, ob der Fehler auf ein elastisches Ereignis aufgrund eines verlorenen Slice zurückzuführen ist. Falls ja, wird der Elastic-Handler in einer Schleife aufgerufen, bis der Vorgang erfolgreich ist oder die maximale Anzahl von Wiederholungsversuchen erreicht wurde. Wenn der Fehler nicht auf ein elastisches Ereignis zurückzuführen ist, wird er noch einmal ausgelöst. Die Rückgabewerte des Elastic-Handlers werden an den Aufrufer übergeben.

Anzahl der Shards erhöhen

Anhand der Konfiguration des Elastic Managers und falls nicht verfügbare Slices vorhanden sind, wird in Pathways geprüft, ob zusätzliche Slices verfügbar geworden sind. Wenn ja, wird sofort ein Snapshot gespeichert (sofern noch kein Snapshot für den aktuellen Schritt vorhanden ist) und der Elastic-Handler wird in einer Schleife aufgerufen, bis der Vorgang erfolgreich abgeschlossen ist oder die maximale Anzahl von Wiederholungsversuchen erreicht wurde. Wenn eine erneute Shardierung erfolgt, werden die Rückgabewerte des Elastic-Handlers an den Aufrufer weitergeleitet. Andernfalls wird None zurückgegeben.

Hot-Swap

Hot-Swap bezieht sich auf eine Funktion der GKE JobSet API, bei der ein Job mit höherer Priorität schnell Ressourcen von einem Job mit niedrigerer Priorität übernehmen kann. So werden Ausfallzeiten minimiert und eine schnellere Wiederherstellung ermöglicht.

Wenn ein JobSet erstellt wird, plant GKE die Arbeitslast auf mehreren Slices, wie in der JobSet-Konfiguration angegeben. Wenn auf einem oder mehreren Slices ein Hardwarefehler auftritt, werden die betroffenen Pods als fehlgeschlagen markiert. Wenn Sie dieses Jobset neu planen und sich dafür entschieden haben, einen Ersatz-Slice in Ihrem GKE-Cluster zu behalten, der für einen Job mit niedrigerer Priorität verwendet werden könnte, ordnet das Jobset-System die Arbeitslast des fehlgeschlagenen Slice des Jobs mit höherer Priorität dem Ersatz-Slice zu, der vom Job mit niedrigerer Priorität im selben GKE-Cluster verwendet wird. Das Neuzuordnen dauert in der Regel weniger als eine Minute.

Beim Neustart eines JobSet kann es in den folgenden Situationen zu einem Hot-Swap kommen:

Standardmodus: Wenn im selben Cluster freie, inaktive TPU-Slices verfügbar sind, priorisiert der Kubernetes-Scheduler die Planung der neu gestarteten Jobs auf diesen Slices, anstatt darauf zu warten, dass die fehlgeschlagenen Slices repariert werden. Dadurch ist eine schnellere Wiederherstellung möglich.
Heterogene Arbeitslasten: In Clustern, in denen mehrere Arbeitslasten mit einer konfigurierten Kubernetes-PriorityClass ausgeführt werden, kann ein Neustart eines JobSets einen Hot-Swap auslösen. Wenn die Affinität des neu gestarteten Jobs mit den Ressourcen eines Jobs mit niedrigerer Priorität übereinstimmt, wird der Job mit niedrigerer Priorität von Kubernetes unterbrochen, sodass der Job mit höherer Priorität sofort gestartet werden kann. Sie können beispielsweise Ihre Pathways-Worker-Pods mit unterschiedlichen Prioritäten konfigurieren, indem Sie PriorityClass verwenden.

Wenn Sie Prioritäten in Ihrem Cluster verwenden möchten, definieren Sie eine Prioritätsklasse, z. B.:

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

Wenden Sie diese YAML-Datei auf Ihren GKE-Cluster an:

kubectl apply -f high-prior-job.yaml

Hängen Sie als Nächstes die neue Prioritätsklasse an den Pathways-Worker-Job an, indem Sie den folgenden Text zum Podspec Ihres pathways-worker-Pods hinzufügen.

priorityClassName: high-prior-job