Wartungsereignisse mit Cloud TPU Pods verwalten

Übersicht

TPU-Knoten und TPU-VMs sind Instanzen von Compute Engine-VMs mit angehängter TPU-Hardware. Compute Engine-VMs unterliegen den VM-Wartungsereignissen der Compute Engine. Jede TPU ist mit einer Compute Engine-VM verbunden, sodass die Verwendung von mehr TPUs (z. B. in einem TPU-Pod) die Wahrscheinlichkeit erhöht, dass eine Ihrer VMs einem Wartungsereignis auftritt.

In diesem Dokument werden verschiedene Ansätze zur Verarbeitung von Wartungsereignissen für lang andauernde Trainingsjobs auf Cloud TPUs erläutert.

Prüfpunkte für eine schnelle Wiederherstellung nach Wartungsereignissen verwenden

Prüfpunkte sind der Schlüssel zu kurzen Wiederherstellungen nach Wartungsereignissen und sollten häufig gespeichert werden. Als Faustregel gilt, dass Prüfpunkte etwa stündlich gespeichert werden. Prüfpunkte nicht oft genug zu speichern birgt das Risiko, dass viele Trainingsfortschritte aufgrund von Wartungsereignissen oder anderen Trainingsunterbrechungen verloren gehen.

Prüfpunkte beziehen sich im Allgemeinen auf alle gespeicherten Parameter, die für das Training verwendet werden (z. B. Modellgewichtungen). Die Zeit zum Speichern eines Prüfpunkts kann zwischen Sekunden und einigen Minuten liegen.

Obwohl nach die meisten Wartungsereignisse automatisch wiederhergestellt wird und Trainingsjobs ohne manuelles Eingreifen fortgesetzt werden, kann es in Grenzfällen vorkommen, dass der Job nicht neu gestartet und automatisch fortgesetzt wird. In diesem Fall müssen Sie die TPU-Ressourcen löschen und neu erstellen und den Trainingsjob von einem gespeicherten Prüfpunkt aus neu starten. Informationen zum Erkennen von und Wiederherstellen nach automatischen Wiederherstellungsfehlern finden Sie unten unter Erkennung von und Wiederherstellung nach TPU-Fehlern weiter unten.

Die Verfahren zum Speichern und Laden von Prüfpunkten unterscheiden sich je nach ML-Framework. Unterstützte Cloud TPU-Modelle haben im Allgemeinen eine Prüfpunktausführung. Weitere Informationen zum Checkpointing finden Sie unter TensorFlow 2.x, PyTorch oder JAX/flax.

Wartungsereignisse erkennen

Mit dem folgenden gcloud describe-Befehl können Sie feststellen, ob und wann ein Wartungsereignis auf Ihrer TPU auftritt:

TPU-VMs

$ gcloud compute tpus tpu-vm describe tpu-name --zone=zone  | grep 'health'

TPU-Knoten

$ gcloud compute tpus describe tpu-name --zone=zone | grep 'health'

Die Ausgabe dieses Befehls enthält den aktuellen Status der TPU und eine Beschreibung des letzten Wartungsereignisses. Die Ausgabe sollte in etwa so aussehen:

health: HEALTHY
healthDescription: The TPU had a maintenance event at 2022-01-26T03:44:36.265703305Z

Wartungsereignisse verarbeiten

Es gibt mehrere Möglichkeiten, Unterbrechungen von Wartungsereignissen zu minimieren.

  1. Prüfpunkte regelmäßig speichern

    Im Idealfall wird das Training beim letzten Prüfpunkt fortgesetzt, wenn ein "Unterbrechungsereignis" auftritt.

  2. Wiederholungsversuche von Trainingsskripts

    Das Trainingsskript kann aufgrund eines "Unterbrechungsereignisses" beendet werden. Sie können ein bash-Skript verwenden, um das Trainingsskript kontinuierlich zu wiederholen, bis das Training abgeschlossen ist. Jeder Wiederholungsversuch sollte vom letzten Prüfpunkt aus fortgesetzt werden. Daher sollten Wiederholungsskripts immer in Verbindung mit Prüfpunkten verwendet werden.

    Produktionsreife Trainingspipelines sollten GKE oder andere Ressourcenverwaltungssysteme verwenden. Ansonsten können Sie ein bash-Skript implementieren, um das Trainingsskript kontinuierlich bis zum Trainings-Abschluss zu wiederholen. Beispiel:

    Mit TPU-Knoten:

    (Von Ihrer VM aus) while ! python3 [training command]; do sleep 1; done

    Mit TPU-VM: while ! gcloud compute tpus tpu-vm ssh ${TPU_NAME} --command "python3 [training command]"; do sleep 1; done

    Dabei müssen Sie den TPU-VM-Befehl von der GCP oder von einem Terminal aus ausführen, nicht von der TPU-VM aus.

  3. TPU-Fehler erkennen und davon wiederherstellen

    Wenn eine TPU nicht von einem Wartungsereignis wiederhergestellt wird, können Sie ein TPU-Wiederherstellungsskript verwenden, um den TPU-Status zu erkennen und die TPU zu löschen und neu zu erstellen. Ein Beispiel für dieses Skript finden Sie hier. Weitere Informationen zum manuellen Löschen und Neuerstellen von TPUs finden Sie unter TPUs verwalten.

    Beim Erstellen oder Neuerstellen einer TPU-VM können Sie ein Startskript mit dem Parameter --metadata startup-script angeben. Ein Startskript wird immer dann ausgeführt, wenn eine TPU-VM erstellt wird. Weitere Informationen finden Sie unter Standard-Installationsskripts ausführen.