Was ist Gradient Checkpointing?

Gradient-Checkpointing ist eine Speicheroptimierungstechnik beim Training tiefer neuronaler Netze. Die Idee ist, nicht alle Zwischenaktivierungen während des Vorwärtsdurchlaufs zu speichern, sondern nur einige Kontrollpunkte (Checkpoints). Bei der Rückwärtsausbreitung werden die fehlenden Aktivierungen durch erneute Berechnung eines Teils des Vorwärtsdurchlaufs rekonstruiert.

Dies ermöglicht eine erhebliche Reduzierung des Speicherverbrauchs, was besonders beim Training sehr tiefer Modelle oder bei begrenzten GPU-Ressourcen wichtig ist. Allerdings erhöht sich dadurch die Trainingszeit, da einige Berechnungen wiederholt werden müssen.

Beispiel für die Verwendung in PyTorch:

from torch.utils.checkpoint import checkpoint

def custom_forward(x):
    # komplexer Teil des Modells
    return x * x

output = checkpoint(custom_forward, input_tensor)

Daher ist Gradient-Checkpointing ein Kompromiss zwischen Zeit und Speicher, der das Training großer Modelle mit begrenzten Ressourcen ermöglicht.