Diese Seite wurde von der Cloud Translation API übersetzt.

Trainingsfortschritt mit Autocheckpoint beibehalten

Bisher wurde die Wartung einer TPU-VM sofort gestartet, ohne dass Nutzer Zeit hatten, Aktionen auszuführen, die den Fortschritt sichern, z. B. das Speichern eines Checkpoints. Dies ist in Abbildung 1(a) dargestellt.

Diagramm mit den Auswirkungen der Hostwartung mit und ohne automatisches Checkpointing

Abbildung 1 Abbildung der Funktion „Autocheckpoint“: (a) Ohne Autocheckpoint geht der Trainingsfortschritt ab dem letzten Checkpoint verloren, wenn ein Wartungsereignis bevorsteht. (b) Mit der automatischen Checkpoint-Funktion kann der Trainingsfortschritt seit dem letzten Checkpoint bei einem bevorstehenden Wartungsereignis beibehalten werden.

Mit dem automatischen Checkpoint (Abbildung 1(b)) können Sie den Trainingsfortschritt beibehalten, indem Sie Ihren Code so konfigurieren, dass ein nicht geplanter Checkpoint gespeichert wird, wenn ein Wartungsereignis auftritt. Wenn ein Wartungsereignis auftritt, wird der Fortschritt seit dem letzten Checkpoint automatisch gespeichert. Die Funktion funktioniert sowohl für einzelne Slices als auch für Multi-Slices.

Die Funktion „Autocheckpoint“ funktioniert mit Frameworks, die SIGTERM-Signale erfassen und anschließend einen Checkpoint speichern können. Zu den unterstützten Frameworks gehören:

MaxText,
Pax,
JAX mit Orbax

Autocheckpoint verwenden

Die Funktion „Automatische Checkpoints“ ist standardmäßig deaktiviert. Wenn Sie eine TPU erstellen oder eine in die Warteschlange gestellte Ressource anfordern, können Sie Autocheckpoint aktivieren, indem Sie bei der Bereitstellung der TPU das Flag --autocheckpoint-enabled hinzufügen. Wenn die Funktion aktiviert ist, führt Cloud TPU die folgenden Schritte aus, sobald eine Benachrichtigung zu einem Wartungsereignis empfangen wird:

SIGTERM-Signal erfassen, das über das TPU-Gerät an den Prozess gesendet wird
Warten Sie, bis der Vorgang beendet ist oder fünf Minuten vergangen sind, je nachdem, was zuerst eintritt.
Wartung für die betroffenen Segmente ausführen

Die von Autocheckpoint verwendete Infrastruktur ist unabhängig vom ML-Framework. Jedes ML-Framework kann Autocheckpoint unterstützen, wenn es das SIGTERM-Signal erfassen und einen Prüfpunktprozess initiieren kann.

Im Anwendungscode müssen Sie die vom ML-Framework bereitgestellten Funktionen für automatische Checkpoints aktivieren. In Pax bedeutet das beispielsweise, dass Befehlszeilen-Flags beim Starten des Trainings aktiviert werden müssen. Weitere Informationen finden Sie in der Kurzanleitung für Autocheckpoints mit Pax. Im Hintergrund speichern die Frameworks einen nicht geplanten Checkpoint, wenn ein SIGTERM-Signal empfangen wird. Die betroffene TPU-VM wird dann gewartet, wenn die TPU nicht mehr verwendet wird.

Kurzanleitung: Autocheckpoint mit MaxText

MaxText ist ein leistungsstarker, beliebig skalierbarer, Open-Source-LLM, der in reiner Python/JAX geschrieben wurde und auf Cloud TPUs ausgerichtet ist. MaxText enthält alle erforderlichen Einstellungen für die Verwendung der Funktion „Autocheckpoint“.

In der MaxText-README-Datei werden zwei Möglichkeiten beschrieben, wie Sie MaxText im großen Maßstab ausführen können:

multihost_runner.py wird verwendet, was für Tests empfohlen wird
multihost_job.py wird für die Produktion empfohlen

Wenn Sie multihost_runner.py verwenden, aktivieren Sie „Autocheckpoint“, indem Sie das Flag autocheckpoint-enabled beim Bereitstellen der Ressourcen in der Warteschlange setzen.

Wenn Sie multihost_job.py verwenden, aktivieren Sie den automatischen Checkpoint, indem Sie beim Starten des Jobs das Befehlszeilenflag ENABLE_AUTOCHECKPOINT=true angeben.

Kurzanleitung: Automatischer Checkpoint mit Pax auf einer einzelnen Scheibe

In diesem Abschnitt wird ein Beispiel für die Einrichtung und Verwendung von Autocheckpoint mit Pax auf einem einzelnen Slice beschrieben. Bei entsprechender Einrichtung:

Wenn ein Wartungsereignis auftritt, wird ein Checkpoint gespeichert.
Nach dem Speichern des Checkpoints führt Cloud TPU Wartungsarbeiten an den betroffenen TPU-VMs durch.
Sobald die Wartung von Cloud TPU abgeschlossen ist, können Sie die TPU-VM wie gewohnt verwenden.

Verwenden Sie das Flag autocheckpoint-enabled, wenn Sie die TPU-VM erstellen oder eine in der Warteschlange befindliche Ressource anfordern.

Beispiel:

Legen Sie Umgebungsvariablen fest:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=zone-you-want-to-use
export ACCELERATOR_TYPE=your-accelerator-type
export RUNTIME_VERSION=tpu-ubuntu2204-base

Beschreibungen von Umgebungsvariablen

Variable	Beschreibung
`PROJECT_ID`	Ihre Google Cloud Projekt-ID. Verwenden Sie ein vorhandenes Projekt oder erstellen Sie ein neues.
`TPU_NAME`	Der Name der TPU.
`ZONE`	Die Zone, in der die TPU-VM erstellt werden soll. Weitere Informationen zu unterstützten Zonen finden Sie unter TPU-Regionen und ‑Zonen.
`ACCELERATOR_TYPE`	Der Beschleunigertyp gibt die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Weitere Informationen zu den unterstützten Beschleunigertypen für jede TPU-Version finden Sie unter TPU-Versionen.
`RUNTIME_VERSION`	Die Cloud TPU-Softwareversion.

Legen Sie Ihre Projekt-ID und Zone in Ihrer aktiven Konfiguration fest:

gcloud config set project $PROJECT_ID
gcloud config set compute/zone $ZONE

So erstellen Sie eine TPU:

gcloud alpha compute tpus tpu-vm create $TPU_NAME \
    --accelerator-type $ACCELERATOR_TYPE \
    --version $RUNTIME_VERSION \
    --autocheckpoint-enabled

Stellen Sie eine SSH-Verbindung zur TPU her:

gcloud compute tpus tpu-vm ssh $TPU_NAME

Pax auf einem einzelnen Slice installieren

Die Funktion „Autocheckpoint“ funktioniert mit Pax-Versionen ab 1.1.0. Installieren Sie auf der TPU-VM jax[tpu] und die neueste paxml:
```
pip install paxml && pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```

Konfigurieren Sie das LmCloudSpmd2B-Modell. Ändern Sie vor dem Ausführen des Trainingsscripts ICI_MESH_SHAPE in [1, 8, 1]:

@experiment_registry.register
class LmCloudSpmd2B(LmCloudSpmd):

    """SPMD model with 2B params.

    Global batch size = 2 * 2 * 1 * 32 = 128
    """
    PERCORE_BATCH_SIZE = 8

    NUM_LAYERS = 18
    MODEL_DIMS = 3072
    HIDDEN_DIMS = MODEL_DIMS * 4

    CHECKPOINT_POLICY = layers.AutodiffCheckpointType.SAVE_NOTHING
    ICI_MESH_SHAPE = [1, 8, 1]

Starten Sie das Training mit der entsprechenden Konfiguration.

Im folgenden Beispiel wird gezeigt, wie Sie das LmCloudSpmd2B-Modell so konfigurieren, dass von Autocheckpoint ausgelöste Checkpoints in einem Cloud Storage-Bucket gespeichert werden. Ersetzen Sie your-storage-bucket durch den Namen eines vorhandenen Buckets oder erstellen Sie einen neuen Bucket.
```
export JOB_LOG_DIR=gs://your-storage-bucket

{ python3 .local/lib/python3.10/site-packages/paxml/main.py \
    --jax_fully_async_checkpoint=1 \
    --exit_after_ondemand_checkpoint=1 \
    --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2B \
    --job_log_dir=$JOB_LOG_DIR; } 2>&1 | tee pax_logs.txt
```
Beachten Sie die beiden Flags, die an den Befehl übergeben werden:
- jax_fully_async_checkpoint: Wenn dieses Flag aktiviert ist, wird orbax.checkpoint.AsyncCheckpointer verwendet. Die Klasse AsyncCheckpointer speichert automatisch einen Checkpoint, wenn das Trainingsscript ein SIGTERM-Signal empfängt.
- exit_after_ondemand_checkpoint: Wenn dieses Flag aktiviert ist, wird der TPU-Prozess beendet, nachdem der Autocheckpoint erfolgreich gespeichert wurde. Dadurch wird die Wartung sofort ausgeführt. Wenn Sie dieses Flag nicht verwenden, wird das Training nach dem Speichern des Checkpoints fortgesetzt und Cloud TPU wartet 5 Minuten, bevor die erforderliche Wartung durchgeführt wird.

Automatischer Checkpoint mit Orbax

Die Funktion „Autocheckpoint“ ist nicht auf MaxText oder Pax beschränkt. Jedes Framework, das das SIGTERM-Signal erfassen und einen Checkpoint-Prozess initiieren kann, funktioniert mit der von Autocheckpoint bereitgestellten Infrastruktur. Diese Funktionen bietet Orbax, ein Namespace mit gängigen Dienstbibliotheken für JAX-Nutzer.

Wie in der Orbax-Dokumentation erläutert, sind diese Funktionen für Nutzer von orbax.checkpoint.CheckpointManager standardmäßig aktiviert. Die Methode save, die nach jedem Schritt aufgerufen wird, prüft automatisch, ob ein Wartungsereignis bevorsteht. Falls ja, wird ein Checkpoint gespeichert, auch wenn die Schrittnummer kein Vielfaches von save_interval_steps ist. In der GitHub-Dokumentation wird auch veranschaulicht, wie das Training nach dem Speichern eines automatischen Checkpoints beendet werden kann, indem der Nutzercode geändert wird.