Gradient Checkpointing

Was es ist

Ein Trainingsschritt besteht aus zwei Hälften. Im Forward Pass wandelt jede Schicht ihre Eingabe in eine Ausgabe um und erzeugt dabei eine Reihe von Zwischenergebnissen: die Aktivierungen. Der Backward Pass durchläuft die Schichten dann in umgekehrter Reihenfolge, um Gradienten zu berechnen, und die Gradienten jeder Schicht benötigen die Aktivierungen dieser Schicht. Der einfache Ansatz behält sie alle im Speicher, bis der Backward Pass sie erreicht. Bei langen Beispielen können Aktivierungen mehr Speicherplatz beanspruchen als die eigenen Gewichte des Modells.

Gradient Checkpointing, auch Activation Checkpointing oder Recomputation genannt, behält nur einige Aktivierungen, die Checkpoints, und verwirft den Rest. Wenn der Backward Pass eine verworfene Aktivierung benötigt, führt er den Forward Pass ab dem nächstgelegenen Checkpoint erneut aus, um sie wiederherzustellen. Chen et al. (2016) zeigten, dass dies den Aktivierungsspeicher für ein n-Schichten-Netzwerk auf die Größenordnung von √n reduziert, zum Preis von etwa einem zusätzlichen Forward Pass pro Batch. In ihren Experimenten schrumpfte der Speicherbedarf eines 1.000-Schichten-Netzwerks von 48 GB auf 7 GB bei etwa 30 % längerer Laufzeit.

Wie es in einem Sprachmodell funktioniert

In Transformer-Bibliotheken befinden sich die Checkpoints in der Regel an den Schichtgrenzen. Die Eingabe jeder Schicht wird beibehalten; alles, was innerhalb der Schicht berechnet wird, wird verworfen und neu berechnet, wenn der Backward Pass dort ankommt. So existiert zu jedem Zeitpunkt nur der vollständige Satz an Aktivierungen einer einzigen Schicht, anstatt aller.

activations per GPT-style layer, 16-bit, efficient attention   (Korthikanti et al., 2022)
  ≈ 34 × batch × tokens × hidden size   bytes

illustrative 2B-class model: hidden size 2,048 · 28 layers · batch 2 × 1,024 tokens
  all kept                 34 × 2 × 1,024 × 2,048 × 28   ≈ 4.0 GB
  checkpointed, inputs      2 × 2 × 1,024 × 2,048 × 28   ≈ 0.23 GB
  + one layer rebuilt at a time                          ≈ 0.14 GB

Wie viele Checkpoints beibehalten werden, ist eine Stellschraube. Mehr beizubehalten bedeutet weniger Neuberechnung, aber mehr Speicherbedarf; weniger beizubehalten bedeutet das Gegenteil. Das Setzen von Checkpoints an jeder Schichtgrenze ist der übliche Standard, da es einfach ist und bereits den Großteil des Aktivierungsspeichers freigibt: Was bleibt, ist eine kleine Eingabe pro Schicht plus die eine Schicht, die neu berechnet wird.

Die Kosten bestehen in der Rechenleistung. Der Forward Pass läuft faktisch zweimal ab: einmal normal und einmal Stück für Stück während des Backward Pass. Der Backward Pass kostet ohnehin schon etwa doppelt so viel wie der Forward Pass, sodass der Mehraufwand höchstens etwa ein Drittel eines Schritts ausmacht; in der Praxis läuft das Training etwa 20–30 % langsamer. Das Ergebnis ändert sich nicht: Die neu berechneten Werte sind dieselben, sodass die Loss-Kurve bis auf winzige numerische Unterschiede einem Durchlauf ohne Checkpointing entspricht.

Warum große Vokabulare es wichtiger machen

Ganz am Ende des Modells erhält jede Position in jedem Beispiel eine Bewertung für jedes Token im Vokabular: die Logits. Deren Größe hängt nicht von der Tiefe des Modells ab, sondern nur davon, wie viele Positionen trainiert werden und wie groß das Vokabular ist. Neuere Modelle haben Vokabulare von 150.000 bis 250.000 Tokens, sodass allein die Logits Gigabytes beanspruchen können.

logits = batch × tokens × vocabulary size × bytes per number

2 × 1,024 × 248,000 × 4 bytes (32-bit)   ≈ 2.0 GB   before their gradients

Es richtig einsetzen

• Nutzen Sie es, bevor Sie Daten kürzen. Das Aktivieren von Checkpointing kostet Zeit; das Verringern der maximalen Länge schneidet die Enden langer Beispiele ab. Wenn der Speicherplatz das Limit ist, ist die Verlangsamung meist der bessere Kompromiss.
• Nutzen Sie den freigewordenen Speicher. Eine größere Batch-Größe oder eine längere maximale Länge kann einen Teil der verlorenen Geschwindigkeit pro Epoche zurückgewinnen.
• Kombinieren Sie es mit QLoRA für die größten Modelle, die eine Maschine fassen kann: QLoRA verkleinert die Gewichte, Checkpointing verkleinert die Aktivierungen.
• Wissen Sie, was davon unberührt bleibt. Gewichte, Gradienten und Optimizer-Status behalten ihre Größe, und es hat keine Auswirkungen auf die Ausführung des fertigen Modells.
• Prüfen Sie die Standardeinstellung. Viele Fine-Tuning-Rezepte aktivieren es bereits. Wenn ausreichend Speicher vorhanden ist und Geschwindigkeit zählt, macht das Deaktivieren jeden Schritt schneller.

Häufige Fehler

• Zuerst die maximale Länge zu reduzieren und die Enden langer Beispiele abzuschneiden, wenn Checkpointing allein ausgereicht hätte, damit sie in den Speicher passen.
• Zu erwarten, dass es die Gewichte oder den Optimizer-Status verkleinert. Es wirkt sich nur auf die Aktivierungen aus.
• Die Verlangsamung von 20–30 % als Fehler der Maschine zu interpretieren. Es ist der Preis für den eingesparten Speicher.
• Anzunehmen, dass es die Kosten der Logits beseitigt. Bei einem riesigen Vokabular kann dennoch ein kleinerer Batch oder eine geringere maximale Länge erforderlich sein.

Quellen und weiterführende Literatur

Externe Quellen werden in einem neuen Tab geöffnet.

1. 01Training Deep Nets with Sublinear Memory Cost↗Chen et al., 2016 · arXiv · arxiv.orgDas Paper, das die Neuberechnung zum Standard machte: √n Speicher für einen zusätzlichen Forward Pass.
2. 02Reducing Activation Recomputation in Large Transformer Models↗Korthikanti et al., 2022 · arXiv · arxiv.orgWoher die Schätzung der Aktivierungen pro Schicht stammt und wie man nur die rechenintensiven Teile neu berechnet.
3. 03torch.utils.checkpoint↗PyTorch docs · docs.pytorch.org
4. 04GPU memory usage↗Hugging Face Transformers docs · huggingface.coGewichte, Gradienten, Optimizer-Status und Aktivierungen: Wofür der Trainingsspeicher verwendet wird.

Verwandte Konzepte