Schnelles, fehlertolerantes PyTorch-Training auf der AI Runtime

Bei großen Systemen wird Ihre Trainingseffizienz durch eine einzige Kennzahl bestimmt: „Goodput“, also der Anteil der Zeit, den Ihre GPUs mit produktiven Berechnungen verbringen, anstatt zu warten oder sich von Ausfällen zu erholen. Da GPU-Ausfälle bei großen Systemen der Regelfall sind, ist die Fähigkeit zur schnellen und automatischen Wiederherstellung nach einem Ausfall der einzige Weg, um einen hohen Goodput aufrechtzuerhalten und Ihre GPU-Gesamtausgaben zu kontrollieren.

Zwei Subsysteme entscheiden über Erfolg oder Misserfolg dieser Wiederherstellung, werden jedoch routinemäßig vernachlässigt: die Datenpipeline, die Ihre Beschleuniger speist, und der Checkpointing-Mechanismus, der den Zustand per Snapshot sichert, damit ein Job fortgesetzt werden kann. Wenn Sie bei einem von beiden einen Fehler machen, kostet Sie jeder Ausfall weitaus mehr GPU-Leerlaufzeit, als es sein müsste. Selbst außerhalb von Ausfallszenarien wird eine Datenpipeline, die nicht mit Ihren Beschleunigern Schritt halten kann, Ihre GPUs unbemerkt unterversorgen und den Goodput genauso sicher verringern wie ein Absturz. Wir gehen die Mechanismen und Kompromisse von beiden durch und zeigen, wie sich die einzelnen Faktoren auf Ihren Goodput und Ihre GPU-Gesamtausgaben auswirken. Im begleitenden Leitfaden für Trainingsleistung und Ausfallsicherheit finden Sie Code-Hinweise und Beispiele.

Für die Infrastrukturseite desselben Problems – wie eine Flotte fehlerhafte GPUs erkennt und isoliert, bevor sie einen Job zum Absturz bringen – lesen Sie den begleitenden Beitrag Wie wir GPUs in Databricks AI zuverlässig halten.

Warum Ausfälle bei großen Systemen der Regelfall sind

Mit zunehmender Anzahl von GPUs in einem Job sinkt die Wahrscheinlichkeit, dass er seine gesamte Laufzeit ohne Unterbrechung übersteht, rapide. Ein nützliches Überschlagsmodell aus dem begleitenden Databricks-Beitrag geht davon aus, dass jede GPU eine jährliche Ausfallrate von etwa 1 % aufweist. Unter dieser Annahme stellt der Beitrag fest, dass „ein Job mit 256 GPUs, der 30 Tage lang läuft, eine Wahrscheinlichkeit von etwa 19 % für einen Ausfall hat. Bei 1.024 GPUs steigt diese auf 57 %.“ Und das sind nur Probleme auf Infrastrukturebene.

Um diese Schätzung an der Realität zu messen: Der Delta-Supercomputer mit 608 H100-GPUs verzeichnete alle 1,9 Stunden Ausfälle. Das bedeutet, dass bei einem Job mit 32 GPUs die durchschnittliche Zeit bis zum Ausfall 36 Stunden betragen würde. Die wichtigste Erkenntnis ist, dass Ihr Trainingsjob wahrscheinlich irgendwann fehlschlagen wird. Die richtigen Entscheidungen können Ihr Modell widerstandsfähig machen und die verlorene Gesamtzeit reduzieren, wenn es passiert.

Auswirkung 1: Das Checkpoint-Format entscheidet darüber, wie oft Sie sich das Speichern leisten können

Beim Checkpointing entscheidet sich, ob Ausfallsicherheit gelingt oder scheitert, und der von Ihnen gewählte Mechanismus hat direkten Einfluss darauf, wie häufig Sie speichern können. Dies ist der größte Hebel für Ihren Goodput: Wenn Sie einmal am Tag einen Checkpoint erstellen, müssen Sie bei einem Ausfall im Durchschnitt 12 Stunden doppelte Arbeit investieren, um wieder den Zustand vor dem Ausfall zu erreichen.

Der monolithische torch.save-Engpass

Der erste Checkpoint, den die meisten Teams schreiben, ist ein einfaches torch.save auf Rank 0. Je nachdem, wie Ihr Modell trainiert wird, gibt es potenziell zwei Probleme:

1. Beim verteilten Training werden alle Zustände auf Rank 0 gesammelt und in einer einzigen Datei gespeichert.
2. Ein einzelner Prozess schreibt den gesamten Checkpoint synchron. Dies kann durch Dinge wie Netzwerkübertragungen blockiert werden, wenn in Remote-Objektspeichern wie Unity Catalog (UC) gespeichert wird.

Dieses blockierende Verhalten führt zu GPU-Leerlaufzeiten und verringert Ihren Goodput. Es gibt jedoch eine Möglichkeit, die Zeit zu verkürzen, die Ihre GPU für das Checkpointing aufwendet: die Distributed-Checkpoint-API von Torch.

Distributed Checkpoint (DCP): Jeder Rank schreibt sein eigenes Shard

Der Distributed Checkpoint von PyTorch kehrt dieses Prinzip um. Jeder Rank schreibt parallel sein eigenes, separates Shard, zusammen mit einer kleinen .metadata-Datei, die beschreibt, wie sich die Shards zu den vollständigen Tensoren zusammensetzen.

Die Speicherzeit sinkt mit der Anzahl der Ranks um etwa 1/N. Da die .metadata-Datei das globale Layout aufzeichnet, kann derselbe Checkpoint auf einer anderen Anzahl von GPUs neu geladen werden. DCP plant neu, welche Bytes jeder neue Rank benötigt, sodass die Wiederherstellung auf einem Cluster mit reduzierter Kapazität nach dem Verlust von Nodes problemlos funktioniert.

DCP lohnt sich selbst bei einfachen datenparallelen Jobs

Eine häufige Annahme ist, dass DCP nur für Sharded-Modelle gedacht ist und dass ein datenparalleler (DDP) Job, bei dem jeder Rank eine identische Replik der Gewichte enthält, davon nicht profitiert. Das stimmt nicht: DCP teilt den Modellzustand in Shards auf und schreibt ihn selbst bei DDP-Trainingsaufgaben parallel über alle Worker hinweg.

Es ist außerdem dieselbe API, die Sie benötigen, wenn Sie auf FSDP oder Tensor-Parallelität umsteigen. Eine frühzeitige Einführung bedeutet also, dass Sie Ihren Code für die Ausfallsicherheit nie im ungünstigsten Moment neu schreiben müssen.

Asynchrones Speichern macht die Häufigkeit nahezu kostenlos

Selbst bei parallelem Schreiben blockiert ein synchrones Speichern das Training, bis die Bytes dauerhaft im Speicher abgelegt sind – bei einem großen Checkpoint auf einem Remote-Volume bedeutet das Dutzende Sekunden Leerlaufzeit für die Beschleuniger. async_save teilt den Vorgang auf: ein schnelles Kopieren in einen Staging-Puffer, gefolgt von einem Hintergrund-Upload, der sich mit dem fortlaufenden Training überschneidet.

Der Trainings-Loop zahlt nur für die Staging-Kopie, nicht für den Upload. Ein Checkpoint, der früher Dutzende Sekunden Leerlaufzeit kostete, kostet jetzt fast nichts mehr. Genau das macht das häufige Checkpointing im nächsten Abschnitt erschwinglich.

Auf der AI Runtime implementieren UCVolumeWriter and UCVolumeReader DCP für UC-Volumes, wobei I/O über lokales NVMe zwischengespeichert wird und ein Checkpoint erst dann als abgeschlossen markiert wird, wenn seine Daten vollständig übertragen wurden. Weitere Details und Codebeispiele finden Sie im Leitfaden für Leistung und Ausfallsicherheit.

Die obigen Angaben enthalten nicht die Netzwerkspeicherzeit für torch.save.

Auswirkung 2: Die Checkpoint-Häufigkeit entscheidet über Ihre Wiederherstellungskosten

Hier summieren sich die Effekte. Wenn ein Job fehlschlagt, geht alles verloren, was seit dem letzten gültigen Checkpoint berechnet wurde, und muss neu berechnet werden. Der erwartete Arbeitsverlust pro Ausfall entspricht also etwa der Hälfte des Checkpoint-Intervalls, und durch günstiges asynchrones Speichern können Sie dieses Intervall kurz halten.

Eine Verkürzung des Intervalls um den Faktor 10 verkürzt die erwartete Wiederherstellungszeit um den Faktor 10. Denken Sie an die Llama 3-Zahl von ca. 8,6 Unterbrechungen pro Tag: Bei dieser Ausfallrate bedeutet ein Checkpoint alle 2 Stunden, dass Sie voraussichtlich 8,6 Stunden pro Tag mit erneutem Training verschwenden, was einem Goodput von 64 % entspricht. Bei einem Checkpoint alle 30 Minuten verlieren Sie nur 2,15 Stunden, was einem Goodput von 91 % entspricht.

Die Wiederherstellung muss zudem automatisch erfolgen. Bei einem Neustart sollte der Job den neuesten vollständig geschriebenen Checkpoint finden, alle durch den Absturz nur halb geschriebenen Checkpoints überspringen und ohne menschliches Zutun fortgesetzt werden. DCP macht dies zuverlässig: Die .metadata-Datei wird erst geschrieben, nachdem alle Shards übertragen wurden. Ihr Vorhandensein ist also ein vertrauenswürdiger Indikator dafür, dass dieser Speichervorgang abgeschlossen ist.

Auswirkung 3: Das Laden von Daten entscheidet darüber, ob Ihre GPUs jemals stillstehen

Ein Trainingsjob läuft nur so schnell wie seine langsamste Eingabe. Wenn Beschleuniger auf den nächsten Batch warten, sinkt Ihr Goodput, da Ihre GPUs schlichtweg stillstehen. Die einzige Möglichkeit, dieses Problem zu beheben, besteht darin, sicherzustellen, dass sich Ihre Eingabepipeline bei der Datenvorbereitung für den nächsten Schritt mit der Berechnung des aktuellen Schritts überschneidet, wie in der folgenden Abbildung dargestellt:

Wir stellen häufig fest, dass Kunden, die das Laden von Daten mit den Berechnungen überschneiden lassen, eine Reduzierung der tatsächlichen Laufzeit (Wall-Clock-Time) um 20 bis 50 % erzielen.

Die Kosten für das direkte Lesen aus dem Remote-Speicher

Auf einer verwalteten Plattform befinden sich die Trainingsdaten in einem Remote-Objektspeicher. Auf der AI Runtime werden Unity Catalog (UC)-Volumes als Netzwerk-Mounts bereitgestellt.

Das direkte Lesen von Dateien von diesem Mount bei jedem Zugriff bindet Ihre Schrittzeit an die Netzwerklatenz und lädt dieselben Dateien in jeder Epoche erneut herunter. Die Lösung ist ein Dataloader, der jede Datei beim ersten Zugriff in einen schnellen lokalen Speicher kopiert, nachfolgende Lesevorgänge aus diesem lokalen Cache bedient und anstehende Dateien parallel abruft, während die GPU Berechnungen durchführt.

Mit AI Runtime tun UCVolumeDataset und DataLoader genau das (Codebeispiele finden Sie im Leitfaden). UCVolumeDataset streamt Dateien aus einem UC-Volume, speichert jede beim ersten Zugriff im lokalen NVMe-Cache und partitioniert Dateien über Ranks und Worker hinweg, sodass jeder Beschleuniger ein disjunktes, überschneidungsfreies Segment erhält. Unser DataLoader ist eine Drop-in-Unterklasse des PyTorch-DataLoader, deren Standardeinstellungen für diesen Pfad optimiert sind. So werden Dateien gleichzeitig abgerufen und zwischengespeichert, während die GPU rechnet, anstatt nacheinander im Trainings-Thread.

Beispiel: Trainieren eines Bildmodells mit UC-Dateien

Betrachten wir einen einfachen Workload zur Bildklassifizierung: JPEGs aus einem UC-Volume decodieren, augmentieren und ein Vision-Modell trainieren. Sehen wir uns zwei Möglichkeiten an, dies auf derselben GPU, demselben Modell und derselben Batch-Größe zu tun: der standardmäßige PyTorch-Dataset, der aus einem UC-Volume liest, im Vergleich zu UCVolumeDataset plus den Databricks-DataLoaderStandardeinstellungen.

Sie müssen nicht raten, wohin die Zeit fließt

Als Teil der Entwicklung von DataLoader haben wir sichergestellt, dass es seine Metriken in MLflow protokolliert. So sehen Sie auf einen Blick, ob Ihre Datenpipeline das Training blockiert.

Die Metrik fetch_seconds misst explizit, wie lange der DataLoader braucht, um einen Batch zu erstellen, während Ihre GPU in dieser Zeit ungenutzt bleibt.

Auswirkung 4: Das Vergessen der Datenpipeline führt zu einer unbemerkten Beschädigung Ihres Modells

Es gibt einen letzten Resilienz-Fehler, der keine Fehlermeldung, keinen Absturz und keinen fehlgeschlagenen Job verursacht, sondern lediglich ein Modell, das minimal schlechter ist, als es sein sollte. Dies passiert, wenn Sie das Modell, den Optimizer und den Schritt per Checkpoint sichern, aber nicht die Position Ihrer Datenpipeline innerhalb des Datasets.

Stellen Sie sich vor, ein Job wird mitten in einer Epoche unterbrochen. Er stellt das Modell korrekt wieder her und setzt die Trainingsschleife fort, aber der DataLoader beginnt wieder von vorne im Dataset.

Der fortgesetzte Job trainiert erneut mit Beispielen, die er in dieser Epoche bereits gesehen hat, und überspringt potenziell diejenigen, die er noch nicht erreicht hatte. Bei den vielen Neustarts, die bei großen Skalierungen üblich sind, verzerrt dies unbemerkt Ihre Datenverteilung. Das Modell trainiert zwar weiter, aber mit der falschen Stichprobe Ihrer Daten – genau die Art von unbemerktem Fehler, die am kostspieligsten ist, da der Job erfolgreich abgeschlossen wird und niemand ein Problem bemerkt, bis die Metriken enttäuschen.

Die Lösung besteht darin, die Datenposition als Teil des Checkpoints zu behandeln. Je nach Pipeline bedeutet dies, einen Sample- oder Shard-Offset zu verfolgen und beim Fortsetzen nach vorne zu springen, ein benutzerdefiniertes Dataset seine eigene Position serialisieren zu lassen oder Checkpoints an Epochengrenzen zu erstellen. All dies setzt eine Voraussetzung voraus: Determinismus. Shuffling und Augmentation nutzen Zufallszahlengeneratoren. Daher müssen diese Seeds und RNG-Zustände ebenfalls Teil des Checkpoints sein, da andernfalls die Datenreihenfolge nach einem Neustart nicht mit der Reihenfolge davor übereinstimmt und eine gespeicherte Position auf die falschen Samples verweist.

Seed, reproduzierbare Reihenfolge und fortsetzbare Datenpipeline sind drei Ausdrücke derselben Idee. Der Leitfaden beschreibt jede Strategie mit Code.

Zusammenfassung

Schnelles, fehlertolerantes Training resultiert aus einer Handvoll Entscheidungen, die sich gegenseitig verstärken:

1. Verwenden Sie Distributed Checkpoint anstelle von torch.save, selbst für DDP, damit Speichervorgänge parallel und ressourcenschonend ablaufen, anstatt einen seriellen Engpass darzustellen.
2. Speichern Sie asynchron, sodass Checkpoints fast kostenlos sind, was es Ihnen ermöglicht, häufig zu speichern.
3. Führen Sie eine automatische Wiederherstellung auf den neuesten gültigen Checkpoint durch, sodass ein Ausfall nur Minuten an Neuberechnung kostet, nicht Stunden.
4. Überschneiden Sie das Laden von Daten mit der Berechnung durch Caching und Prefetching aus dem Remote-Speicher, sodass Beschleuniger nie ungenutzt auf Eingaben warten. Dadurch werden bei jedem Schritt GPU-Stunden eingespart.
5. Sichern Sie die Datenpipeline und den RNG-Zustand per Checkpoint, damit ein fortgesetzter Job mit den richtigen Daten fortfährt, anstatt Ihr Modell unbemerkt zu beschädigen.

Das vereinheitlichende Prinzip: Häufige, kostengünstige und vollständige Checkpoints verwandeln einen Hardwareausfall von einem jobbeendenden Ereignis in einen Rundungsfehler, und eine überlappende Eingabepipeline hält die Beschleuniger in der Zwischenzeit ausgelastet. Günstiges (asynchrones) Speichern macht eine hohe Frequenz erschwinglich; vollständiges Speichern (Modell, Daten und RNG) stellt eine korrekte Wiederherstellung sicher. Wenn beides gegeben ist und eine Flotte vorhanden ist, die fehlerhafte Hardware erkennt und isoliert, nähert sich Ihre effektive Trainingszeit dem Maximum, das die Hardware zulässt – unabhängig davon, wie instabil der zugrunde liegende Cluster ist.

Referenzen

• Meta, The Llama 3 Herd of Models (2024): §3.3.1, Trainingszuverlässigkeit und Aufschlüsselung von Unterbrechungen auf bis zu 16.384 H100-GPUs.
• Characterizing the Resilience of Hopper H100 and Ampere A100 GPUs (2025): 2,5-jährige Feldstudie des Delta-Systems; MTBE und Ausfallwahrscheinlichkeit nach Jobgröße.
• Databricks, Wie wir GPUs in Databricks AI zuverlässig halten: GPU-Integritätsprüfung auf Flottenebene und das probabilistische Ausfallmodell.

Bereit, es auszuprobieren? Den vollständigen Code finden Sie im Leitfaden für Trainingsleistung und -resilienz in der Dokumentation zur Databricks AI Runtime. Die infrastrukturelle Seite der Geschichte können Sie unter Wie wir GPUs in Databricks AI zuverlässig halten nachlesen.

(Dieser Blogbeitrag wurde mit KI-gestützten Tools übersetzt.) Originalbeitrag