Umfangreiche Modelle für maschinelles Lernen in GKE mit Multi-Tier Checkpointing trainieren

Auf dieser Seite erfahren Sie, wie Sie mit Multi-Tier Checkpointing Prüfpunkte während des Trainings Ihres Machine-Learning-Modells in GKE zuverlässig speichern und verwalten können. Die Speicherung und Verwaltung von Prüfpunkten ist für umfangreiche Trainingsjobs, die Tausende von Knoten nutzen, von entscheidender Bedeutung. Unterbrechungen dieser umfangreichen Jobs sind häufig (potenziell stündlich) und die Wiederherstellung kann langsam sein.

Vorteile

Die Verwendung von mehrstufigen Prüfpunkten bietet folgende Vorteile:

  • Vollständig orchestrierte Verwaltung von Prüfpunktdaten, einschließlich Sicherungen, Replikation und automatischer Wiederherstellung für die folgenden Arbeitslasten:
  • Schnelle Wiederherstellung von Trainingsjobs von einem Prüfpunkt, der auf dem lokalen Knoten gespeichert ist. Sie können auch Checkpoints verwenden, die auf einem anderen Knoten im Trainingscluster gespeichert sind.
  • Schnelle Wiederherstellung von Trainingsjobs aus einem Prüfpunkt, der in einem Cloud Storage-Backup gespeichert ist, im schlimmsten Fall, wenn keine Prüfpunkte im Cluster vorhanden sind.

Hinweis

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

  • Aktivieren Sie die Google Kubernetes Engine API.
    • Google Kubernetes Engine API aktivieren
  • Wenn Sie die Google Cloud CLI für diese Aufgabe verwenden möchten, müssen Sie die gcloud CLI installieren und dann initialisieren. Wenn Sie die gcloud CLI bereits installiert haben, rufen Sie die neueste Version mit dem Befehl gcloud components update ab. In früheren gcloud CLI-Versionen werden die Befehle in diesem Dokument möglicherweise nicht unterstützt.

Voraussetzungen

Für mehrstufiges Checkpointing ist die GKE-Clusterversion 1.32.4-gke.1415000 oder höher erforderlich.

Beschränkungen

  • Autopilot-Cluster werden nicht unterstützt.

GKE-Knoten für die Verwendung von mehrstufigem Checkpointing konfigurieren

In diesem Abschnitt wird beschrieben, wie Sie GKE-Knoten in neuen und vorhandenen Clustern konfigurieren.

Knoten in einem neuen Cluster konfigurieren

  1. Erstellen Sie einen Cluster mit aktiviertem mehrstufigem Checkpointing, dem Cloud Storage FUSE-CSI-Treiber und der Workload Identity Federation for GKE. Wenn Sie TPU-Slices für Ihre ML-Arbeitslast verwenden, müssen Sie den Befehl zum Erstellen des Clusters anpassen, um die Konfiguration für einen TPU-Slice-Knotenpool einzuschließen.

    gcloud container clusters create CLUSTER_NAME \
    --addons=HighScaleCheckpointing,GcsFuseCsiDriver  \
    --node-locations=NODE_LOCATION \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --cluster-version=CLUSTER_VERSION
    --location=CLUSTER_LOCATION \
    --machine-type=MACHINE_TYPE \
    --num-nodes=NUM_NODES

    Ersetzen Sie die folgenden Werte:

    • CLUSTER_NAME ist der Name des Clusters.
    • NODE_LOCATION: die Zone für Ihre Clusterknoten. Hier wird Ihre TPU-Kapazität angezeigt.
    • PROJECT_ID: Ihre Projekt-ID in Google Cloud .
    • CLUSTER_VERSION: Die Version Ihres Clusters. 1.32.4-gke.1415000 ist die niedrigste unterstützte Version.
    • CLUSTER_LOCATION: die Region, in der Sie den Cluster erstellen möchten.
    • MACHINE_TYPE: Der Maschinentyp, der für Knoten verwendet wird, auf denen Komponenten wie der JobSet-Controller und der mehrstufige Prüfpunkt-Controller ausgeführt werden. Für das Training im großen Maßstab empfehlen wir, mindestens e2-standard-4 Maschinen zu verwenden. Sie verwenden diesen Maschinentyp nicht für das Modelltraining. Stattdessen erstellen Sie dafür separate Knotenpools, häufig mit beschleunigungsoptimierten VM-Familien.
    • NUM_NODES: Die Anzahl der Knoten, die in jeder der Clusterzonen erstellt werden sollen.

Knoten in einem vorhandenen Cluster konfigurieren

Wenn Sie Multi-Tier Checkpointing mit einem vorhandenen Cluster verwenden möchten, aktivieren Sie es zusammen mit dem CSI-Treiber für Cloud Storage FUSE und der Workload Identity-Föderation für GKE mit den folgenden Befehlen. Ihre vorhandene Clusterversion muss höher als 1.32.3-gke.1170000 sein.

  1. Aktivieren Sie die Workload Identity-Föderation für GKE:

    gcloud container clusters update CLUSTER_NAME \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --location=CLUSTER_LOCATION

    Ersetzen Sie die folgenden Werte:

    • CLUSTER_NAME ist der Name des Clusters.
    • PROJECT_ID: Ihre Google Cloud Projekt-ID.
    • CLUSTER_LOCATION: Die Region des Clusters.

  2. Aktivieren Sie das Multi-Tier Checkpointing und den CSI-Treiber für Cloud Storage FUSE:

    gcloud container clusters update CLUSTER_NAME \
    --update-addons=HighScaleCheckpointing=ENABLED,GcsFuseCsiDriver=ENABLED \
    --location=CLUSTER_LOCATION

Berechtigungen für die Verwendung von mehrstufigem Checkpointing konfigurieren

In diesem Abschnitt wird beschrieben, wie Sie Berechtigungen für die Verwendung von mehrstufigen Prüfpunkten konfigurieren.

Zugriff auf Cloud Storage-Buckets gewähren

Für die temporären Volumes, die vom Multi-Tier Checkpointing-CSI-Treiber verwendet werden, müssen vorhandene Cloud Storage-Buckets verwendet werden.

Damit Prüfpunkte in einem Cloud Storage-Bucket gespeichert werden können, benötigt Multi-Tier Checkpointing Zugriff auf den Bucket. Weisen Sie dem Kubernetes-Dienstkonto für Multi-Tier Checkpointing die IAM-Rolle „Storage Object User“ (roles/storage.objectUser) für den Bucket zu.

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member "principal://iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/gke-managed-checkpointing/sa/gke-checkpointing-multitier-node" \
    --role "roles/storage.objectUser"

Ersetzen Sie die folgenden Werte:

(Optional) Dem Compute Engine-Standarddienstkonto Zugriff gewähren

Wenn Ihre Compute Engine-Instanzen Lesezugriff auf den Cloud Storage-Bucket benötigen, weisen Sie dem Compute Engine-Standarddienstkonto die IAM-Rolle „Storage-Objekt-Betrachter“ (roles/storage.objectViewer) zu.

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --role roles/storage.objectViewer

JobSet-Controller bereitstellen

Der JobSet-Controller ist für die Verwaltung der Batchjobs verantwortlich, die Ihr Modelltraining in GKE ausführen. Seine Ressourcenzuweisung wird angepasst, um die Arbeitslast effizient zu verarbeiten. Achten Sie darauf, dass Ihr Launcher für Trainingsjobs JobSet bereitstellt und verwendet.

Führen Sie den folgenden Patch-Befehl aus, um die Speicheranfrage für den Manager-Container in Ihrer JobSet-Bereitstellung auf 1 Gi, das Speicherlimit auf 2 Gi und die CPU-Anfrage auf 1 zu erhöhen:

kubectl patch -n jobset-system deploy jobset-controller-manager --type json \
    --patch '[{"op": "add", "path": "/spec/template/spec/containers/0/resources", "value": {"limits": {"memory": "2Gi"}, "requests": {"cpu": "1", "memory": "1Gi"}}}]'

CSI-Treiber für Multi-Tier Checkpointing initialisieren

In diesem Abschnitt wird beschrieben, wie Sie den CSI-Treiber für mehrstufiges Checkpointing auf Knoten initialisieren, auf denen Ihre Arbeitslasten ausgeführt werden. Der CSI-Treiber ist für die Speicherung und Verwaltung von Prüfpunkten während des Modelltrainings verantwortlich.

CheckpointConfiguration erstellen

Eine CheckpointConfiguration ist eine benutzerdefinierte Kubernetes-Ressource, die Attribute für die Bereitstellung des CSI-Treibers für mehrstufige Prüfpunkte angibt. Diese Ressource ist clusterbezogen.

  1. Erstellen Sie das folgende checkpoint.yaml-Manifest:

    apiVersion: checkpointing.gke.io/v1
kind: CheckpointConfiguration
metadata:
  name: MTC_CONFIG_NAME-configuration
spec:
    cloudStorageBucketName: GCS_BUCKET
    nodeSelector:
        node.kubernetes.io/instance-type: MACHINE_TYPE
    tolerations:
    - key: TOLERATION_KEY
        operator: Exists
        effect: NoSchedule
    inMemoryVolumeSize: IN_MEMORY_VOLUME_SIZE
    gcsFuseMountOptions:
    - implicit-dirs
    - metadata-cache:negative-ttl-secs:0
    - metadata-cache:ttl-secs:-1
    - metadata-cache:stat-cache-max-size-mb:-1
    - metadata-cache:type-cache-max-size-mb:-1
    - file-system:kernel-list-cache-ttl-secs:0
    - read_ahead_kb=1024
    - write:enable-streaming-writes:true

    Ersetzen Sie Folgendes:

    • MTC_CONFIG_NAME: der Name Ihrer CheckpointConfiguration. Dieser Name ist global für den Cluster und nicht jobspezifisch.
    • GCS_BUCKET: der Name des Cloud Storage-Bucket, in dem Sie Prüfpunktdaten speichern. Verwenden Sie den Bucket, den Sie im Schritt Cloud Storage-Bucket mit Berechtigungen einrichten eingerichtet haben.

    • MACHINE_TYPE: der Maschinentyp für die entsprechenden Beschleuniger. Der Wert kann einer der folgenden sein:

      • TPU v5p: ct5p-hightpu-4t
      • TPU v5e: ct5e-hightpu-4t
      • TPU v6e: ct6e-standard-4t
      • Ironwood (TPU7x): tpu7x-standard-4t
      • NVIDIA H100 80 GB-GPUs (A3-Serie):

      Weitere Informationen zum Ausführen verteilter Arbeitslasten auf GPUs mit GKE finden Sie unter GPUs mit mehreren Instanzen ausführen. Informationen zu TPUs finden Sie unter TPU-Slice-Knotenpool erstellen.

    • TOLERATION_KEY: Mit diesem Feld kann der CSI-Treiber auf Knoten mit übereinstimmenden Markierungen geplant werden. Weitere Informationen zur Funktionsweise von Taints bei verschiedenen Beschleunigertypen finden Sie auf den folgenden Seiten:

    • IN_MEMORY_VOLUME_SIZE: die Größe des In-Memory-Checkpointing-Cache. Geben Sie die Menge und die Einheit an, z. B. „200 Gi“. Dieser Wert sollte:

      • Die Größe des lokalen Prüfpunkts für TPUs multipliziert mit 2,2
      • Die Größe des lokalen Prüfpunkts für GPUs mit einem einzelnen Peer multipliziert mit 6,6.

  2. Wenden Sie das Manifest an:

    kubectl apply -f checkpoint.yaml

  3. Prüfen Sie, ob der CSI-Treiber ausgeführt wird:

    kubectl get pod -n gke-managed-checkpointing

    Die Ausgabe sollte in etwa so aussehen: Es gibt mehrere Einträge, einen pro beschleunigtem Knoten.

    NAME                                                          READY   STATUS    RESTARTS   AGE
multitier-driver-e2b033a7-a4e7-496a-87a3-ffd7fcc2e57b-2d4fz   5/5     Running   0          114s

CSI-Treiber für Multi-Tier Checkpointing deinstallieren

Wenn Sie die Bereitstellung des CSI-Treibers für mehrstufiges Checkpointing aufheben möchten, löschen Sie die Ressource CheckpointConfiguration. Der Controller für Prüfpunkte auf mehreren Ebenen entfernt den CSI-Treiber von den Knoten. Dadurch werden die RAM-Laufwerke entfernt und Arbeitsspeicher für andere Arbeitslasten freigegeben. Beispiel:

kubectl delete -f checkpoint.yaml

Datenaufbewahrung und automatische Speicherbereinigung für Cloud Storage-Sicherungen verwalten

Sie sind für die Implementierung von Aufbewahrungsrichtlinien für die Cloud Storage-Sicherungen von Prüfpunkten verantwortlich. Bei der mehrstufigen Prüfpunkterstellung werden Prüfpunktsicherungen nur in Cloud Storage geschrieben und nie geändert oder gelöscht.

Viele Open-Source-Tools können die Aufbewahrung und die automatische Speicherbereinigung verarbeiten, darunter die folgenden:

Im folgenden Beispiel wird backup-warden verwendet, wobei das Verzeichnis backup an einem Sicherungsspeicherort bereitgestellt wird, der Cloud Storage FUSE verwendet:

# Add --delete option to actually delete the backups, as is it only shows what would be deleted (dry-run)
backup-warden -p backup \
    --hourly 24 \
    --daily 7 \
    --weekly 5 \
    --monthly always \
    --yearly always \
    --prefer-recent

JobSet-Manifest der Arbeitslast aktualisieren

Aktualisieren Sie das JobSet-Manifest für Ihren Job, um das Volume für Prüfpunkte im großen Maßstab einzuschließen. Die Details hängen von Ihrer Arbeitslast ab.

Wenn Sie beispielsweise das Beispiel-JobSet aus TPU-Multiplikationen in GKE bereitstellen erweitern möchten, gehen Sie so vor:

  1. Fügen Sie dem Container jax-tpu die folgenden Zeilen hinzu.

    volumeMounts:
- name: checkpoint
  mountPath: CHECKPOINT_DIR

    Ersetzen Sie CHECKPOINT_DIR durch den Pfad zu Ihrem Prüfpunktverzeichnis. An diesem Speicherort wird der replicator.yaml generiert und die mehrstufige Prüfpunkterstellung führt den Speichervorgang für den Prüfpunkt aus. Weitere Informationen finden Sie unter Mehrstufige Prüfpunkte in Ihre Anwendung einbinden.

  2. Fügen Sie dem Feld spec.template.spec der Job-Spezifikation die folgenden Zeilen hinzu.

    volumes:
- name: checkpoint
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

Multi-Tier-Checkpointing in Ihre Anwendung einbinden

Wenn Sie Informationen zu Prüfpunktstandorten und zur Replikationsbereitschaft freigeben möchten, ändern Sie Ihre Anwendung so, dass sie das folgende Protokoll für die Kommunikation mit Multi-Tier Checkpointing verwendet.

Die folgenden Frameworks und Bibliotheken implementieren diese Integration bereits:

  1. MaxText (mit Orbax): MaxText bietet integrierte Unterstützung für Multi-Tier Checkpointing (MTC). Weitere Informationen finden Sie in der MaxText-MTC-Dokumentation.
  2. Alle Modelle, die Orbax verwenden: Die Integration besteht aus zwei Hauptkomponenten: dem MTC-Initialisierungsmodul für die Ersteinrichtung (Beispiel siehe MaxText) und dem Replicator Checkpoint Manager zum Speichern und Wiederherstellen von Prüfpunkten von einer lokalen RAM-Disk (Beispiel siehe MaxText).
  3. PyTorch-Referenzbeispiel: Dieses Beispiel ist für Arbeitslasten verfügbar, die NeMo mit PyTorch verwenden.

Starten

In diesem Abschnitt werden die ersten Schritte beschrieben, die die Anwendung für die Interaktion mit der mehrstufigen Checkpointing-Funktion ausführen muss.

Der Replicator ist eine Kernkomponente von Multi-Tier Checkpointing, die auf jedem Knoten als Teil des CSI-Treibers ausgeführt wird. Der Replikator verwaltet die Prüfpunktreplikation über Speicherebenen hinweg, von der lokalen RAM-Disk zu Peer-Knoten und zu externem Speicher wie Cloud Storage.

Die Datei replicator.yaml fungiert als dynamische Steuerungsebene zwischen Ihrem ML-Trainingsjob (Framework-Code) und der Replikator-Komponente. Ihre ML-Anwendung generiert diese Datei programmatisch auf dem lokalen Volume (RAMDisk), auf das sowohl der Trainingsjob als auch der Replikationsdienst zugreifen können. Mit diesem Manifest kann das ML-Framework dem Replikator Laufzeitkonfigurations- und Lebenszyklusverwaltungsanweisungen bereitstellen, die sich von statischen Infrastrukturparametern (z. B. Cloud Storage-Uploadhäufigkeit) unterscheiden, die während der Backend-Einrichtung definiert werden.

Ein konkretes Beispiel für diese Interaktion finden Sie unter:

Ihre Anwendung sollte beim Start die folgenden Schritte ausführen:

  1. Warten Sie, bis die Datei replicator.yaml nicht mehr vorhanden ist. Das bedeutet, dass der Replikator von Ihrer Anwendung konfiguriert werden kann. Die Datei replicator.yaml wird am Speicherort CHECKPOINT_DIR generiert, den Sie im Abschnitt JobSet-Manifest der Arbeitslast aktualisieren konfiguriert haben.

    Wenn der Job zum Trainieren des Modells zum ersten Mal erstellt wird, ist die Datei replicator.yaml nicht vorhanden und Ihre Anwendung kann sofort fortfahren. Wenn der Job jedoch neu gestartet wurde (z. B. aufgrund eines Fehlers oder eines manuellen Eingriffs), verarbeitet das System möglicherweise noch die vorherige Jobinstanz und die replicator.yaml dieser Instanz ist möglicherweise noch auf dem lokalen Volume vorhanden.

  2. Ihre Anwendung oder Ihr ML-Job erstellt die Datei replicator.yaml mit einer Konfiguration, die in etwa so aussieht:

    Orbax

    job-name: orbax
framework: orbax
assume-data-parallelism: 3
node-rank: 0
nodes: 32
peer-ranks: [1, 16] or peers-per-node: 2
backup-interval-minutes: 30

    PyTorch

    job-name: nemo
framework: pytorch.distributed
node-rank: 0
nodes: 32
peer-ranks: [1, 16] or peers-per-node: 2
backup-interval-minutes: 30

    Diese Beispielkonfiguration enthält die folgenden Felder:

    • name: Der Name des Trainingsjobs.
    • framework: das ML-Framework, das vom Trainingsjob verwendet wird.
    • node-rank: die eindeutige Kennung des aktuellen Knotens innerhalb des verteilten Trainingsjobs. Dies ist der Knotenrang des Knotens, der diese Datei erstellt hat. Jeder Knoten, der am Durchlauf beteiligt ist, hat einen eigenen Rang.
    • nodes: Die Gesamtzahl der Knoten, die am verteilten Trainingsjob beteiligt sind. Dieser Wert stammt aus den Metadaten des Pods. Der ML-Trainingsjob kann diesen Wert ebenfalls sehen.
    • peer-ranks oder peers-per-node: zwei alternative Möglichkeiten, die Replikationstopologie anzugeben. Nur einer dieser beiden Parameter sollte vorhanden sein.
      • peer-ranks: Explizite Ränge von Peer-Knoten, auf die die Prüfpunktdaten des aktuellen Knotens repliziert werden sollen. So können Sie genau steuern, welche Knoten als Replikationspartner dienen.
      • peers-per-node: Die Anzahl der Peer-Knoten pro Knoten, die der Replicator automatisch für die Replikation auswählen soll.
    • backup-interval-minutes: die Häufigkeit in Minuten, mit der Prüfpunkte in Cloud Storage gesichert werden. Wir empfehlen, diesen Wert auf mindestens 30 Minuten festzulegen.

  3. Warten Sie, bis die neue replicator.yaml-Datei vom System gelöscht wurde. Das bedeutet, dass der Replikator neu gestartet wurde und die Bereinigung durchgeführt hat. Mit diesem Schritt können Sie vermeiden, dass auf dem lokalen Volume veraltete oder temporäre Dateien vorhanden sind, wenn Ihre Anwendung die Schritte im nächsten Abschnitt ausführt.

Wiederherstellung ab dem letzten bekannten funktionsfähigen Prüfpunkt

  1. Nachdem der Replikator initialisiert wurde, erstellt Multi-Tier Checkpointing einen symbolischen Link pro TPU- oder GPU-Worker. Diese symbolischen Links werden im selben eingebundenen lokalen Volume wie die Datei replicator.yaml erstellt, in der der Job Checkpoints speichert.

    Die symbolischen Links haben das Format <job-name>-s{step}-n<node-rank>-w<worker-index>.restore.

  2. Stellen Sie jeden Worker aus der entsprechenden .restore-Datei wieder her. Ein Beispiel finden Sie im nächsten Abschnitt.

Standardmäßig ist der LKG-Prüfpunkt der letzte vollständige Prüfpunkt, den Multi-Tier Checkpointing finden kann. Clusterinterne Prüfpunkte (auf der lokalen RAM-Disk und auf Peer-Knoten) werden bevorzugt. Cloud Storage-Sicherungen werden nur verwendet, wenn kein vollständiger clusterinterner Prüfpunkt vorhanden ist. Informationen zum Überschreiben dieser Auswahl finden Sie unter Aus einer bestimmten Sicherung wiederherstellen.

Aus einer bestimmten Sicherung wiederherstellen

Sie können einen Job aus einer bestimmten Cloud Storage-Sicherung anstelle des LKG-Prüfpunkts wiederherstellen, z. B. um das Training auf einen früheren Zustand zurückzusetzen. Führen Sie vor Beginn des Jobs die folgenden Schritte aus:

  1. Suchen Sie die Sicherung, aus der Sie Daten wiederherstellen möchten. Bei der mehrstufigen Prüfpunkterstellung werden Sicherungen in den Bucket geschrieben, den Sie in der CheckpointConfiguration konfiguriert haben, und zwar in ein Verzeichnis, das nach dem Job benannt ist. Jede Sicherung ist ein Unterverzeichnis, das nach der UTC-Zeit benannt ist, zu der die Sicherung erstellt wurde. Das Format ist YYYY-MM-DD_HH-MM:

    gcloud storage ls gs://GCS_PATH/JOB_NAME/

    Ersetzen Sie Folgendes:

    • GCS_PATH: der Cloud Storage-Pfad, in den die mehrstufige Prüfpunkterstellung Sicherungen schreibt. Dieser Wert ist der Bucket, den Sie im Feld cloudStorageBucketName der CheckpointConfiguration angegeben haben, einschließlich aller Verzeichnisse im Bucket.
    • JOB_NAME: Der Name des Trainingsjobs, wie im Feld job-name der Datei replicator.yaml angegeben.

  2. Erstellen Sie im Jobverzeichnis im Bucket eine Datei mit dem Namen restore.version. Die Datei darf nur den Namen des Sicherungsverzeichnisses enthalten, z. B.:

    echo "BACKUP_NAME" | gcloud storage cp - gs://GCS_PATH/JOB_NAME/restore.version

    Ersetzen Sie BACKUP_NAME durch den Namen des Sicherungsverzeichnisses, z. B. 2025-06-30_14-30.

  3. Starten Sie den Job.

Wenn die Datei restore.version vorhanden ist, verhält sich Multi-Tier Checkpointing so:

  • Alle Knoten werden aus der angegebenen Sicherung wiederhergestellt. In-Cluster-Prüfpunkte werden ignoriert, auch wenn sie aktueller als die Sicherung sind.
  • Die Datei restore.version wird automatisch gelöscht, nachdem der Job den ersten Prüfpunkt nach der Wiederherstellung gespeichert hat. Die Überschreibung gilt nur einmal: Wenn der Job später neu gestartet wird, wird er wie gewohnt vom LKG-Prüfpunkt wiederhergestellt.
  • Wenn die angegebene Sicherung nicht vorhanden oder unvollständig ist, wird der Job ohne Checkpoint gestartet. Bei der mehrstufigen Prüfpunkterstellung wird nicht auf clusterinterne Prüfpunkte oder andere Sicherungen zurückgegriffen.

Prüfpunkt speichern

Ihre Anwendung führt diese Schritte während des Trainingsjobs mehrmals aus. Die Speichervorgänge erfolgen am Speicherort CHECKPOINT_DIR, den Sie im JobSet-Manifest der Arbeitslast aktualisieren konfiguriert haben.

Orbax

Orbax-Checkpoints erstellen: Das Verzeichnis ist nach der Schrittnummer benannt. Der Replikator erkennt das neu erstellte Prüfpunktverzeichnis, führt bei Bedarf die Replikation oder Sicherung durch und bereinigt automatisch.

Weitere Informationen zur Verwendung des Orbax-Replicator-Checkpoint-Managers finden Sie in der MaxText-Datei checkpointing. Ein Beispiel für die Interaktion mit dem Replikator-Dienst finden Sie in der MaxText-Datei max_utils.

PyTorch

Verwenden Sie InClusterLocalCheckpointIO als benutzerdefiniertes pytorch_lightning.CheckpointIO, um das korrekte verteilte Checkpointing mit lokalem Speicher zu ermöglichen. Mit dem folgenden Beispielbefehl wird die mehrstufige Prüfpunkterstellung mithilfe einer Referenzimplementierung aktiviert, die auf dem NVIDIA NeMo-Framework basiert:

torchrun train.py <other_train_flags> \
    --local-ckpt-dir=CHECKPOINT_DIR \
    --local-ckpt-interval=20 \
    --job-name=JOB_NAME \
    --enable-high-scale-ckpt

Ersetzen Sie Folgendes:

  • CHECKPOINT_DIR: der Pfad zu Ihrem Prüfpunktverzeichnis.
  • JOB_NAME: Der Name der Arbeitslast Ihres Trainingsjobs.

Clusterupgrades

Bei Clusterupgrades können Sie das CheckpointConfiguration-Objekt entweder vor oder nach dem Upgrade löschen und neu erstellen. Diese Aktion ist erforderlich, da das DaemonSet des Knotensicherungstreibers, das von diesem Objekt dynamisch bereitgestellt wird, nicht automatisch aktualisiert wird.

Wenn Sie die DaemonSet-Spezifikation beibehalten möchten, müssen Sie nichts weiter tun.

Fehlerbehebung

Dieser Abschnitt enthält Hinweise zur Fehlerbehebung bei Problemen mit Multi-Tier Checkpointing. Allgemeine Informationen zur Fehlerbehebung bei Speicherproblemen finden Sie unter Fehlerbehebung bei Cloud Storage in GKE.

Multi-Tier-Prüfpunkte nicht aktiviert

Der folgende Fehler weist darauf hin, dass Multi-Tier Checkpointing in Ihrem Cluster nicht aktiviert ist:

error: unable to recognize "checkpoint.yaml": no matches for kind "CheckpointConfiguration" in version "checkpointing.gke.io/v1"

Dieser Fehler kann auftreten, nachdem Sie kubectl apply -f checkpoint.yaml im Schritt CheckpointConfiguration erstellen ausgeführt haben.

Prüfen Sie mit dem folgenden Befehl, ob Sie Multi-Tier Checkpointing in Ihrem Cluster aktiviert haben, um dieses Problem zu beheben:

gcloud container clusters describe CLUSTER_NAME \
    --project PROJECT_ID
    --location CLUSTER_LOCATION

Wenn mehrstufiges Checkpointing aktiviert ist, sollte die Ausgabe in etwa so aussehen:

addonsConfig:
  gcePersistentDiskCsiDriverConfig:
    enabled: true
  gcsFuseCsiDriverConfig:
    enabled: true
  highScaleCheckpointingConfig:
    enabled: true
  kubernetesDashboard:
    disabled: true
  networkPolicyConfig:
    disabled: true

Wenn Multi-Tier Checkpointing nicht aktiviert ist, aktualisieren Sie Ihren Cluster, um Multi-Tier Checkpointing zu aktivieren.

Der CSI-Treiber für Multi-Tier Checkpointing kann Volumes nicht bereitstellen

Dieses Problem kann auftreten, wenn der CSI-Treiber das Cloud Storage-Volume nicht bereitstellen kann. Es kann mehrere ähnliche Zeilen geben.

kubectl get pod -n gke-managed-checkpointing
NAME                                                          READY   STATUS     RESTARTS   AGE
multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9   0/5     Init:0/1   0          6m32s

Prüfen Sie zur Behebung dieses Problems die CSI-Treiber-Pod-Ereignisse, wie im folgenden Beispiel gezeigt:

kubectl describe pod multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 -n gke-managed-checkpointing

Events:
  Type     Reason       Age                 From               Message
  ----     ------       ----                ----               -------
  Normal   Scheduled    17m                 default-scheduler  Successfully assigned gke-managed-checkpointing/multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 to gke-my-cluster-default-pool-353c773f-6d8q
  Warning  FailedMount  82s (x16 over 17m)  kubelet            MountVolume.SetUp failed for volume "gcs" : rpc error: code = PermissionDenied desc = failed to get GCS bucket "checkpointing-test-bucket": googleapi: Error 403: Caller does not have storage.objects.list access to the Google Cloud Storage bucket. Permission 'storage.objects.list' denied on resource (or it may not exist)., forbidden

Wenn das Problem wie im Beispiel gezeigt auf einen Cloud Storage-Bucket-Fehler PermissionDenied zurückzuführen ist, können Sie es beheben, indem Sie die Berechtigungen richtig einrichten.

Nächste Schritte