🤖 machine learning

Accelerating RL Post-Training Rollouts via System-Integrated Speculative Decoding

Dieses Papier stellt einen systemintegrierten Rahmen für spekulatives Decodieren innerhalb von NeMo-RL vor, der als verlustlose Beschleunigungsprimitive für Rollouts nach dem RL-Training dient und bei einer Skalierung von 8 Milliarden Parametern eine 1,8-fache Steigerung des Durchsatzes sowie bei einer Skalierung von 235 Milliarden Parametern in Kombination mit asynchroner Ausführung eine End-to-End-Beschleunigung des Trainings von bis zu 2,5-fach erzielt.

Ursprüngliche Autoren: Hayate Iso, Tiyasa Mitra, Sudipta Mondal, Rasoul Shafipour, Venmugil Elango, Terry Kong, Yuki Huang, Seonjin Na, Izzy Putterman, Benjamin Chislett, Maor Ashkenazi, Joseph Guman, Gerald Shen, Tugrul Ko

Veröffentlicht 2026-04-30

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Hayate Iso, Tiyasa Mitra, Sudipta Mondal, Rasoul Shafipour, Venmugil Elango, Terry Kong, Yuki Huang, Seonjin Na, Izzy Putterman, Benjamin Chislett, Maor Ashkenazi, Joseph Guman, Gerald Shen, Tugrul Konuk, Ashwath Aithal, Ritika Borkar, Ran Zilberstein, Bita Rouhani

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

Stellen Sie sich vor, Sie trainieren einen brillanten, aber langsam denkenden Schüler (das KI-Modell), um komplexe Mathematikaufgaben zu lösen. Um ihn zu unterrichten, müssen Sie ihm eine Frage stellen, warten, bis er seinen gesamten Denkprozess schrittweise aufschreibt, und dann prüfen, ob er es richtig gelöst hat.

Das Problem ist, dass der Teil des „Prüfens" schnell ist, der Teil des „Schreibens" jedoch unglaublich langsam. Der Schüler schreibt ein Wort, macht eine Denkpause, schreibt das nächste Wort, macht erneut eine Pause und so weiter. Dieses langsame, wortweise Schreiben ist die größte Engstelle beim Training dieser KI-Modelle.

Diese Arbeit stellt einen cleveren Trick vor, der Spekulatives Decodieren (Speculative Decoding) genannt wird, um diesen Schreibprozess zu beschleunigen, ohne zu verändern, wie der Schüler denkt oder was er lernt.

Die Kernidee: Der „Entwurfs-Assistent"

Stellen Sie sich das KI-Modell als einen sehr präzisen, aber langsamen Küchenchef vor. Um die Dinge zu beschleunigen, stellen Sie einen schnellen, energiegeladenen Sous-Chef ein (das „Entwurfs-Modell").

Der alte Weg (Autoregressiv): Der Küchenchef schreibt ein Wort, hält inne, denkt nach, schreibt das nächste, hält inne, denkt nach. Das dauert ewig.
Der neue Weg (Spekulatives Decodieren): Der Sous-Chef errät schnell die nächsten 3 oder 4 Wörter, die der Chef schreiben könnte. Der Sous-Chef schreibt sie schnell auf.
Die Prüfung: Der Küchenchef wirft dann einen schnellen Blick auf die Notizen des Sous-Chefs.
- Wenn die Notizen richtig sind, sagt der Chef: „Toll!" und akzeptiert alle 4 Wörter auf einmal.
- Wenn die Notizen falsch sind, streicht der Chef sie durch, schreibt das korrekte Wort und beginnt von vorne.

Die Magie: Da der Sous-Chef schnell ist, kann der Chef die meisten der Zeit mehrere Wörter auf einmal akzeptieren. Das Endergebnis ist exakt dasselbe, als hätte der Chef allein geschrieben, aber es geschieht viel schneller.

Was die Arbeit tatsächlich getan hat

Die Forscher haben dieses System in ein reales Trainingsframework namens NeMo-RL integriert. Sie haben es nicht nur an einfachen Aufgaben getestet, sondern an „Denkaufgaben", bei denen die KI hart nachdenken muss (wie beim Lösen von Mathematikaufgaben).

Hier sind ihre wichtigsten Erkenntnisse, in alltägliche Sprache übersetzt:

Es funktioniert ohne Betrug: Einige Beschleunigungsmethoden versuchen, Abkürzungen zu nehmen (wie die Verwendung von Mathematik geringerer Qualität oder das Überspringen von Schritten), was das Lernen des Schülers ruinieren kann. Diese Methode ist „verlustfrei". Sie garantiert, dass die KI exakt so lernt, wie sie es ohne den Assistenten getan hätte, nur schneller.
Der Geschwindigkeitsschub:
- Bei einem mittelgroßen Modell (8 Milliarden Parameter) sahen sie eine Beschleunigung des Schreibprozesses um das 1,5- bis 1,8-fache.
- Da das Schreiben etwa 70 % der gesamten Trainingszeit in Anspruch nimmt, wurde der gesamte Trainingsprozess etwa 1,35- bis 1,4-mal schneller.
Der „Entwurf" ist wichtig: Der Sous-Chef muss ein guter Rater sein.
- Wenn Sie den Sous-Chef auf exakt dieselbe Art von Mathematikaufgaben trainieren, die der Schüler lernt, ist der Geschwindigkeitsschub enorm.
- Wenn Sie einen allgemeinen Sous-Chef verwenden, der nur allgemeine Gespräche kennt, ist der Geschwindigkeitsschub geringer.
- Nicht zu weit voraus raten: Wenn der Sous-Chef versucht, 7 Wörter auf einmal zu raten, macht er zu viele Fehler, und der Küchenchef verbringt zu viel Zeit mit der Korrektur. Das Raten von 3 Wörtern auf einmal war der „Sweet Spot".
Die Zukunft (Große Modelle): Sie verwendeten einen hochpräzisen Computersimulator, um vorherzusagen, was bei massiven Modellen (235 Milliarden Parameter) passiert, wenn Tausende von Computern zusammenarbeiten.
- Sie prognostizieren, dass für diese riesigen Modelle diese Technik den gesamten Trainingsprozess um das 2,5-fache beschleunigen könnte.

Warum das wichtig ist

In der Welt der KI ist Zeit Geld. Wenn Sie ein Modell 2,5-mal schneller trainieren können, können Sie entweder in der gleichen Zeit ein intelligenteres Modell erhalten oder dasselbe Modell für einen Bruchteil der Kosten bekommen.

Die Arbeit beweist, dass Sie das „Gehirn" der KI oder die Spielregeln nicht ändern müssen, um diese Geschwindigkeit zu erreichen. Sie müssen lediglich einen intelligenten, schnellen Assistenten hinzufügen, der beim Entwerfen hilft, und das Hauptmodell die endgültige Verifizierung durchführen lassen. Es ist ein System-Upgrade, das die gesamte Trainingspipeline reibungsloser laufen lässt, ohne etwas zu zerstören.

1. Problemstellung

Das Nachtrainieren von Reinforcement Learning (RL) für fortschrittliche Large Language Models (LLMs) wird zunehmend durch die autoregressive Rollout-Generierung ausgebremst. Bei rechenintensiven Workloads (z. B. Mathematik, Code) dominiert die Zeit, die für die Generierung von Trajektorien (Rollouts) aufgewendet wird, die gesamte Trainingszeit und übersteigt oft die Zeit für die Gradientenberechnung.

Bestehende Effizienzmethoden versuchen, dieses Problem zu lösen, indem sie die Trainingsdynamik verändern, wie zum Beispiel:

Asynchrone Ausführung: Überlappt die Generierung mit dem Lernen, führt jedoch zu einer Verzögerung der Policy.
Off-Policy-Replay: Wiederverwendung veralteter Trajektorien, erfordert jedoch Korrekturen durch Importance Sampling.
Rollouts mit niedriger Präzision: Reduziert die Rechenlast, führt jedoch zu einer Diskrepanz in der Verteilung.

Diese Methoden opfern die Wirksamkeit (die Qualität des Lernsignals) zugunsten des Durchsatzes. Das Papier identifiziert einen kritischen Bedarf an einer Methode, die die Generierung beschleunigt, ohne die Stichprobenverteilung der Ziel-Policy zu verändern, und damit die exakten Trainingssemantiken erhält, die für eine stabile RL-Konvergenz erforderlich sind.

2. Methodik

Die Autoren schlagen vor, Speculative Decoding als verlustlose Beschleunigungsprimitive direkt in das NeMo-RL-Trainingsframework (unter Verwendung eines vLLM-Backends) zu integrieren.

Kernmechanismus

Drafting: Ein „Draft"-Modell (oder -Kopf) schlägt mehrere Tokens ( $k$ ) parallel vor.
Verifikation: Das Ziel-„Verifier"-Modell (die RL-Policy) verifiziert diese Tokens in einem einzigen Vorwärtsdurchlauf.
Ablehnungsstichprobe (Rejection Sampling): Tokens werden basierend auf der Wahrscheinlichkeitsverteilung des Verifiers akzeptiert oder abgelehnt. Dies garantiert, dass die finale Rollout-Verteilung identisch mit der Ziel-Policy bleibt und somit keine Änderung des RL-Optimierungsziels erfolgt.

Systemintegration

Die Integration geht über den Standard-Inferenzdienst hinaus, indem sie spezifische RL-Einschränkungen adressiert:

Gewichts-Synchronisation: Das Draft-Modell muss mit den sich schnell ändernden Policy-Gewichten während des RL-Trainings synchronisiert bleiben.
Zwei Drafting-Pfade:
- Allgemeiner Pfad (EAGLE-3): Verwendet ein externes Draft-Modell, das auf den Ausgaben der Policy trainiert wurde. Dies funktioniert für jedes vortrainierte Modell.
- Nativer Pfad: Verwendet integrierte Multi-Token-Prediction (MTP)-Köpfe, falls die Modellarchitektur diese unterstützt.
Gradienten-Trennung (Gradient Detachment): Um zu verhindern, dass das Training des Draft-Modells das Gradientensignal der Policy beeinträchtigt, werden versteckte Zustände und Log-Wahrscheinlichkeiten während des Vorwärtsdurchlaufs der Policy zwischengespeichert und über eine detach()-Grenze zum Draft-Kopf geleitet.
Pipeline-Kompatibilität: Das System unterstützt sowohl synchrone (Generierung wartet auf Training) als auch asynchrone (Generierung überlappt mit Training) RL-Pipelines.

3. Hauptbeiträge

Systemintegration: Erste Implementierung von Speculative Decoding innerhalb eines produktionsreifen, Open-Source-RL-Trainingsstapels (NeMo-RL), der verifier-exakte Trainingssemantiken bewahrt.
Operative Analyse: Eine umfassende Studie der Faktoren, die die Beschleunigung bestimmen, einschließlich Draft-Initialisierung, Draft-Länge und Online-Anpassung.
Skalierbarkeitsprojektionen: Ein hochpräziser Leistungssimulator, der verwendet wird, um Gewinne im Bereich der Frontier-Skalen (bis zu 235 Milliarden Parameter und 2048 GPUs) zu projizieren und damit über die Grenzen aktueller 8-Milliarden-Experiment-Setups hinauszugehen.

4. Experimentelle Ergebnisse

Experimentelles Setup

Modelle: Qwen3-8B (RL-Think) und Qwen3-8B-Base (RL-Zero).
Aufgaben: Mathematisches Reasoning (DAPO-Math-17K-Training, AIME-2024-Validierung).
Hardware: 8 GB200 NVL72-Knoten (insgesamt 32 GPUs).
Basislinie: Standard-autoregressive Decodierung.

Wichtige Erkenntnisse

Durchsatzgewinne:
- Generierungs-Latenz: Speculative Decoding reduzierte die Generierungszeit um das 1,54-fache (RL-Think) bzw. 1,77-fache (RL-Zero).
- End-to-End-Schrittzeit: Die Gesamtbeschleunigung des Trainingsschritts erreichte 1,35-fach bzw. 1,41-fach. Die Gewinne wurden durch nicht-Generierungsphasen (Berechnung von Log-Wahrscheinlichkeiten und Training) begrenzt, die durch Speculative Decoding nicht beschleunigt werden können.
- Genauigkeit: Die Validierungsgenauigkeitskurven waren zwischen autoregressiver und speculative Decodierung nicht unterscheidbar, was die verlustlose Natur der Beschleunigung bestätigt.
Operative Erkenntnisse:
- Initialisierung: Die Initialisierung des Draft-Modells mit domänenspezifischen Daten (DAPO-Math) schnitt deutlich besser ab als eine generische Chat-Domain-Initialisierung (UltraChat) und führte zu längeren Akzeptanzlängen und höheren Beschleunigungen.
- Draft-Länge ( $k$ ): Entgegen der Intuition erzielten kürzere Drafts ( $k=3$ ) die beste End-to-End-Beschleunigung. Längere Drafts ( $k=5, 7$ ) erhöhten den Verifikationsaufwand und machten den Prozess manchmal sogar langsamer als die autoregressive Decodierung.
- Online-Anpassung: Die Online-Aktualisierung des Draft-Modells brachte nur marginale Gewinne, wenn der initiale Draft schlecht abgestimmt war. Für gut initialisierte Drafts war ein Offline-Training ausreichend.
- Asynchrones RL: In asynchronen Modi bleibt Speculative Decoding komplementär. Während asynchrones Überlappen einige Generierungskosten verbirgt, reduziert die Spekulation zusätzlich die „exponierte" Generierungszeit auf dem kritischen Pfad und erzielt in spezifischen asynchronen Konfigurationen eine Beschleunigung von 1,24-fach.

Simulationsprojektionen (Einsatzskala)

Unter Verwendung eines Simulators für ein 235-Milliarden-Parameter-Modell auf 2048 GPUs:

Rollout-Beschleunigung: Kann unter optimalen Bedingungen (spezifische Draft-/Akzeptanzlängen und moderate Policy-Verzögerung) 3,5-fach überschreiten.
End-to-End-Beschleunigung: Wird auf ~2,5-fach projiziert, bedingt durch den hohen Anteil der Zeit, die bei Frontier-Skalen-Modellen für die Generierung aufgewendet wird.
Skalierungsempfindlichkeit: Größere Modelle profitieren in großskaligen Bereitstellungen stärker von der Spekulation, da sie mehr unter Latenzen mit langer Schwanzverteilung und Unterauslastung leiden, die durch Spekulation gemildert werden.

5. Bedeutung

Dieses Papier zeigt, dass Speculative Decoding eine viable, verlustlose Strategie zur Beschleunigung des RL-Nachtrainings ist. Seine primäre Bedeutung liegt in:

Bewahrung der Semantik: Im Gegensatz zu anderen Beschleunigungstechniken kompromittiert es das RL-Trainingsignal nicht und erfordert keine Off-Policy-Korrekturen.
Systemweite Optimierung: Es schließt die Lücke zwischen Inferenzbeschleunigungstechniken und den komplexen Anforderungen von RL-Trainingsloops (Gewichtssynchronisation, Gradientenbehandlung).
Skalierbarkeit: Die projizierte 2,5-fache End-to-End-Beschleunigung im Frontier-Bereich deutet darauf hin, dass Speculative Decoding zu einem Standardbestandteil zukünftiger RL-Infrastrukturen werden könnte, was die Zeit und Kosten für das Training von Reasoning-fähigen Modellen erheblich reduziert.

Die Arbeit legt einen klaren Einsatzpfad fest: Verwenden Sie domänenspezifisch initialisierte Drafts mit moderaten Draft-Längen (z. B. $k=3$ ), um den Durchsatz zu maximieren, ohne die Trainingsstabilität zu beeinträchtigen.