Gradient Checkpointing

Gradienten-Checkpointing ist ein speichersparendes Trainingsverfahren, bei dem nur ausgewählte Zwischenaktivierungen des Vorwärtsdurchlaufs gespeichert und die übrigen während der Rückpropagierung erneut berechnet werden. Das Verfahren wird auch Aktivierungs-Checkpointing genannt und tauscht zusätzlichen Rechenaufwand gegen einen geringeren Spitzenbedarf an Speicher. Trotz seines Namens setzt es Checkpoints für Aktivierungen, nicht für Parametergradienten oder Modelldateien. Dadurch ist es besonders nützlich, wenn Aktivierungstensoren verhindern, dass ein neuronales Netz in den verfügbaren GPU-Speicher passt.

So funktioniert Gradienten-Checkpointing#

Während eines gewöhnlichen Vorwärtsdurchlaufs berechnet ein neuronales Netz Zwischentensoren, sogenannte Aktivierungen. Das automatische Differenzierungssystem behält die Aktivierungen, die später zur Berechnung der Gradienten benötigt werden, wie in den Autograd-Mechanismen von PyTorch beschrieben. Tiefe Netze, große Stapel, hochauflösende Bilder und lange Eingabesequenzen können dazu führen, dass diese gespeicherten Tensoren viel Speicher belegen.

Gradienten-Checkpointing unterteilt das Netz in Segmente:

1. Der Vorwärtsdurchlauf speichert Eingaben oder Aktivierungen an den Grenzen ausgewählter Segmente.
2. Andere Zwischenaktivierungen innerhalb dieser Segmente werden verworfen.
3. Während des Rückwärtsdurchlaufs wird jedes Segment mit Checkpoint erneut vorwärts ausgeführt, um die fehlenden Werte wiederherzustellen.
4. Die wiederhergestellten Aktivierungen werden sofort zur Berechnung der Gradienten verwendet.

Die PyTorch-API für Aktivierungs-Checkpointing stellt dieses Verhalten über torch.utils.checkpoint bereit. Entsprechende Konzepte sind Gradienten-Checkpointing und Rematerialisierung in JAX sowie Tape-Checkpointing in TensorFlow.

Die Platzierung der Checkpoints bestimmt den Kompromiss. Mehr Checkpoint-Bereiche sparen in der Regel mehr Speicher, erfordern aber auch mehr Wiederholungsberechnungen. Eine gezielte Platzierung um aktivierungsintensive Blöcke kann ausgewogener sein als die erneute Berechnung des gesamten Netzes.

Kompromisse und verwandte Verfahren#

Gradienten-Checkpointing lässt Modellparameter, Gradienten und Optimiererzustände in der Regel unverändert. Es zielt hauptsächlich auf den Aktivierungsspeicher ab und verlangsamt das Training, weil einige Vorwärtsberechnungen zweimal ausgeführt werden. Das genaue Ergebnis hängt von der Architektur, den Checkpoint-Grenzen, der Stapelform und der Hardware ab.

Es unterscheidet sich von mehreren verwandten Verfahren:

• Gradientenakkumulation verarbeitet mehrere Mikrostapel, bevor die Modellgewichte aktualisiert werden. So entsteht ein größerer effektiver Stapel, ohne alle Beispiele gleichzeitig laden zu müssen. Gradienten-Checkpointing verringert stattdessen die Zahl der für jeden Mikrostapel gespeicherten Aktivierungen.
• Gemischte Genauigkeit verwendet für ausgewählte Operationen Datentypen mit geringerer Genauigkeit. Der Ablauf für automatische gemischte Genauigkeit in PyTorch kann den Speicherbedarf senken und kompatible Operationen beschleunigen, während Checkpointing bewusst zusätzlichen Rechenaufwand verursacht.
• Stapelgröße verringern senkt den Speicherbedarf, indem weniger Beispiele gemeinsam verarbeitet werden. Durch Checkpointing können ein größerer Stapel oder eine höhere Eingabeauflösung dennoch praktikabel bleiben.
• Trainings-Checkpoints speichern Gewichte und Optimiererzustände für die Wiederherstellung oder spätere Inferenz. Der Leitfaden zu Modell-Checkpoints in PyTorch beschreibt diesen Speichermechanismus, der nichts mit der Neuberechnung von Aktivierungen zu tun hat.

Checkpointierter Code sollte bei der ursprünglichen und der erneut ausgeführten Vorwärtsberechnung funktional konsistent sein. Veränderlicher Zustand, Geräteübertragungen oder nicht kontrollierter Zufall innerhalb eines Checkpoint-Bereichs können Fehler oder falsche Gradienten verursachen. Bei Transformer-Decodern muss möglicherweise auch ein Schlüssel-Wert-Cache während des Checkpoint-Trainings deaktiviert werden, da ein zwischengespeicherter Inferenzzustand mit dem erneuten Aufbau des Vorwärtsberechnungsgraphen in Konflikt geraten kann.

PyTorch-Beispiel#

Das folgende Beispiel setzt während eines Trainingsschritts einen Checkpoint für einen speicherintensiven Block:

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

torch.manual_seed(0)

block = nn.Sequential(
    nn.Linear(1024, 4096),
    nn.ReLU(),
    nn.Linear(4096, 1024),
)
optimizer = torch.optim.AdamW(block.parameters())
inputs = torch.randn(8, 1024, requires_grad=True)
targets = torch.zeros_like(inputs)

optimizer.zero_grad(set_to_none=True)
outputs = checkpoint(block, inputs, use_reentrant=False)
loss = nn.functional.mse_loss(outputs, targets)
loss.backward()
optimizer.step()

Es werden nur die Eingaben des Blocks und die erforderlichen Informationen an den Grenzen gespeichert; seine internen Aktivierungen werden während loss.backward() wiederhergestellt. In realen Projekten sollten Läufe mit und ohne Checkpointing anhand einer Messgröße wie dem maximal zugewiesenen GPU-Speicher in PyTorch verglichen werden.

Anwendungen in der Praxis#

• Hochauflösende Bildverarbeitung: Beim Training für medizinische Segmentierung, Luftbilderkennung und industrielle Prüfung können große Bilder zum Einsatz kommen, deren Merkmalskarten mehr Speicher beanspruchen als die Modellgewichte. Checkpointing für ausgewählte Stufen des Rückgrats kann die Bildauflösung erhalten oder mehr Beispiele pro Stapel ermöglichen, statt die Eingaben stark verkleinern zu müssen.

• Training von Transformern mit langen Sequenzen: Der Speicherbedarf für Aktivierungen steigt rasch mit zunehmender Sequenzlänge und Zahl der Schichten. Durch die Neuberechnung von Transformer-Blöcken können längere Kontexte oder größere Mikrostapel auf demselben Beschleuniger verarbeitet werden. Der NVIDIA-Leitfaden zur Neuberechnung von Aktivierungen veranschaulicht die vollständige und selektive Neuberechnung für Transformer-Schichten.

Praktische Empfehlungen#

Nutze Gradienten-Checkpointing, wenn die Profilerstellung zeigt, dass Aktivierungen und nicht Parameter oder Optimiererzustände den Speicherbedarf dominieren. Beginne mit großen, wiederholten Blöcken und vergleiche den Spitzenspeicherbedarf, die Iterationsdauer und das Validierungsverhalten, bevor du den Umfang erweiterst.

Bei dokumentierten Ultralytics-YOLO-Trainingskonfigurationen gehören automatische gemischte Genauigkeit, Bildgröße und physische Stapelgröße zu den ersten Stellschrauben für den Speicherbedarf. Die Ultralytics-AutoBatch-Referenz erläutert die automatische Auswahl der Stapelgröße anhand des verfügbaren GPU-Speichers. Reicht die lokale Hardware weiterhin nicht aus, stellt das Cloud-Training der Ultralytics Platform konfigurierbare Cloud-GPUs für verwaltete Trainingsläufe bereit. Bei benutzerdefinierten PyTorch-Architekturen kann Gradienten-Checkpointing diese Stellschrauben ergänzen und eine genauere Verwaltung des Aktivierungsspeichers ermöglichen.