🤖 machine learning

FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training

FORGE ist eine speichereffiziente LLM-Trainingsmethode, die den Optimierungsschritt in den Backward-Pass integriert, um materialisierte Gradienten zu eliminieren, indem sie diese vollständig in Registern verarbeitet, wodurch der Speicherbedarf des Optimierers halbiert, das Training beschleunigt und die volle Präzisionsgenauigkeit bewahrt wird, ohne die zugrunde liegende Update-Logik zu verändern.

Ursprüngliche Autoren: Dikshant Kukreja, Kritarth Prasad, Avinash Anand, Zhengkui Wang, Erik Cambria, Timothy Liu, Aik Beng Ng, Simon See, Bapi Chatterjee

Veröffentlicht 2026-06-23

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Dikshant Kukreja, Kritarth Prasad, Avinash Anand, Zhengkui Wang, Erik Cambria, Timothy Liu, Aik Beng Ng, Simon See, Bapi Chatterjee

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Das große Problem: Die „überladene Werkstatt“

Stellen Sie sich vor, Sie trainieren ein riesiges KI-Modell (wie ein Robotergehirn) auf einem Computer. Um diesen Roboter zu lehren, muss der Computer eine gewaltige Menge an Mathematik betreiben.

Auf die Standard-Art und Weise (genannt „Reverse-mode differentiation“) folgt der Computer einem strengen zweistufigen Prozess:

Die Fehler berechnen: Er schaut sich die Daten an, findet heraus, wo der Roboter einen Fehler gemacht hat, und schreibt eine riesige Liste mit „Korrekturhinweisen“ (Gradienten) für jeden einzelnen Teil des Robotergehirns auf. Er schreibt diese Notizen auf ein riesiges Whiteboard (den Speicher des Computers).
Die Korrekturen anwenden: Erst nachdem das gesamte Whiteboard vollgeschrieben ist, nimmt der Computer einen Radiergummi und einen Marker in die Hand, um das Robotergehirn basament auf Basis dieser Notizen tatsächlich zu aktualisieren.

Der Flaschenhals: Das Problem ist, dass der Computer dieses riesige Whiteboard voller Notizen in seinem Speicher halten muss, während er gleichzeitig auch das Robotergehirn und die Notizen für den nächsten Schritt hält. Dieses „Whiteboard“ nimmt so viel Platz ein, dass der Speicher des Computers oft schon voll ist, bevor das Training überhaupt richtig beginnen kann. Es ist, als würde man versuchen, ein Auto in einer Garage zu reparieren, aber man muss gleichzeitig den gesamten Bauplan des Autos, die Werkzeuge und das Reparaturhandbuch auf dem Boden ausgebreitet halten. Wenn die Garage zu klein ist, passt das Auto gar nicht erst hinein.

Die Lösung: FORGE (Die „Sofort-Fix“-Methode)

Die Autoren dieses Papers, FORGE (Fused On-Register Gradient Elimination), sagen: „Warum überhaupt Notizen auf dem Whiteboard aufschreiben?“

Sie argumentieren, dass die riesige Liste von Notizen nur ein Nebenprodukt dessen ist, wie wir die Mathematik betreiben, und nicht etwas, das der Lernprozess eigentlich benötigt.

Wie FORGE funktioniert:
Anstatt die Notizen aufzuschreiben und sie dann wieder abzulesen, erledigt FORGE die Mathematik und das Reparieren gleichzeitig im schnellsten, kleinsten Speicherbereich des Computers (den sogenannten „Registern“).

Die Analogy: Stellen Sie sich einen Meisterkoch vor, der ein komplexes Gericht kocht.
- Der alte Weg: Der Koch schneidet eine Zwiebel, schreibt „Zwiebel geschnitten“ auf einen Notizblock, legt die Zwiebel in eine Schüssel und geht zum nächsten Zutatenpunkt über. Sobald alle Zutaten geschnitten und aufgelistet sind, liest der Koch den Notizblock und mischt alles zusammen. Der Notizblock nimmt Arbeitsfläche weg.
- Der FORGE-Weg: Der Koch schneidet eine Zwiebel und wirft sie sofort in den Topf. Dann schneidet er eine Karotte und wirft sie sofort hinein. Er schreibt nie etwas auf. Er braucht keinen Notizblock. Die Arbeitsfläche bleibt frei.

Warum das eine große Sache ist

Das Paper behauptet drei wesentliche Vorteile durch diesen „No-Notes“-Ansatz:

1. Es spart massiv Platz (Speicher)
Da der Computer die riesige Liste von Notizen nie in den Hauptspeicher schreibt, wird eine enorme Menge an Platz frei.

Das Ergebnis: Auf einem Standard-Chip (H200) konnten sie ein Modell mit 8 Milliarden Parametern trainieren, wobei sie 53 % weniger Speicher benötigten.
Die Analogie: Es ist, als ob man einen Esstisch für 10 Personen in ein Einzimmerapartment stellen könnte, weil man aufgehört hat, die zusätzlichen Stühle im Flur zu lagern.

2. Es ist tatsächlich schneller
Da der Computer nicht anhalten muss, um Notizen zu schreiben, zu warten und sie dann wieder abzulesen, beschleunigt sich der gesamte Prozess.

Das Ergebnis: Die Trainingsschritte laufen etwa 1,5-mal schneller.
Die Analogie: Es ist der Unterschied zwischen einem Lieferfahrer, der ein Paket abliefert, zurück zum LKW rennt, um das nächste zu holen, und das wiederholt (alter Weg), gegenüber einem Fahrer, der ein Förderband hat, das das Paket direkt auf den LKW lädt, während er es greift (FORGE).

3. Es verliert nicht an Genauigkeit
Normalerweise verliert man etwas Präzision (der Roboter lernt etwas schlechter), wenn man versucht, Platz zu sparen, indem man Schritte überspringt. Die Autoren beweisen, dass sie, da sie die Mathematik in den hochwertigsten „Registern“ (voller Präzision) durchführen, bevor sie die Daten verwerfen, das Lernen mathematisch identisch mit der alten, langsamen Methode ist.

Das Ergebnis: Der Roboter lernt genauso gut, aber viel effizienter.

Der „Tile“-Trick

Wie machen sie das ohne Chaos? Sie zerlegen das riesige mathematische Problem in kleine, handhabbare Stücke, sogenannte „Tiles“ (wie 128x128 Quadrate).

Sie bearbeiten ein Tile: Berechnen den Fehler, reparieren das Gehirn und werfen den Fehler sofort weg.
Dann bewegen sie sich zum nächsten Tile.
Da sie dies Tile für Tile tun, muss der Computer niemals die gesamte Liste der Fehler gleichzeitig halten.

Was uns das ermöglicht

Das Paper zeigt, dass man mit FORGE:

Größere Modelle auf demselben Computer trainieren kann.
Größere „Batches“ verwenden kann (dem Roboter mehr Beispiele gleichzeitig beibringen), ohne dass der Speicher voll läuft.
In einem speziellen Test konnten sie ein Modell auf vier GPUs trainieren, das mit der alten Methode acht GPUs benötigt hätte.

Zusammenfassung

FORGE ist eine neue Art, KI zu trainieren, die verhindert, dass der Computer seinen Speicher mit temporären „Korrekturhinweisen“ überlädt. Indem es die Fehler berechnet und das Modell sofort im schnellsten Teil des Chips repariert, halbiert es den Speicherverbrauch und beschleunigt das Training, ohne die KI weniger intelligent zu machen. Es verwandelt eine überladene, langsame Werkstatt in ein effizientes, Hochgeschwindigkeits-Fließband.

Technisches Resümee: FORGE (Fused On-Register Gradient Elimination)

1. Das Problem: Der Gradienten-Boden (Gradient Floor)

Aktuelle speichereffiziente Trainingsmethoden für Large Language Models (LLMs) zielen primär auf drei der vier Speicherpools auf einer GPU ab: Optimizer-Zustände (durch Kompression), Aktivierungen (durch Recomputation) und Modellgewichte. Sie lassen jedoch den vierten Pool – den materialisierten Gewichtsgradienten – unberührt.

Im Standard-Reverse-Mode-Differenzierungsprozess berechnet der Backward-Pass den vollständigen Gradiententensor ( $\nabla W$ ) für jede Schicht, schreibt ihn in den High-Bandwidth Memory (HBM) und erst danach liest der Optimizer ihn zurück, um die Gewichte zu aktualisieren. Dieser zweiphasige Zeitplan erzeugt einen „Gradienten-Boden“, bei dem die Gradienten für alle Schichten an der Grenze zwischen dem Backward-Pass und dem Optimizer-Schritt gleichzeitig im Speicher existieren müssen. Für ein 8B-Parameter-Modell verbraucht dieser transiente Gradienten-Pool allein etwa 16 GiB und entscheidet oft darüber, ob eine Trainingskonfiguration auf eine einzelne GPU passt. Bestehende Methoden, die Backward- und Optimizer-Schritte fusionieren (z. B. LOMO, FlashOptim), operieren auf der Granularität ganzer Parameter oder Schichten, was bedeutet, dass der vollständige Gradiententensor für einen Parameter dennoch vor dem Update im Speicher materialisiert werden muss.

2. Methodik: FORGE

FORGE (Fused On-Register Gradient Elimination) eliminiert den materialisierten Gradiententensor vollständig, indem es den Optimizer-Schritt direkt in die Gewichtsgradienten-Berechnung auf Tile-Granularität (speziell $128 \times 128$ Tiles) fusioniert.

Kernmechanismus

Anstatt die vollständige Gradientenmatrix zu berechnen und zu speichern, verarbeitet FORGE die Gewichtstabelle Tile für Tile:

Tile-weise Gradienten-Akkumulation: Für ein spezifisches Gewicht-Tile akkumuliert der Kernel die Teilprodukte des Upstream-Gradienten ( $\Delta Y$ ) und der Input-Aktivierung ( $X$ ) direkt in FP32-GPU-Registern.
In-Register-Optimierung: Sobald der Gradient des Tiles in den Registern vollständig akkumuliert ist, wird das Optimizer-Update (z. B. AdamW) sofort auf dieses Tile angewendet, ebenfalls innerhalb der Register.
Sofortiges Verwerfen: Das Gradient-Tile wird in dem Moment, in dem es erzeugt wird, konsumiert. Es wird niemals als Tensor in den HBM geschrieben.
Write-Back: Nur das aktualisierte Gewicht-Tile und die (optional quantisierten) Optimizer-State-Tiles (Momente) werden zurück in den Global Memory geschrieben.

Präzision und Treue (Fidelity)

Exaktheit: Für jeden elementweise-separablen Optimizer (bei dem das Update für ein Gewicht nur von dessen Koordinate abhängt, z. B. AdamW, SGD, Lion), ist das pro-Tile-Update mathematisch identisch mit einem Standard-Vollgradienten-Schritt. Die einzige Abweichung ist die Reihenfolge der Gleitkomma-Summation über die Token, was inhärent zu jeder GEMM-Implementierung ist und keine Approximation der Optimizer-Logik darstellt.
Low-Precision-Treue: In BF16- oder INT8-Trainingsregimen bewahrt FORGE die Full-Precision-Treue auf zwei Arten:
- Der Gradient wird vor der Optimierung nicht auf BF16 gekürzt; er wird in FP32-Registern berechnet und konsumiert.
- Um die „Dead Zone“ von BF16 (in der kleine Updates auf Null gerundet werden) zu verhindern, nutzt FORGE stochastisches Rounding beim Zurückschreiben der Gewichte, was eine unverzerrte Akkumulation von Sub-ULP-Updates gewährleistet.
- Für INT8-Optimizer-Zustände verwendet FORGE kleinere Blockgrößen (64 statt 2048 bei Standard-8-Bit-Optimierern), um Quantisierungsfehler zu minimieren.

Kompatibilität mit verteinfertem Training

FORGE ist unter Tensor Parallelism (TP) und Sequence Parallelism (SP) nachweislich exakt, da der lokale Gradient für einen Shard zum Zeitpunkt des Backward-Passes vollständig ist. Es ist jedoch in seiner jetzigen Form inkompatibel mit Data Parallelism (DP) oder Context Parallelism, da diese ein All-Reduce der partiellen Gradienten erfordern, bevor das Update angewendet werden kann, was die Materialisierung des Gradiententensors notwendig machen würde.

3. Wichtigste Beiträge

Strukturelle Eliminierung, nicht Kompression: FORGE identifiziert den materialisierten Gradienten als Artefakt des Trainings-Schedulings statt als mathematische Notwendigkeit. Es eliminiert den $2P$ (Parameteranzahl) Gradienten-Speicherterm und reduziert den Gradienten-Pool von $O(P)$ auf $O(T_R T_C)$ (Tile-Größe), unabhängig von der Modellgröße.
Nachweisbare Exaktheit: Das Paper beweist, dass für elementweise Optimierer der Tile-fusionierte Zeitplan Updates liefert, die identisch mit einem Standard-AdamW auf FP32-Gradienten sind. Diese Exaktheit gilt unter Tensor- und Sequence-Parallelism.
Full-Precision-Treue bei niedriger Präzision: Durch die Beibehaltung der Gradienten in FP32-Registern bis zur Optimierung vermeidet FORGE die Verzerrung, die durch das Kürzen der Gradienten auf BF16 vor der Optimierung entsteht. Kombiniert mit stochastischem Rounding erreicht es die Konvergenz eines FP32-Master-Modells selbst bei Verwendung von BF16- oder INT8-Zuständen.
Generalität: Die Methode ist architekturagnostisch für lineare Schichten und unterstützt acht vollständig fusionierte elementweise Optimierer (AdamW, NAdam, RAdam, Lion, RMSprop, AdaGrad, SGD, SGD-momentum).

4. Experimentelle Ergebnisse

Evaluierung auf NVIDIA H200 und B200 GPUs über Modelle mit 0,6B bis 14B Parametern (Llama-3.1-8B, Qwen3-Familie):

Speicherreduktion: FORGE reduziert den Spitzenwert des Speichers um 38–53 % im Vergleich zu fusioniertem AdamW. Für Llama-3.1-8B sinkt der Spitzenwert des Speichers von ~75 GiB auf 35,4 GiB (eine Reduktion um 53 %).
Durchsatz: FORGE erreicht eine 1,52-fache Beschleunigung in der Schrittzeit (99,1 ms vs. 150,9 ms), indem es den separaten Optimizer-Kernel-Aufruf und den HBM-Round-Trip für Gradienten eliminiert.
Skalierbarkeit:
- Auf einer einzelnen H200 trainiert FORGE Qwen3-14B (was mit Standard-AdamW zu Out-of-Memory-Fehlern führt) mit 63,2 GiB.
- In einem 4-GPU A100-40GB Setup ermöglicht FORGE eine Micro-Batch-Größe von 4 (passt in 24,6 GiB), während Standard-AdamW nur Micro-Batch 1 zulässt. Dies erlaubt es FORGE, einen 2,42-mal höheren Durchsatz als die Standard-Baseline auf dersselben Hardware zu erreichen.
Konvergenz: FORGE folgt der Konvergenz von fusioniertem AdamW innerhalb von 0,001 nats (durchschnittliche Abweichung) bei Continued-Pretraining-Aufgaben, was bestätigt, dass die Low-Precision-Abweichungen innerhalb des Run-to-Run-Rauschens liegen.
Komponierbarkeit: FORGE lässt sich sauber mit FP8-Quantisierung und Activation Checkpointing kombinieren und absorbiert die Dequantisierungskosten, die typischerweise Standard-Optimierer benachteiligen.

5. Bedeutung und Ansprüche

Das Paper behauptet, dass FORGE eine strukturelle Verschiebung darstellt und keine inkrementelle Optimierung ist. Durch die Entfernung des Gradienten-Bodens wird der Optimizer-Speicheraufwand in nutzbaren Spielraum für größere Batches, längere Sequenzen oder größere Modelle auf bestehender Hardware umgewandelt.

Kapazitätsgrenze: Die Autoren stellen fest, dass die Eliminierung des Gradienten-Bodens nicht nur eine konstante Verbesserung, sondern eine Kapazitätsgrenze ist. Beispielsweise sind Standard-Trainingskonfigurationen für 14B-Modelle auf einer einzelnen H200 aufgrund des Gradienten-Bodens unmöglich, während FORGE sie machbar macht.
Orthogonalität: Die Methode ist orthogonal zu anderen speichereffizienten Techniken (wie State-Quantisierung oder Activation Recomputation) und kann mit diesen kombiniert werden, um den Speicherbedarf weiter zu reduzieren.
Limitierungen: Die Methode ist derzeit auf elementweise separable Optimierer beschränkt und unterstützt Data Parallelism ohne Modifikation nicht. Zudem erfordert sie, dass der Input-Gradient gelesen wird, bevor das Weight-Update erfolgt, was Tied-Embeddings oder MoE-Experten mit mehreren Gradienten-Konsumenten in einem einzigen Schritt ausschließt.

Zusammenfassend zeigt FORGE, dass der „Gradienten-Boden“ ein unnötiges Artefakt aktueller Trainingsschleifen ist. Durch die Fusion des Optimizers in den Backward-Pass auf der Tile-Ebene erreicht es signifikante Speicherersparnisse und Geschwindigkeitssteigerungen bei gleichzeitiger Wahrung der mathematischen Exaktheit für eine breite Klasse von Optimierern.