Gradient Checkpointing: Wie spart es GPU-Speicher?
Automatische Übersetzung
Dieser Artikel wurde automatisch aus der englischen Originalversion übersetzt.
Gradient Checkpointing speichert beim Vorwärtsdurchlauf weniger intermediäre Activations und berechnet fehlende Werte bei der Rückwärtsberechnung erneut. Verwende es, wenn Activations verhindern, dass der Trainingsbatch oder die Sequenzlänge in den GPU-Speicher passt, und der zusätzliche Rechenaufwand vertretbar ist.
Es reduziert den Speicherbedarf gespeicherter Activations. Es entfernt weder Model Weights noch den Optimizer-Zustand und auch nicht alle temporären Puffer.
Was wird neu berechnet?
Die Rückwärtsberechnung benötigt Zwischenergebnisse des Vorwärtsdurchlaufs, um Gradienten zu berechnen. Ohne Checkpointing bewahrt das Framework die benötigten Ergebnisse auf. Mit Checkpointing speichert es ausgewählte Eingaben oder Zwischenergebnisse und führt die betroffenen Operationen erneut aus, sobald deren Werte benötigt werden.
Ein Framework kann beispielsweise die Eingabe einer Gruppe von Transformer-Schichten speichern und deren interne Activations verwerfen. Bei der Rückwärtsberechnung führt es diese Gruppe mit der gespeicherten Eingabe erneut aus. Die gespeicherte Eingabe ist ein Activation Checkpoint. Dieser unterscheidet sich von einem Model Checkpoint, der zur Wiederherstellung auf Datenträger geschrieben wird.
Die ursprüngliche Analyse von Chen et al. zeigt eine Strategie mit gleichmäßiger Segmentierung, die bei einem Netzwerk mit n Schichten O(sqrt(n)) Speicher für Activations benötigt und ungefähr einen zusätzlichen Vorwärtsdurchlauf pro Trainingsschritt erfordert. Reale Transformer haben unterschiedlich große Activations, verschiedene Attention-Implementierungen und Framework-Regeln. Miss den tatsächlichen Speicherbedarf und die Laufzeit, statt die asymptotische Schätzung als Deployment-Prognose zu übernehmen.
Aktiviere die Integration, die dein Stack erwartet
Bei unterstützten Models im Hugging Face Trainer aktiviert gradient_checkpointing=True den allgemeinen Mechanismus. Prüfe bei FSDP stattdessen die Integration des Frameworks für Activation Checkpointing. Die aktuelle Dokumentation des Transformers Trainer empfiehlt für FSDP activation_checkpointing, weil allgemeines Gradient Checkpointing bei der Rückwärtsberechnung einen zusätzlichen All-Gather verursachen kann.
Vergleiche dieselbe Trainingskonfiguration mit und ohne Checkpointing. Erfasse die maximale GPU-Speicherbelegung, Beispiele oder Tokens pro Sekunde und ob die gewünschte Sequenzlänge in den Speicher passt. Prüfe außerdem, ob die Gradienten gültig bleiben, wenn dein Model benutzerdefinierte Operationen oder Adapter-Training verwendet.
Wenn der Optimizer-Zustand den größten Teil des Speichers belegt, ziehe stattdessen State Sharding oder parameter-effizientes Tuning in Betracht. Das Neuberechnen von Activations beseitigt diesen Speicherbedarf nicht.
Der Abschnitt zu Gradient Checkpointing im LLM Engineering Guide erklärt die segmentierte Berechnung.