Resilienztraining mit Pathways

Pathways bietet folgende Vorteile in Bezug auf die Resilienz:

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

    Hinweise

    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 vorzeitige Beendigung eingeht, prüft Pathways zuerst, ob die aktuelle Arbeitslast wiederhergestellt werden kann, d. h., ob Pathways die Arbeitslast transparent speichern und wiederherstellen kann. Wenn dies der Fall ist, wird versucht, Ihre ML-Arbeitslast transparent zu pausieren, indem der aktuelle Status in einem nichtflüchtigen Speicher wie Cloud Storage gespeichert wird, bevor GKE Ihre Accelerator-Jobs 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 mithilfe eines HBM-Snapshots (High Bandwidth Memory) wiederhergestellt werden können. Bestimmte ML-Arbeitslasten, z. B. solche, die mit der JAX-API für die gemeinsame Ausführung von Python-Code definiert werden, basieren auf zustandsbehafteten Pathways-Komponenten. Diese können nicht wiederhergestellt werden.

    Elastisches Training

    Mit dem elastischen 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 neu konfigurieren, um die Verarbeitung fortzusetzen oder ordnungsgemäß zu beenden.
    • 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:
      • Ersatz von Segmenten: 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.
    • Dienstprogramme in pathwaysutils: Eine Reihe von Pathways-Dienstprogrammen, 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:

    1. Elastischen Manager erstellen
    2. Kapseln Sie Ihre Trainingsschleife in try-except-Blöcken, die jax.errors.JaxRuntimeErrors verarbeiten.
    3. 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, oder löst ihn anderweitig neu aus.
    4. maybe_snapshot und maybe_reshard_up am Ende des Trainingsdurchlaufs aufrufen
    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 Neupartitionierung wiederholen soll. Der Manager ist ein Singleton-Objekt und sollte nur einmal erstellt werden.

    Snapshots

    Je nach Konfiguration des Elastic Managers 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. Wenn ja, wird der elastische 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

    Basierend auf der Konfiguration des Elastic Managers und wenn 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 nicht bereits ein Snapshot für den aktuellen Schritt vorhanden ist) und der Elastic-Handler wird in einer Schleife aufgerufen, bis der Vorgang erfolgreich ist oder die maximale Anzahl von Wiederholungsversuchen erreicht wurde. Wenn eine erneute Shardierung erfolgt, werden die Rückgabewerte des Elastic-Handlers an den Aufrufer übergeben. 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 für mehrere 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:

    1. 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.
    2. 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, beendet Kubernetes den Job mit niedrigerer Priorität vorzeitig, 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-Workerjob an, indem Sie den folgenden Text zum Podspec Ihres pathways-worker-Pods hinzufügen.

    priorityClassName: high-prior-job

    Nächste Schritte