Optimierung von Reinforcement Learning durch Jet-RL und einheitliche FP8-Präzision

Optimierung des Reinforcement Learnings: Jet-RL und der einheitliche FP8-Präzisionsfluss

Reinforcement Learning (RL) hat sich als eine Schlüsseltechnologie zur Steigerung der komplexen Denkfähigkeiten großer Sprachmodelle (LLMs) etabliert. Insbesondere bei der Generierung von "Chain-of-Thought" (CoT) ermöglicht RL den Modellen, detaillierte Analysen durchzuführen und strukturierte logische Schlussfolgerungen zu ziehen. Trotz seiner Bedeutung ist das RL-Training jedoch bekanntermaßen ressourcenintensiv. Ein aktuelles Forschungspapier, entwickelt von einem Team unter Beteiligung von Haocheng Xi, Charlie Ruan und Peiyuan Liao, beleuchtet diese Herausforderungen und stellt einen innovativen Ansatz zur Effizienzsteigerung vor.

Der Engpass der Rollout-Phase im RL-Training

Die Rollout-Phase, in der der Akteur des LLM autoregressiv Antworten generiert, stellt einen erheblichen Engpass im RL-Trainingsprozess dar. Wie Analysen zeigen, kann diese Phase über 70 % der gesamten Trainingszeit beanspruchen, insbesondere wenn lange Token-Sequenzen erforderlich sind. Dies ist eine direkte Folge der autoregressiven Natur von LLMs und der Notwendigkeit, umfangreiche Lösungspfade zu erkunden. Die Beschleunigung des Rollouts ist daher von höchster Priorität für die Effizienz des gesamten Trainingsprozesses.

Ein vielversprechender Ansatz zur Effizienzsteigerung ist die FP8-Quantisierung, die bereits erfolgreich in der LLM-Inferenz eingesetzt wird. Eine gängige Strategie bestand darin, FP8-Präzision während der Rollout-Phase zu nutzen, während die Trainingsphase weiterhin in BF16-Präzision durchgeführt wurde. Diese Methode, oft als "BF16-Train-FP8-Rollout" bezeichnet, wurde in modernen RL-Frameworks wie VeRL, SLIME, Nemo-RL und OpenRLHF implementiert.

Grenzen der BF16-Train-FP8-Rollout-Strategie

Die umfassende Untersuchung des Forschungsteams zeigt jedoch, dass die BF16-Train-FP8-Rollout-Strategie erhebliche Einschränkungen aufweist. Sie leidet unter schwerwiegender Trainingsinstabilität und einem katastrophalen Genauigkeitsverlust, insbesondere bei langen Rollout-Sequenzen und anspruchsvollen Aufgaben. Die Analyse legt dar, dass diese Probleme aus einem numerischen Missverhältnis zwischen Training und Inferenz resultieren, das durch den Off-Policy-Charakter des Ansatzes entsteht.

Kleinere numerische Diskrepanzen akkumulieren sich während langer CoT-Begründungen, was zu einer zunehmenden Divergenz zwischen Rollout- und Trainingsverteilungen führt. Dies verstärkt die Off-Policy-Problematik im Reinforcement Learning unter FP8-Rollouts und führt zu Instabilität. Darüber hinaus erweist sich die Methode als fragil bei anspruchsvollen Aufgaben oder mit schwächeren Basismodellen, bei denen quantisierungsbedingte Fehler die Rollout-Trajektorie erheblich verzerren können.

Jet-RL: Ein einheitlicher FP8-Präzisionsfluss

Um diese Einschränkungen zu überwinden und FP8-Rollouts zu einer zuverlässigen Beschleunigungsstrategie zu machen, wurde das Framework Jet-RL entwickelt. Die Kernidee von Jet-RL ist die Durchsetzung eines truly on-policy FP8-Trainingsparadigmas, das das RL-Training stabilisiert. Dies wird durch die Verwendung eines identischen Quantisierungs-Präzisionsflusses für sowohl Training als auch Inferenz erreicht, wodurch das Policy-Missverhältnis eliminiert und die Notwendigkeit ineffizienter Zwischenschritt-Kalibrierung umgangen wird.

Jet-RL modelliert die Ausbreitung der Quantisierungspräzision als gerichteten Graphen, in dem Knoten Operatoren oder Gewichte darstellen und Kanten die Tensorausbreitung mit zugehöriger Präzision und Granularität beschreiben. Im Gegensatz zum BF16-Train-FP8-Rollout-Ansatz, bei dem der Inferenzgraph (FP8-Gewichte) vom Trainings-Forward-Graph (BF16-Gewichte und -Aktivierungen) abweicht, stellt Jet-RL sicher, dass der Inferenzgraph ein Subgraph des Trainings-Forward-Graphen ist. Alle relevanten Berechnungen, einschließlich Aktivierungen und Gewichte, die in lineare Schichten eingehen, werden in FP8-Präzision durchgeführt. Eine Masterkopie der Gewichte in höherer Präzision (BF16) wird beibehalten, um das Training zu stabilisieren.

Detaillierte Quantisierungsstrategie

Jet-RL verwendet ein gemischtes Per-Group- und Per-Block-Quantisierungsschema und nutzt fortschrittliche FP8-GEMM-Kernel, um die Beschleunigung des End-to-End-RL-Trainings zu ermöglichen:

• Gewichte: Werden mit Per-Block-Quantisierung von 128x128 quantisiert. Die Gewichtungsquantisierung für die Inferenz erfolgt während der Parameteraktualisierungsphase, um den Overhead zu minimieren.
• Aktivierungen und Gradienten: Werden mit Per-Group-Quantisierung von 1x128 quantisiert.
• FProp (Forward Pass): Die Eingabeaktivierung wird mit 1x128 Per-Group-Quantisierung und das Gewicht mit 128x128 Per-Block-Quantisierung quantisiert. Beide werden zeilenweise gespeichert, um den FP8-TensorCore-Kernel-Anforderungen gerecht zu werden.
• DGrad (Backward Pass): Ähnlich wie FProp, involviert es eine 1x128 quantisierte Matrix multipliziert mit einer 128x128 quantisierten Matrix.
• WGrad (Backward Pass): Die erste Matrix wird 1x128 quantisiert, die zweite 128x1 quantisiert.

Die Gradienten, die während des Backward-Passes zwischen Operatoren transportiert werden, bleiben in BF16-Präzision, um die Modellgenauigkeit zu erhalten und Probleme wie Gradienten-Underflow oder Quantisierungsrauschen zu vermeiden. Die GEMM-Operationen im Backward-Pass (DGrad und WGrad) werden jedoch ebenfalls in FP8-Präzision quantisiert, um eine Beschleunigung zu erzielen.

Bewertung der Effektivität von Jet-RL

Umfassende Experimente mit verschiedenen Modellen (Llama3.1-8B, Qwen2.5-7B, Qwen3-8B-Base), Datensätzen (GSM8K, MATH, DeepMATH) und Rollout-Konfigurationen (8K, 16K Token-Länge) validieren die Wirksamkeit von Jet-RL:

• Genauigkeit: Die BF16-Train-FP8-Rollout-Methode zeigte durchweg erhebliche Instabilität und konnte in einigen Szenarien nicht konvergieren oder führte zu erheblichen Leistungseinbußen (z.B. durchschnittlich 9,8 % bei Llama3.1-8B). Jet-RL hingegen konvergierte in allen Szenarien robust und reduzierte den Genauigkeitsverlust im Vergleich zur BF16-Baseline auf typischerweise 1–3 %. In einigen Fällen übertraf Jet-RL sogar die BF16-Baseline.
• Effizienz: Jet-RL erzielte signifikante Beschleunigungen:
  • Rollout-Phase: 1,07- bis 1,33-fache Beschleunigung gegenüber BF16 in Tokens/s. Größere Modelle (z.B. 32B) zeigten dabei die größten Beschleunigungen.
  • Trainingsphase (8B-Modell): 1,41-fache Beschleunigung, zurückzuführen auf 1,54-fache Beschleunigung bei der Akteur-Aktualisierung und 1,80-fache Beschleunigung bei der Referenzmodell-Inferenz.
  • End-to-End: Eine Gesamtbeschleunigung der Schrittzeit um das 1,16-fache für das 8B-Modelltraining.

Diese Ergebnisse bestätigen, dass Jet-RL eine robuste Lösung für effizientes, niedrig-präzises RL-Training bietet. Es ermöglicht eine signifikante Beschleunigung, ohne die Modellgenauigkeit zu beeinträchtigen, und ebnet den Weg für die Anwendung von FP8-Berechnungen zur Beschleunigung des großskaligen RL-Trainings.

Zukünftige Perspektiven

Die Forschung identifiziert weiterhin Bereiche für zukünftige Untersuchungen, darunter eine tiefere Analyse der Ursachen von Train-Inferenz-Inkonsistenzen, die Entwicklung von Quantisierungsstrategien mit geringeren Fehlern und die Verbesserung der Effizienz von niedrig-präzisem Training durch bessere Algorithmen, Frameworks und Hardware-Software-Co-Design. Diese Fortschritte sind entscheidend, um das volle Potenzial von LLMs für immer komplexere Aufgaben auszuschöpfen.

Jet-RL stellt einen bedeutenden Schritt dar, um die Leistungsfähigkeit des Reinforcement Learnings für große Sprachmodelle durch eine präzise und effiziente Handhabung numerischer Präzision zu maximieren. Die Fähigkeit, Training und Rollout in einem einheitlichen FP8-Fluss zu integrieren, überwindet bisherige Stabilitätsprobleme und eröffnet neue Möglichkeiten für die Skalierung und Anwendung von LLMs in anspruchsvollen Szenarien.

Bibliography: - Xi, H., Ruan, C., Liao, P., Lin, Y., Cai, H., Zhao, Y., Yang, S., Keutzer, K., Han, S., & Zhu, L. (2026). Jet-RL: Enabling On-Policy FP8 Reinforcement Learning with Unified Training and Rollout Precision Flow. arXiv. - T., R. (2026, January 23). Jet-RL Achieves 41% Faster FP8 Reinforcement Learning With Unified Precision Flow. Quantum Zeitgeist. - LMSYS Org. (2025, November 25). Unified FP8: Moving Beyond Mixed Precision for Stable and Accelerated MoE RL. - AI Papers Slop. (2026, January 24). Jet-RL: Enabling On-Policy FP8 Reinforcement Learning with Unified Training and Rollout Precision Flow [Video]. YouTube. - Hugging Face. (2026, January 26). Daily Papers. - Corca, Inc. (2026). [Literature Review] Jet-RL: Enabling On-Policy FP8 Reinforcement Learning with Unified Training and Rollout Precision Flow. TheMoonlight.io. - Qi, P., Liu, Z., Zhou, X., Pang, T., Du, C., Lee, W. S., & Lin, M. (2025). Defeating the Training-Inference Mismatch via FP16. arXiv.