Google verwendet KI-Technologie, um Inhalte in Ihre bevorzugte Sprache zu übersetzen. KI-Übersetzungen können Fehler enthalten.

Resilienztraining mit Pathways

Pathways bietet folgende Vorteile in Bezug auf die Ausfallsicherheit:

Anhalten/Fortsetzen: Ermöglicht es Ihrer Arbeitslast, geplante Unterbrechungen wie Benachrichtigungen über vorzeitiges Beenden zu tolerieren, ohne dass benutzerdefinierter Code für die Verarbeitung von vorzeitigem Beenden erforderlich ist.
Elastisches Training: Mit der Wiederherstellungslogik in Ihrem Trainingscode kann Ihre Arbeitslast ungeplante Hardwarefehler tolerieren, ohne dass der Client abstürzt.

Hinweis

Führen Sie die folgenden Aufgaben aus, bevor Sie beginnen:

Anhalten/Fortsetzen

Die Google Kubernetes Engine (GKE) sendet in der Regel eine Benachrichtigung über die vorzeitige Beendigung an einen Beschleuniger-Pod, bevor der Pod vorzeitig beendet wird. Die Pathways-Toleranz für vorzeitiges Beenden ist standardmäßig für alle Cloud Deploy-Bereitstellungen aktiviert und Pathways-Beschleunigerjobs warten auf diese Benachrichtigungen.

Wenn eine Benachrichtigung über eine vorzeitige Beendigung eingeht, prüft Pathways zuerst, ob die Arbeitslast transparent gespeichert und wiederhergestellt werden kann. In diesem Fall wird versucht, Ihre ML-Arbeitslast transparent zu sperren, indem der aktuelle Status in einen nichtflüchtigen Speicher wie Cloud Storage geschrieben wird, bevor GKE Ihre Beschleunigerjobs entfernt. Wenn GKE Ihre Jobs später neu plant, setzt Pathways Ihre ML-Arbeitslast fort, indem der persistente Status zurückgelesen wird.

Wenn die Arbeitslast nicht wiederhergestellt werden kann, beendet Pathways den Beschleunigerjob und leitet den Fehler an Ihren Job weiter, wenn elastisches Training konfiguriert ist. Wenn das elastische Training nicht konfiguriert ist, startet die GKE die gesamte Arbeitslast gemäß der JobSet-Neustartrichtlinie 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 colocated Python API von JAX definiert wurden, basieren auf zustandsorientierten Pathways-Komponenten. Diese Arbeitslasten können nicht wiederhergestellt werden.

Elastisches Training

Mit dem elastischen Training kann Ihr Trainingsjob bei Hardwarefehlern oder vorzeitigen Beendigungen von Slices fortgesetzt werden, ohne dass die gesamte Arbeitslast neu gestartet werden muss. In einem Trainingsjob mit mehreren Slices hält Pathways den Clientprozess mit einem einzelnen Controller und die fehlerfreien TPU-Slices aktiv, wenn ein einzelner Slice nicht mehr verfügbar ist.

Das elastische Training kombiniert die Funktionen des Pathways-Systems mit der Wiederherstellungslogik in Ihrem Trainingscode:

Fehlererkennung: Wenn auf einem TPU-Slice ein Hardwarefehler oder ein vorzeitiges Beenden auftritt, hebt Pathways die Zuweisung des betroffenen Slice auf und löst einen jax.errors.JaxRuntimeError aus, wenn Ihre Arbeitslast das nächste Mal auf Daten auf diesem Slice zugreift. Da der Clientprozess aktiv bleibt, kann Ihr Code diesen Fehler abfangen, Ihre Ressourcen neu konfigurieren und das Training fortsetzen.
Wiederherstellungsmodi: Die Klasse pathwaysutils.elastic.manager.Manager und der zugehörige Decorator elastic_retry unterstützen zwei Wiederherstellungsstrategien:
- Pausieren und fortsetzen: Wenn minimum_slice_count gleich None oder gleich der Gesamtzahl der Slices im Job ist, wird das Training pausiert, wenn ein Slice verloren geht. Es wird erst fortgesetzt, wenn alle Slices wieder fehlerfrei sind. Da sich die Anzahl der Slices bei den Versuchen nicht ändert, kann Ihre Arbeitslast das vorhandene Gerätenetz und die kompilierten ausführbaren Dateien wiederverwenden.
- Größenänderung von Replikaten: Wenn minimum_slice_count kleiner als die Gesamtzahl der Slices ist, wird das Training sofort auf den verbleibenden aktiven Slices fortgesetzt. Dazu wird das Gerätenetz neu erstellt und für die kleinere Topologie neu kompiliert. Ein Hintergrund-Monitoring-Thread prüft regelmäßig inaktive Slices und zeichnet neu fehlerfreie Slices in elastic_manager.available_inactive_slices auf. Ihre Trainingsschleife kann dann das Geräte-Mesh-Netzwerk mit den Ersatz-Slices neu erstellen, für die erweiterte Topologie neu initialisieren und den aktuellen Status im neuen Mesh-Netzwerk platzieren.
Zustandswiederherstellung: Um das Training nach einem elastischen Ereignis fortzusetzen, stellt Ihre Arbeitslast Modellparameter, Optimierungszustand und Dateniteratorzustand mithilfe von In-Memory-Snapshots wieder her. Die Klasse Snapshotter in orbax.checkpoint kopiert den replizierten Status asynchron vom TPU-HBM in den CPU-Arbeitsspeicher von pinned_host auf jedem Slice. Während der Wiederherstellung rekonstruiert Snapshotter den Status aus den verbleibenden Hostreplikaten und verteilt ihn direkt auf das aktive TPU-Mesh, ohne aus dem Remotespeicher zu lesen.

Elastizität auf dem Pathways-Proxyserver aktivieren

Wenn Sie elastische Gerätesätze für Pathways in GKE aktivieren möchten, fügen Sie dem Feld args des Containers pathways-proxy in Ihrem JobSet-Manifest das Flag --num_elastic_instances hinzu und geben Sie die Anzahl der elastischen TPU-Slices in Ihrer Arbeitslast an:

- name: pathways-proxy
  image: us-docker.pkg.dev/cloud-tpu-v2-images/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

Ersetzen Sie Folgendes:

BUCKET_NAME: der Cloud Storage-Bucket, der für temporäre Pathways-Dateien verwendet wird.
NUM_SLICES: die Anzahl der TPU-Slices in Ihrer Arbeitslast.

Aktive Geräte filtern und das Gerätenetzwerk erstellen

In einer Pathways-Arbeitslast gibt jax.devices() alle Geräte in der ursprünglichen Zuweisung für die Lebensdauer des Clientprozesses zurück. Wenn Sie elastisches Training verwenden, verwenden Sie elastic_manager.active_devices, damit Ihr Gerätenetz und Ihre Batchgröße nur fehlerfreie, aktive Slices verwenden:

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"))

Wenn sich die Anzahl der aktiven Slices im Modus zur Größenänderung von Replikaten ändert, verwenden Sie einen elastischen Daten-Loader wie grain.experimental.ElasticIterator. Standard-Dateniteratoren pro Host speichern den Status pro Host und können nicht auf einer anderen Anzahl von Hosts wiederhergestellt werden. ElasticIterator verfolgt dagegen eine einzelne globale Dataset-Position, die auf einer beliebigen Anzahl aktiver Slices wiederhergestellt werden kann.

Elastische Trainingsschleife

Eine elastische Trainingsschleife koordiniert das Slice-Monitoring, die Fehlerbehebung und das Erstellen von In-Memory-Snapshots, wie im folgenden Codebeispiel gezeigt:

Ihre train-Funktion erstellt eine einzelne manager.Manager()-Instanz und einen Orbax-Snapshotter außerhalb des Wiederholungs-Wrappers, damit das Slice-Tracking und die Host-Speicher-Snapshots über die Versuche hinweg erhalten bleiben.
Wenn Sie die innere run_training_loop-Funktion mit @elastic_manager.elastic_retry dekorieren, wird auf fehlerfreie Slices gewartet, jax.default_device auf ein aktives Slice gesetzt, der Hintergrund-Monitor-Thread über die Schritte hinweg aktiv gehalten und die Schleife wiederholt, wenn ein jax.errors.JaxRuntimeError auftritt.
Wenn Sie sich von einem Slice-Fehler zu Beginn eines Versuchs erholen, muss Ihr Code das Gerätenetzwerk von elastic_manager.active_devices neu erstellen und snapshotter.load aufrufen, um den neuesten Snapshot auf den aktualisierten Netzwerk-Shards wiederherzustellen. Dabei werden Gerätepuffer aus dem fehlgeschlagenen Versuch aus dem Bereich entfernt.
In jedem Schritt der Schleife ruft Ihr Code regelmäßig snapshotter.save auf, um den Trainingsstatus zu sichern, und prüft elastic_manager.available_inactive_slices, um fehlerfreie Ersatz-Slices zu erkennen. Wenn Ersatz-Slices verfügbar sind, werden sie von Ihrem Code zu elastic_manager.active_slice_indices hinzugefügt, das Mesh-Netzwerk wird neu erstellt, die Trainingsschleife wird neu initialisiert und der Status wird auf das erweiterte Mesh-Netzwerk übertragen.

Das folgende Beispiel zeigt, wie Sie eine elastische Trainingsschleife implementieren, indem Sie manager.Manager, elastic_retry und den speicherinternen Snapshotter von Orbax kombinieren:

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

Elastic Manager und Wiederholungsverhalten konfigurieren

Sie instanziieren manager.Manager(devices=None) einmal pro Prozess. Wenn Sie devices weglassen, verwendet der Manager jax.devices() und gruppiert Geräte nach ihrem slice_index-Attribut.

Die Manager-Instanz stellt die folgenden Eigenschaften bereit, um die Clustertopologie während des Trainings zu prüfen:

slice_to_devices: Eine Zuordnung von jedem slice_index zur Sequenz von jax.Device-Objekten, die zu diesem Segment gehören.
all_slice_indices: Ein frozenset aller Slice-Indizes in der Zuweisung.
active_slice_indices: Ein frozenset von Slice-Indizes, die am aktiven Berechnungs-Mesh beteiligt sind.
active_devices: Eine Liste von jax.Device-Objekten, die zu den aktiven Slices gehören, sortiert nach slice_index.
inactive_slice_indices: Eine frozenset von Slice-Indizes in der Zuweisung, die nicht am Berechnungs-Mesh beteiligt sind.
available_inactive_slices: Ein frozenset inaktiver Slice-Indizes, die der Hintergrund-Monitoring-Thread als fehlerfrei und bereit für die Wiederaufnahme erkannt hat.
total_slice_count und active_slice_count: Die Gesamtzahl der zugewiesenen Slices und die Anzahl der aktiven Slices.
default_device: Gerät 0 aus einem aktiven Slice, das elastic_retry automatisch als Standard-JAX-Gerät für jeden Versuch festlegt.

Sie können das Wiederholungs- und Skalierungsverhalten anpassen, indem Sie die folgenden Argumente an @elastic_manager.elastic_retry übergeben:

minimum_slice_count: Die Mindestanzahl fehlerfreier Slices, die erforderlich sind, bevor die erste Ausführung oder ein Wiederholungsversuch gestartet wird. Wenn der Wert auf None (Standard) gesetzt ist, wartet elastic_retry auf alle Slices in der Zuweisung (total_slice_count) und arbeitet im Pausen- und Fortsetzungsmodus. Wenn Sie eine kleinere Ganzzahl festlegen (z. B. 1), wird elastic_retry im Modus zur Größenanpassung von Replikaten ausgeführt.
poll_interval: Das Intervall in Sekunden zwischen den Slice-Systemdiagnosen, die vom Elastic Training Manager (Manager) während des Trainings durchgeführt werden. Die Standardeinstellung ist 10.
timeout: Die maximale Anzahl von Sekunden, die gewartet werden soll, bis minimum_slice_count Slices aktiv werden, bevor ein TimeoutError ausgelöst wird. Die Standardeinstellung ist None (kein Zeitlimit).
pre_callback: Ein optionaler Callback, der zu Beginn jedes Versuchs ausgeführt wird, nachdem active_slice_indices aktualisiert wurde und bevor die umschlossene Trainingsschleifenfunktion ausgeführt wird.
on_elastic_event_callback: Ein optionaler Callback, der nach dem Abfangen eines Slice-Down-jax.errors.JaxRuntimeError ausgeführt wird, bevor der nächste Wiederholungsversuch beginnt. Mit diesem Callback können Sie unvollständige Prüfpunkte bereinigen oder Dateniteratoren zurücksetzen.
retry_policy: Ein aufrufbares (attempt: int, error: Exception) -> bool, das bestimmt, ob nach einem elastischen Ereignis weitere Wiederholungsversuche unternommen werden sollen. Verwenden Sie manager.ElasticRetryLimit(max_attempts), um die Gesamtzahl der Versuche zu begrenzen. Wenn keine Angabe gemacht wird, wird elastic_retry unbegrenzt wiederholt.

Bei einem elastischen Ereignis gibt JAX die Gerätepuffer des fehlgeschlagenen Versuchs frei, sobald der Stapelframe aufgelöst wird und keine Variablen mehr auf diese Arrays verweisen. Vermeiden Sie es, Verweise auf Geräte-Arrays aus einem fehlgeschlagenen Versuch in globalen Variablen oder langlebigen Callbacks zu speichern, damit der Arbeitsspeicher vor dem nächsten Versuch freigegeben wird.

In-Memory-Snapshots mit Orbax

Die Klasse ocp.v1.training.pathways.Snapshotter (verfügbar ab orbax-checkpoint-Version 0.12.6) verwaltet eine In-Memory-Sicherung Ihres Trainingsstatus im pinned_host-Speicher über Slices hinweg:

Snapshotter(replica_axis_index=0): Initialisiert den Snapshotter und gibt an, welche Mesh-Achse datenparallele Slice-Replikate darstellt. Alle JAX-Arrays im gespeicherten Zustand müssen entlang dieser Mesh-Achse repliziert werden, damit die verbleibenden Slices eine vollständige Kopie des Modell- und Optimierungszustands behalten, wenn ein Slice ausfällt. Wenn Ihr Status Arrays mit einem einzelnen Gerät enthält, replizieren Sie sie im gesamten Mesh-Netzwerk, bevor Sie save aufrufen.
save(step, state): Überträgt state asynchron in den pinned_host-Speicher und teilt es entlang replica_axis_index in Bäume pro Replikat auf. Wenn eine vorherige Snapshot-Übertragung noch läuft, überspringt save die neue Anfrage, um ein Blockieren der Trainingsschleife zu vermeiden.
load(abstract_state, reset_snapshot_state=True): Prüft die gespeicherten Host-Replikate, um überlebende Slices zu identifizieren, verkettet die fehlerfreien Replikate und verteilt den Zustand neu auf die durch abstract_state angegebenen Ziel-Shards.
join(): Blockiert, bis alle Snapshot-Vorgänge in der Warteschlange abgeschlossen sind.
latest: Gibt Metadaten (einschließlich step) für den zuletzt abgeschlossenen Snapshot zurück oder None, wenn kein Snapshot abgeschlossen wurde.

Sharding bei Slice-Fehler reduzieren

Wenn ein TPU-Slice während der Ausführung fehlschlägt oder vorzeitig beendet wird, löst Pathways eine jax.errors.JaxRuntimeError aus (z. B. DATA_LOSS, DEADLINE_EXCEEDED, NOT_FOUND oder INTERNAL). Der elastic_retry-Decorator fängt diese Ausnahme ab und prüft elastic.is_error_due_to_slice_down(error):

Wenn der Fehler nicht durch ein Slice-Down-Ereignis verursacht wird, löst elastic_retry die Ausnahme sofort noch einmal aus.
Wenn der Fehler durch ein Slice-Down-Ereignis verursacht wird, beendet elastic_retry alle aktiven JAX-Profiler-Traces, führt on_elastic_event_callback aus und wertet retry_policy(attempt, error) aus.
Vor dem nächsten Versuch ruft elastic_retry einmal elastic.wait_for_slices auf. Dadurch wird der Zustand der Slices abgefragt, bis mindestens minimum_slice_count Slices aktiv sind. Außerdem wird elastic_manager.active_slice_indices aktualisiert und die Trainingsschleifenfunktion wird auf den verbleibenden Slices ausgeführt.

Sharding erhöhen, wenn Slices wieder zusammengeführt werden

Wenn Ihre Arbeitslast im Modus zur Größenanpassung von Replikaten auf einer reduzierten Anzahl von Slices ausgeführt wird, führt der Hintergrund-Monitoring-Thread regelmäßig Systemdiagnosen für elastic_manager.inactive_slice_indices aus. Wenn reparierte oder neu geplante Slices die Systemdiagnose bestehen, fügt der Monitor-Thread ihre Slice-Indizes zu elastic_manager.available_inactive_slices hinzu.

Wenn elastic_manager.available_inactive_slices am Ende eines Trainingsschritts nicht leer ist, wird die Arbeitslast direkt wieder hochskaliert, ohne dass der Trainingsfortschritt verloren geht:

Wenn Sie die Slice-Indizes in elastic_manager.available_inactive_slices zu elastic_manager.active_slice_indices hinzufügen, werden die wieder zusammengefügten Slices in elastic_manager.active_devices aufgenommen und aus available_inactive_slices gelöscht.
Durch das Neuerstellen des Gerätenetzes aus elastic_manager.active_devices und das erneute Initialisieren des Trainingsschritts und des Daten-Iterators wird die Berechnung an die erweiterte Topologie angepasst.
Durch den Aufruf von jax.device_put wird der aktuelle Trainingsstatus auf die neuen Mesh-Shards übertragen, ohne dass eine Wiederherstellung aus einem Snapshot erforderlich ist.

Hot-Swap

Hot-Swap ist eine Funktion der GKE JobSet API, mit der ein Job mit höherer Priorität schnell Ressourcen von einem Job mit niedrigerer Priorität übernehmen kann. Dadurch 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 fehlerhaft markiert. Wenn Ihr GKE-Cluster einen freien Slice hat, auf dem ein Job mit niedrigerer Priorität ausgeführt wird, ordnet das JobSet-System die Arbeitslast des fehlgeschlagenen Slice vom Job mit höherer Priorität diesem freien Slice zu. Diese Neuzuordnung dauert in der Regel weniger als eine Minute.

Wenn ein JobSet neu gestartet wird, kann in den folgenden Situationen ein Hot-Swap auftreten:

Standardmodus: Wenn im selben Cluster freie, inaktive TPU-Slices verfügbar sind, priorisiert der Kubernetes-Planer die Planung der neu gestarteten Jobs auf diesen Slices, anstatt darauf zu warten, dass die fehlerhaften Slices repariert werden. Dieses Planungsverhalten ermöglicht eine schnellere Wiederherstellung.
Heterogene Arbeitslasten: In Clustern, in denen mehrere Arbeitslasten mit einer konfigurierten Kubernetes-PriorityClass-Ressource ausgeführt werden, kann ein neu gestartetes JobSet 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, damit der Job mit höherer Priorität gestartet werden kann, ohne auf die Reparatur des Slices warten zu müssen. Sie können beispielsweise Ihre Pathways-Worker-Pods mit einer PriorityClass-Ressource mit unterschiedlichen Prioritäten konfigurieren.

Wenn Sie Prioritäten in Ihrem Cluster verwenden möchten, definieren Sie ein PriorityClass-Manifest:

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 Ihren Pathways-Worker-Job an, indem Sie der Pod-Spezifikation Ihres pathways-worker-Pods den folgenden Text hinzufügen.

priorityClassName: high-prior-job

Resilienztraining mit Pathways Mit Sammlungen den Überblick behalten Sie können Inhalte basierend auf Ihren Einstellungen speichern und kategorisieren.