Resilienztraining mit Pathways

Auf dieser Seite wird erläutert, wie Sie Ausfallsicherheitsfunktionen für Pathways-Arbeitslasten in Google Kubernetes Engine (GKE) konfigurieren, damit umfangreiche Trainingsjobs nach geplanten Unterbrechungen und ungeplanten Hardwarefehlern wiederhergestellt werden können.

Pathways bietet folgende Vorteile in Bezug auf die Ausfallsicherheit:

  • Anhalten und fortsetzen: Ermöglicht es Ihrer Arbeitslast, geplante Unterbrechungen wie Benachrichtigungen über vorzeitige Beendigungen zu tolerieren, ohne dass benutzerdefinierter Code für die Verarbeitung vorzeitiger Beendigungen erforderlich ist.
  • Elastisches Training: Ermöglicht es Ihrer Arbeitslast, ungeplante Hardwarefehler zu tolerieren, ohne dass der Client abstürzt, indem Wiederherstellungslogik in Ihrem Trainingscode verwendet wird.
  • Multi-Tier Checkpointing: Ermöglicht es Ihrer Arbeitslast, hochfrequente lokale und Peer-replizierte Prüfpunkte über TPU-Worker hinweg mit colocated Python zu speichern, wobei asynchrone Cloud Storage-Sicherungen die Wiederherstellungszeit verkürzen.
  • Hot-Swap: Ermöglicht es einer GKE-JobSet-Arbeitslast mit höherer Priorität, bei einem Hardwarefehler schnell freie TPU-Slices oder TPU-Slices mit niedrigerer Priorität zu übernehmen.

Hinweis

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

Wenn Sie elastisches Training verwenden, installieren Sie die folgenden Pakete in Ihrem Arbeitslast-Container-Image:

  • pathwaysutils
  • Orbax (orbax-checkpoint): >= 0.12.6 (erforderlich für In-Memory-Snapshots mit ocp.v1.training.pathways.Snapshotter)
  • Körnung (grain): erforderlich, wenn Sie grain.experimental.ElasticIterator verwenden

Wenn Sie mehrstufige Prüfpunkte verwenden, müssen Sie auch die folgenden Voraussetzungen erfüllen:

  • Konfigurieren Sie Ihren GKE-Cluster für mehrstufiges Checkpointing, indem Sie der Anleitung unter Große Modelle für maschinelles Lernen in GKE mit mehrstufigem Checkpointing trainieren folgen:
    • Aktivieren Sie die Add-ons für mehrstufiges Checkpointing und Cloud Storage FUSE (GcsFuseCsiDriver) sowie die Workload Identity-Föderation für GKE in Ihrem Cluster.
    • Erstellen Sie einen Cloud Storage-Bucket mit aktiviertem hierarchischen Namespace.
    • Weisen Sie dem Kubernetes-Dienstkonto für Multi-Tier Checkpointing die IAM-Rollen (Identity and Access Management) roles/storage.objectUser und roles/storage.insightsCollectorService zu.
    • Stellen Sie eine benutzerdefinierte CheckpointConfiguration-Ressource bereit, in der Sie Ihren Cloud Storage-Bucket, die TPU-Knotenauswahl und die Größe des In-Memory-Volumes (inMemoryVolumeSize) angeben.
  • Erstellen Sie ein Image für einen am selben Ort befindlichen Python-Sidecar-Container und übertragen Sie es per Push, z. B. mit dem Beispiel-Sidecar-Dockerfile. Verwenden Sie in Ihren Arbeitslast- und Sidecar-Images übereinstimmende jax- und jaxlib-Versionen, die mit Ihrem Pathways-Release übereinstimmen und die folgenden Mindestversionsanforderungen erfüllen:
    • Orbax-Version: >= 0.12.4
    • JAX-Version (einschließlich colocated Python): >= 0.10.1

Anhalten/Fortsetzen

GKE sendet in der Regel eine Benachrichtigung über die vorzeitige Beendigung an einen Beschleuniger-Pod, bevor der Pod vorzeitig beendet wird. Pathways aktiviert standardmäßig die Toleranz für Preemption bei allen Cloud-Bereitstellungen. 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. Wenn möglich, versucht Pathways, 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 die 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 Ihre Arbeitslast weiter, wenn elastisches Training konfiguriert ist. Wenn das elastische Training nicht konfiguriert ist, startet 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 und 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 Sie minimum_slice_count auf None oder auf die Gesamtzahl der Slices festlegen, wird das Training pausiert, wenn ein Slice ausfällt, und fortgesetzt, wenn alle Slices wieder fehlerfrei sind. Da sich die Anzahl der Slices bei den Versuchen nicht ändert, kann Ihre Arbeitslast das vorhandene Gerätemesh und die kompilierten ausführbaren Dateien wiederverwenden.
    • Größenänderung von Replikaten: Wenn Sie minimum_slice_count auf einen Wert festlegen, der kleiner als die Gesamtzahl der Slices ist, wird das Training auf den verbleibenden aktiven Slices fortgesetzt, wenn ein Slice ausfällt. 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ätenetz mit den Ersatz-Slices neu erstellen, für die erweiterte Topologie neu initialisieren und den aktuellen Status auf dem neuen Netz 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/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

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. Im Gegensatz dazu verfolgt ElasticIterator eine einzelne globale Dataset-Position, die Sie über eine beliebige Anzahl aktiver Slices wiederherstellen können.

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 @elastic_manager.elastic_retry auf run_training_loop anwenden, wird die Funktion so konfiguriert, dass sie auf fehlerfreie Slices wartet, jax.default_device auf ein aktives Slice setzt, den Hintergrund-Monitoring-Thread aktiv hält und bei einem jax.errors.JaxRuntimeError einen neuen Versuch startet.
  • Zu Beginn eines Wiederherstellungsversuchs muss Ihr Code das Geräte-Mesh aus elastic_manager.active_devices neu erstellen und snapshotter.load aufrufen, um den neuesten Snapshot auf den aktualisierten Mesh-Shards wiederherzustellen. Lassen Sie die Gerätepuffer des fehlgeschlagenen Versuchs aus dem Bereich fallen, damit JAX den Speicher freigeben kann.
  • 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 der Manager active_slice_indices aktualisiert hat und bevor die umschlossene Trainingsschleifenfunktion ausgeführt wird.
  • on_elastic_event_callback: Ein optionaler Callback, der ausgeführt wird, nachdem elastic_retry einen Slice-Down-jax.errors.JaxRuntimeError abfängt, 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). In einer Funktion, die mit @elastic_manager.elastic_retry versehen ist, fängt elastic_retry 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, um den Zustand des Slices abzufragen, bis mindestens minimum_slice_count Slices aktiv sind. Anschließend wird elastic_manager.active_slice_indices aktualisiert und die Trainingsschleifenfunktion 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.

Mehrstufiges Checkpointing

Bei umfangreichen Trainingsjobs kommt es häufig zu Unterbrechungen. Das Schreiben oder Wiederherstellen von Prüfpunkten direkt in und aus Cloud Storage kann den Trainingsfortschritt beeinträchtigen. Mehrstufiges Checkpointing reduziert den Aufwand für das Checkpointing und die Wiederherstellungszeit, indem drei Speicherebenen in GKE koordiniert werden:

  • Lokale RAM-Disk: TPU-Worker schreiben hochfrequente Prüfpunkte in ein CSI-Volume im Arbeitsspeicher (multitier-checkpoint.csi.storage.gke.io), das auf jedem Worker-Knoten bereitgestellt ist. Dieses Volume bleibt bei Pod-Neustarts auf demselben Knoten erhalten.
  • Clusterinterne Peer-Replikation: Der GKE-Replikationsdienst repliziert lokale Prüfpunkt-Shards auf Peer-TPU-Worker im Cluster.
  • Cloud Storage-Sicherung: Der Replikationsdienst sichert lokale Prüfpunkte asynchron in einem Cloud Storage-Bucket mit aktiviertem hierarchischem Namespace, ohne TPU-Trainingsschritte zu blockieren.

Im Gegensatz zu JAX mit mehreren Controllern, bei dem jeder TPU-Host ein eigenes Trainingsskript ausführt, wird bei Pathways ein einzelner Controller auf dem pathways-head-Pod ausgeführt und die Arbeit an pathways-worker-Pods verteilt. Um Prüfpunkt-Shards direkt auf der lokalen RAM-Disk jedes TPU-Hosts zu lesen und zu schreiben, ohne Prüfpunktdaten über den einzelnen Controller weiterzuleiten, verwendet Pathways colocated Python. In jedem worker-Pod wird ein colocated-python-sidecar-Container neben dem pathways-worker-Container ausgeführt. Beide Container stellen das CSI-Volume für mehrstufige Prüfpunkte und ein Verzeichnis für freigegebenen Arbeitsspeicher (/tmp/sidecar) bereit. Während der Initialisierung und des Trainings sendet Orbax die Einrichtung, Speicherung und Wiederherstellung von Prüfpunkten vom Pathways-Controller an die Worker-Sidecars auf jedem TPU-Host. Multi-Tier Checkpointing auf Pathways unterstützt auch elastisches Training, wenn die Anzahl der aktiven TPU-Slices herunterskaliert oder hochskaliert wird.

Pathways unterstützt mehrstufige Prüfpunkte mit MaxText. Eine Anleitung zur Einrichtung und Konfiguration finden Sie im Leitfaden für mehrstufige Prüfpunkte in MaxText.

Pathways-Arbeitslast konfigurieren

Sie können eine Pathways-Arbeitslast für mehrstufiges Checkpointing entweder mit dem Cluster Toolkit (gcluster job submit) oder durch Konfigurieren eines benutzerdefinierten Pathways-JobSet-Manifests konfigurieren.

Cluster Toolkit verwenden

Wenn Sie einen Job mit gcluster job submit --pathways senden, übergeben Sie die Flags --gke-mtc-enabled, --gke-mtc-ramdisk-dir und --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" \
    ...

Ersetzen Sie SIDECAR_IMAGE durch den URI Ihres am selben Ort befindlichen Python-Sidecar-Container-Images.

Wenn Sie diese Flags festlegen, konfiguriert gcluster automatisch Folgendes:

  • Übergibt --sidecar_name=external an den pathways-proxy-Container.
  • Führt den Container colocated-python-sidecar in jedem worker-Pod aus und stellt ein freigegebenes /tmp/sidecar-Verzeichnis (--cloud_pathways_sidecar_shm_directory=/tmp/sidecar) sowohl auf pathways-worker als auch auf colocated-python-sidecar bereit.
  • Stellt das CSI-Volume multitier-checkpoint.csi.storage.gke.io im Pfad bereit, der durch --gke-mtc-ramdisk-dir in pathways-worker und colocated-python-sidecar auf den worker-Pods angegeben ist, aber nicht auf dem pathways-head-Pod.

Benutzerdefiniertes Pathways-JobSet-Manifest verwenden

Wenn Sie Ihr eigenes Pathways-JobSet-Manifest verwalten, anstatt gcluster zu verwenden, übergeben Sie --sidecar_name=external an pathways-proxy. Stellen Sie in der Pod-Spezifikation für worker sowohl das freigegebene Sidecar-Speichervolume als auch das CSI-Volume für multitier-checkpoint.csi.storage.gke.io in den Containern pathways-worker und colocated-python-sidecar bereit:

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

Orbax-Multi-Tier-Prüfpunkte in einer benutzerdefinierten JAX-Arbeitslast verwenden

Verwenden Sie in einem benutzerdefinierten JAX-Trainingsskript, das auf Pathways ausgeführt wird, das Paket orbax.checkpoint.experimental.emergency.multi_tier_checkpointing und legen Sie use_colocated_python=True sowohl in initialize_multi_tier_checkpointing als auch in ReplicatorCheckpointManagerOptions fest:

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

Hot-Swap

Hot-Swap ist ein Feature der GKE JobSet API, mit dem ein Job mit höherer Priorität schnell Ressourcen von einem Job mit niedrigerer Priorität übernehmen kann. Dadurch werden Ausfallzeiten minimiert und die Wiederherstellung beschleunigt.

Wenn Sie ein JobSet erstellen, plant GKE die Arbeitslast auf mehreren Slices, wie in der JobSet-Konfiguration angegeben. Wenn auf einem oder mehreren Slices ein Hardwarefehler auftritt, markiert GKE die betroffenen Pods als fehlgeschlagen. Wenn Ihr GKE-Cluster einen Ersatz-Slice hat, auf dem ein Job mit niedrigerer Priorität ausgeführt wird, ordnet der JobSet-Controller die Arbeitslast des fehlgeschlagenen Slice vom Job mit höherer Priorität diesem Ersatz-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 und speichern Sie es als 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."

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

kubectl apply -f high-prior-job.yaml

Als Nächstes hängen Sie die neue Prioritätsklasse an Ihre Pathways-Arbeitslast an. Wenn Sie Ihre Arbeitslast mit dem Cluster Toolkit einreichen, übergeben Sie --priority=high-prior-job an den Befehl gcluster job submit, wodurch priorityClassName sowohl für die pathways-head- als auch für die worker-Pods festgelegt wird. Wenn Sie Ihr eigenes JobSet-Manifest schreiben, legen Sie priorityClassName in der Pod-Vorlage des replizierten Jobs worker fest, der den Container pathways-worker ausführt:

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

Wenn Sie Ihren Cluster mit einem Cluster Toolkit-TPU-Blueprint mit aktiviertem Pathways erstellt haben, enthält der Cluster bereits die Prioritätsklassen very-low, low, medium, high und very-high. Sie können eine davon an --priority übergeben, anstatt eine eigene zu erstellen.

Nächste Schritte