Gradient Checkpointing

Gradient Checkpointing ist eine speichersparende Trainingsmethode, die nur ausgewählte intermediäre Aktivierungen aus dem Vorwärtsdurchlauf speichert und die anderen während der backpropagation neu berechnet. Auch als Aktivierungs-Checkpointing bezeichnet, tauscht sie zusätzliche Rechenleistung gegen einen geringeren Spitzenspeicherbedarf ein. Entgegen dem Namen sichert die Technik Aktivierungen anstelle von Parametergradienten oder Modelldateien ab, was sie besonders wertvoll macht, wenn Aktivierungstensoren verhindern, dass ein neuronales Netz in den verfügbaren GPU memory passt.

Wie Gradient Checkpointing funktioniert#

Während eines standardmäßigen Vorwärtsdurchlaufs berechnet ein neuronales Netz intermediäre Tensoren namens Aktivierungen. Das System zur automatischen Differenzierung behält die Aktivierungen bei, die für die spätere Gradientenberechnung erforderlich sind, wie in den PyTorch autograd mechanics beschrieben. Tiefe Netzwerke, große Batches, hochauflösende Bilder und lange Eingabesequenzen können dazu führen, dass diese gespeicherten Tensoren erheblichen Speicherplatz verbrauchen.

Gradient Checkpointing unterteilt das Netzwerk in Segmente:

1. Der Vorwärtsdurchlauf speichert Eingaben oder Randaktivierungen für ausgewählte Segmente.
2. Andere intermediäre Aktivierungen innerhalb dieser Segmente werden verworfen.
3. Während des Rückwärtsdurchlaufs wird jedes mit Checkpoints versehene Segment erneut vorwärts ausgeführt, um die fehlenden Werte zu rekonstruieren.
4. Die rekonstruierten Aktivierungen werden sofort zur Berechnung von Gradienten verwendet.

Die PyTorch activation checkpointing API macht dieses Verhalten über torch.utils.checkpoint zugänglich. Ähnliche Konzepte erscheinen als JAX gradient checkpointing and rematerialization und TensorFlow tape checkpointing.

Die Platzierung von Checkpoints bestimmt den Kompromiss. Das Setzen von Checkpoints in mehr Regionen spart im Allgemeinen mehr Speicher, wiederholt jedoch mehr Operationen. Eine selektive Platzierung um aktivierungsintensive Blöcke herum kann eine bessere Balance bieten als die Neuberechnung des gesamten Netzwerks.

Kompromisse und verwandte Techniken#

Gradient Checkpointing lässt Modellparameter, Gradienten und Optimizer-Zustände normalerweise unverändert. Ihr Hauptziel ist der Aktivierungsspeicher, und das Training wird langsamer, da einige Vorwärtsberechnungen zweimal ausgeführt werden. Das genaue Ergebnis hängt von der Architektur, den Checkpoint-Grenzen, der Batch-Form und der Hardware ab.

Es unterscheidet sich von mehreren verwandten Techniken:

• Gradient accumulation verarbeitet mehrere Mikro-Batches vor der Aktualisierung der Modellgewichte, wodurch ein größerer effektiver Batch erzeugt wird, ohne jedes Sample gleichzeitig zu laden. Gradient Checkpointing reduziert stattdessen die für jeden Mikro-Batch beibehaltenen Aktivierungen.
• Mixed precision verwendet Datentypen mit geringerer Präzision für ausgewählte Operationen. Der PyTorch automatic mixed precision workflow kann den Speicher reduzieren und kompatible Operationen beschleunigen, während Checkpointing bewusst Rechenleistung hinzufügt.
• Batch size-Reduzierung senkt den Speicherbedarf, indem weniger Samples gemeinsam verarbeitet werden. Checkpointing kann es ermöglichen, dass ein größerer Batch oder eine höhere Eingabeauflösung praktikabel bleibt.
• Training-Checkpoints speichern Gewichte und den Optimizer-Zustand für die Wiederherstellung oder spätere Inferenz. Der PyTorch model checkpoint guide beschreibt diesen Persistenzmechanismus, der nichts mit der Aktivierungsneuberechnung zu tun hat.

Mit Checkpoints versehener Code sollte zwischen seinen ursprünglichen und neu berechneten Vorwärtsausführungen funktional konsistent sein. Veränderbare Zustände, Geräteübertragungen oder unkontrollierte Zufälligkeit innerhalb einer mit Checkpoints versehenen Region können zu Fehlern oder falschen Gradienten führen. In Transformer-Modellen im Decoder-Stil muss möglicherweise auch ein key-value cache während des Trainings mit Checkpoints deaktiviert werden, da der im Cache gespeicherte Inferenzzustand im Widerspruch zum Aufbau des Vorwärtsgraphen stehen kann.

PyTorch-Beispiel#

Das folgende Beispiel versieht einen speicherintensiven Block während eines Trainingsschritts mit einem Checkpoint:

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 Randinformationen beibehalten; seine internen Aktivierungen werden während loss.backward() rekonstruiert. Reale Projekte sollten Durchläufe mit und ohne Checkpointing vergleichen, indem sie eine Messung wie den PyTorch peak allocated GPU memory verwenden.

Praxisanwendungen#

• Hochauflösende Computer Vision: Medizinisches Segmentieren, Luftbilderkennung und industrielle Inspektion werden möglicherweise mit großen Bildern trainiert, deren Feature-Maps mehr Speicher als die Modellgewichte verbrauchen. Das Setzen von Checkpoints in ausgewählten Backbone-Stufen kann die Bildauflösung beibehalten oder zusätzliche Samples pro Batch ermöglichen, anstatt Eingaben aggressiv herunterskalieren zu müssen.

• Langes Sequenz-Transformer-Training: Die Aktivierungsspeicherung wächst rasant, wenn Sequenzlänge und Schichtenanzahl zunehmen. Die Neuberechnung von Transformer-Blöcken kann dazu führen, dass längere Kontexte oder größere Mikro-Batches auf demselben Beschleuniger Platz finden. Der NVIDIA activation recomputation guide veranschaulicht die vollständige und selektive Neuberechnung für Transformers-Schichten.

Praktische Hinweise#

Verwende Gradient Checkpointing, wenn das Profiling zeigt, dass Aktivierungen anstelle von Parametern oder dem Optimizer-Zustand den Speicher dominieren. Beginne mit großen, wiederholten Blöcken und benchmarke den Spitzenspeicher, die Iterationszeit und das Validierungsverhalten, bevor du die Abdeckung ausweitest.

Für dokumentierte Ultralytics YOLO training configurations umfassen Speichersteuerungen erster Linie automatische gemischte Präzision, Bildgröße und physische Batch-Größe. Die Ultralytics AutoBatch reference erklärt die automatische Batch-Auswahl basierend auf dem verfügbaren GPU-Speicher. Wenn die lokale Hardware unzureichend bleibt, bietet Ultralytics Platform cloud training konfigurierbare Cloud-GPUs für verwaltete Trainingsläufe. Gradient Checkpointing kann diese Steuerungen ergänzen, wenn eine benutzerdefinierte PyTorch-Architektur eine feinere Aktivierungs-Speicherverwaltung erfordert.