Optimizer-State
Beim Training eines Modells merkt sich der Algorithmus nicht nur die aktuellen Gewichte, sondern auch zusätzliche Statistiken über bisherige Lernschritte. Diese gesammelten Zusatzinformationen heißen Optimizer-State. Sie ermöglichen es, ein unterbrochenes Training präzise fortzusetzen.
Warum ein Optimizer mehr als nur Gewichte speichert
Ein Optimizer passt während des Trainings die Gewichte eines Modells Schritt für Schritt an. Einfache Verfahren wie Stochastic Gradient Descent ohne Erweiterungen benötigen dafür nur den aktuellen Gradienten und eine feste Schrittweite. Fortgeschrittene Optimizer hingegen nutzen Informationen aus früheren Trainingsschritten, um bessere Entscheidungen zu treffen.
Beispiel: Ein Optimizer sieht in Schritt 1000, dass der Gradient für ein bestimmtes Gewicht den Wert 0,3 hat. Ohne Gedächtnis würde er nur diesen einen Wert berücksichtigen. Mit Optimizer-State weiß er zusätzlich, dass die Gradienten für dieses Gewicht in den letzten hundert Schritten im Durchschnitt bei 0,25 lagen und wie stark sie schwankten.
Beispiel: Stochastic Gradient Descent ohne Momentum speichert keinen State. Jeder Schritt hängt ausschließlich vom aktuellen Gradienten ab. Erst mit der Momentum-Erweiterung entsteht ein zusätzlicher Zustandsvektor pro Parameter.
Diese zusätzlich gespeicherten Werte bilden den Optimizer-State. Er existiert parallel zu den Modellgewichten und wird bei jedem Trainingsschritt aktualisiert.
Welche Werte Adam und AdamW speichern
Die meisten modernen Trainingsverfahren nutzen AdamW oder eine Variante davon. Diese Optimizer speichern für jedes einzelne Gewicht im Modell zwei gleitende Durchschnitte:
Der erste Moment (oft als m bezeichnet) ist ein exponentiell gewichteter Durchschnitt der bisherigen Gradienten. Er glättet das Rauschen einzelner Trainingsschritte und gibt dem Training eine Art Trägheit in die Richtung, in die es sich zuletzt konsistent bewegt hat.
Der zweite Moment (oft als v bezeichnet) ist ein exponentiell gewichteter Durchschnitt der quadrierten Gradienten. Er misst, wie stark die Gradienten für ein bestimmtes Gewicht schwanken. Parameter mit hoher Varianz erhalten kleinere Schritte, Parameter mit niedriger Varianz größere.
Beispiel: Ein Gewicht in einer frühen Schicht eines Transformer-Modells erhält konsistent Gradienten um 0,01 mit geringer Streuung. Der erste Moment liegt stabil bei 0,01, der zweite Moment ist klein. Der Optimizer macht hier vergleichsweise große Schritte. Ein Gewicht in einer späteren Schicht schwankt zwischen -0,5 und +0,5. Hier ist der zweite Moment groß, und der Optimizer reduziert die Schrittweite automatisch.
Zusätzlich führt Adam einen Schrittzähler (t). Dieser wird für die sogenannte Bias-Korrektur benötigt: Da die gleitenden Durchschnitte bei null initialisiert werden, wären sie in den ersten Trainingsschritten systematisch zu niedrig. Die Korrektur gleicht diesen Effekt aus.
Beispiel: Ohne Bias-Korrektur würde Adam in den ersten zehn Trainingsschritten die Gradienten systematisch unterschätzen, weil die Momentschätzungen noch nahe null liegen. Die Korrektur teilt jeden Moment durch (1 - β^t), wodurch die Schätzungen in frühen Schritten nach oben skaliert werden.
Fachliche Einordnung: Adam kombiniert zwei Techniken, die unabhängig voneinander entwickelt wurden: Momentum (Polyak, 1964) für den ersten Moment und RMSProp (Hinton, 2012) für den zweiten Moment. Die Verbindung beider Ansätze mit Bias-Korrektur wurde 2014 von Kingma und Ba veröffentlicht. AdamW ergänzt eine entkoppelte Gewichtsregularisierung (Loshchilov und Hutter, 2017).
Speicherbedarf und dessen Auswirkungen
Der Optimizer-State verbraucht bei Adam-basierten Verfahren erheblichen Speicher. Für jedes Gewicht werden zwei Fließkommazahlen (erster und zweiter Moment) gespeichert. In FP32 bedeutet das: Der Optimizer-State belegt doppelt so viel Speicher wie die Modellgewichte selbst.
Beispiel: Ein Modell mit 7 Milliarden Parametern in FP32 belegt etwa 28 GB für die Gewichte. Der Optimizer-State für AdamW benötigt zusätzlich 56 GB (zwei Werte pro Parameter in FP32). Zusammen mit den Gradienten (weitere 28 GB) ergibt sich ein Mindestbedarf von 112 GB allein für Training, ohne Aktivierungen und Eingabedaten.
Dieser Speicherbedarf ist ein zentraler Engpass beim Training großer Modelle auf einzelnen Beschleunigern. Ein VRAM-Budget von 80 GB (wie bei einer A100) reicht für das Training eines 7B-Modells in FP32 mit AdamW nicht aus, selbst wenn man die Aktivierungen außer Acht lässt.
Beispiel: Beim Training von GPT-3 mit 175 Milliarden Parametern würde der AdamW-State allein etwa 1,4 TB an FP32-Speicher erfordern. Das übersteigt die Kapazität jeder einzelnen GPU bei Weitem und erfordert verteiltes Training über viele Beschleuniger hinweg.
Techniken zur Reduktion des State-Speichers
Mehrere Ansätze reduzieren den Speicherbedarf des Optimizer-States, ohne die Trainingsqualität wesentlich zu beeinträchtigen:
8-Bit-Optimizer
Die Momentschätzungen werden statt in FP32 in 8-Bit-Formaten gespeichert. Die Bibliothek bitsandbytes von Tim Dettmers implementiert diesen Ansatz für Adam. Pro Parametergruppe wird ein dynamischer Skalierungsfaktor berechnet, der den Wertebereich auf 8 Bit abbildet. Der Speicherbedarf des Optimizer-States sinkt auf ein Viertel.
Beispiel: Ein 7B-Modell mit 8-Bit-AdamW benötigt statt 56 GB nur etwa 14 GB für den Optimizer-State. Zusammen mit Gewichten und Gradienten in FP16 ergibt sich ein Gesamtbedarf von etwa 42 GB, der auf eine einzelne A100 mit 80 GB passt.
Paged Optimizer
Der Optimizer-State wird nicht vollständig im GPU-Speicher gehalten, sondern bei Bedarf zwischen GPU und CPU hin- und hergeschoben. Seiten, die gerade nicht benötigt werden, lagern im Hauptspeicher. Der Nachteil ist ein erhöhter Datentransfer über den PCIe-Bus.
Beispiel: QLoRA (Dettmers et al., 2023) kombiniert Quantisierung der Gewichte mit einem Paged Optimizer. Das ermöglicht das Feinabstimmen eines 65B-Modells auf einer einzigen 48-GB-GPU, weil nur die trainierbaren LoRA-Parameter einen Optimizer-State benötigen und dieser bei Speicherengpässen auf die CPU ausweichen kann.
Optimizer mit reduziertem State
Adafactor (Shazeer und Stern, 2018) speichert statt einer vollständigen Matrix für den zweiten Moment nur Zeilen- und Spaltenmittelwerte. Das reduziert den Speicher für den zweiten Moment von O(n*m) auf O(n+m). Der erste Moment kann optional ganz entfallen.
Beispiel: Für eine Gewichtsmatrix mit 4096 mal 4096 Einträgen speichert Adam den zweiten Moment als vollständige 4096x4096-Matrix (etwa 67 MB in FP32). Adafactor speichert stattdessen einen Zeilenvektor (4096 Werte) und einen Spaltenvektor (4096 Werte), also insgesamt etwa 32 KB. Das ist eine Reduktion um den Faktor 2000.
Optimizer-State in Checkpoints
Ein Checkpoint sichert den vollständigen Trainingszustand zu einem bestimmten Zeitpunkt. Dazu gehören die Modellgewichte, der Optimizer-State, der aktuelle Trainingsschritt und gegebenenfalls der Zustand des Lernraten-Schedulers.
Wird ein Training aus einem Checkpoint fortgesetzt, der keinen Optimizer-State enthält, beginnt der Optimizer effektiv von vorn. Die Momentschätzungen starten bei null, die Bias-Korrektur setzt den Schrittzähler zurück. Das führt zu einem vorübergehenden Anstieg des Loss, weil der Optimizer in den ersten Schritten nach dem Neustart suboptimale Aktualisierungen berechnet.
Beispiel: Ein Training wird nach 50.000 Schritten unterbrochen und aus einem Checkpoint ohne Optimizer-State fortgesetzt. In den nächsten 1000 bis 5000 Schritten liegt der Loss typischerweise höher als unmittelbar vor der Unterbrechung. Die Momentschätzungen müssen sich erst wieder aufbauen. Nach einigen tausend Schritten erreicht der Loss in der Regel wieder das vorherige Niveau.
Bei verteiltem Training über mehrere Beschleuniger wird der Optimizer-State häufig partitioniert (etwa mit ZeRO von DeepSpeed). Jeder Beschleuniger hält nur einen Teil des States. Beim Speichern eines Checkpoints muss der vollständige State aus allen Partitionen zusammengeführt werden.
Unterschiede zwischen Optimizer-Varianten
Der Umfang des Optimizer-States variiert je nach Verfahren erheblich:
Stochastic Gradient Descent ohne Erweiterungen speichert keinen State. Jeder Schritt berechnet sich ausschließlich aus dem aktuellen Gradienten und der Lernrate. Das ist speichereffizient, aber die Konvergenz ist langsamer und empfindlicher gegenüber der Wahl der Lernrate.
SGD mit Momentum speichert einen Vektor pro Parameter (den Momentum-Puffer). Der Speicherbedarf des States entspricht dem der Modellgewichte.
Adam und AdamW speichern zwei Vektoren pro Parameter (erster und zweiter Moment). Der State belegt das Doppelte der Modellgewichte.
LAMB (Layer-wise Adaptive Moments for Batch Training) speichert dieselben Werte wie Adam, berechnet aber die Aktualisierung pro Schicht mit einem angepassten Skalierungsfaktor. Der State-Umfang ist identisch mit Adam, der zusätzliche Rechenaufwand ist minimal.
Beispiel: Beim Training eines Deep-Learning-Modells mit 1 Milliarde Parametern in FP32 belegt der Optimizer-State bei SGD ohne Momentum 0 GB, bei SGD mit Momentum 4 GB, bei Adam/AdamW 8 GB und bei Adafactor (ohne ersten Moment) deutlich weniger als 4 GB, abhängig von der Matrixstruktur der Gewichte.
Grenzen und offene Fragen
Der Optimizer-State ist ein notwendiger Kompromiss zwischen Trainingsqualität und Ressourcenbedarf. Folgende Einschränkungen bestehen:
Die Komprimierung des States (8-Bit, Faktorisierung) führt zu Approximationsfehlern. Bei den meisten Architekturen und Aufgaben sind diese Fehler vernachlässigbar, aber es gibt Szenarien, in denen die reduzierte Präzision die Konvergenz verlangsamt oder die finale Modellqualität leicht verschlechtert. Eine allgemeingültige Aussage, welche Komprimierung immer sicher ist, existiert nicht.
Der Transfer von Optimizer-States zwischen verschiedenen Modellarchitekturen oder bei Änderungen der Modellgröße ist nicht ohne Weiteres möglich. Die Momentschätzungen sind an die spezifische Parameterstruktur gebunden. Bei Architekturänderungen muss der State in der Regel verworfen werden.
Beispiel: Ein Forschungsteam trainiert ein Modell 100.000 Schritte lang und beschließt dann, die Anzahl der Attention-Heads zu erhöhen. Der vorhandene Optimizer-State passt nicht mehr zur neuen Parameterstruktur und muss verworfen werden. Das Training muss entweder von vorn beginnen oder nach der Architekturänderung eine Phase mit erhöhtem Loss in Kauf nehmen.
Die Wahl des Optimizers und damit des State-Umfangs beeinflusst nicht nur den Speicherbedarf, sondern auch die erreichbare Batch-Size. Ein kleinerer Optimizer-State setzt GPU-Speicher frei, der für größere Batches oder längere Sequenzen genutzt werden kann. Diese Abwägung ist bei der Planung eines Trainings ein zentraler Designparameter.
Fachliche Einordnung: Die Forschung zu speichereffizienten Optimizern ist ein aktives Feld. Neuere Ansätze wie GaLore (Zhao et al., 2024) projizieren den Optimizer-State in einen niedrigdimensionalen Unterraum und wechseln die Projektionsrichtung periodisch. Andere Arbeiten untersuchen, ob der zweite Moment bei bestimmten Aufgaben ganz entfallen kann, ohne die Konvergenz wesentlich zu beeinträchtigen. Eine endgültige Lösung für das Speicherproblem ist noch nicht gefunden.