gpt-oss Reinforcement Learning

Sie können jetzt OpenAI trainieren gpt-oss mit RL und GRPO über Unsloth. Unsloth bietet jetzt das schnellste Inferenz (3x schneller), geringster VRAM-Verbrauch (50% weniger) und längster Kontext (8x länger) für gpt-oss RL im Vergleich zu jeder Implementierung – ohne Genauigkeitsverlust. Da Reinforcement Learning (RL) für gpt-oss noch nicht vLLM-kompatibel ist, mussten wir den Inferenzcode von Transformers umschreiben, um für gpt-oss eine 3x schnellere Inferenz von ~21 Tokens/s zu erreichen. Für BF16 erzielt Unsloth ebenfalls die schnellste Inferenz (~30 Tokens/s), insbesondere relativ zum VRAM-Verbrauch, und nutzt 50% weniger VRAM als jede andere RL-Implementierung. Wir planen, unsere 50% Gewichts-Sharing-Funktion zu unterstützen, sobald vLLM mit RL kompatibel wird.

• Kostenloses Notebook: gpt-oss-20b GRPO Colab-Notebook Dieses Notebook erstellt automatisch schnellere Matrix-Multiplikations-Kernel und verwendet 4 neue Unsloth-Belohnungsfunktionen. Wir zeigen außerdem, wie man Reward-Hacking entgegenwirkt was eine der größten Herausforderungen im RL ist.\

  

Mit Unsloth können Sie gpt-oss-20b mit GRPO auf 15GB VRAM trainieren und kostenlos auf Colab. Wir haben Embedding-Offloading eingeführt, das den Verbrauch zusätzlich um 1GB reduziert via offload_embeddings. Unsloths neue Inferenz läuft schneller auf jedes GPUs einschließlich A100, H100 und älteren T4s. gpt-oss-120b passt gut auf eine GPU mit 120GB VRAM.

Unsloth ist das einzige Framework, das 4-Bit-RL für gpt-oss unterstützt. Alle Leistungsgewinne sind auf Unsloths einzigartige Gewichtsteilung, Flex Attention, Standby und maßgeschneiderte Kernel zurückzuführen.

Erinnerung: Flash Attention 3 (FA3) ist ungeeignet für gpt-oss Training da es derzeit den Backward-Pass für Attention Sinks nicht unterstützt, was falsche Trainingsverlusteverursacht. Wenn Sie brauchen Unsloth verwenden, kann FA3 standardmäßig aktiviertsein, überprüfen Sie also bitte unbedingt, dass es nicht verwendet wird! Das Deaktivieren von FA3 führt allerdings auch zu O(N^2) Speichernutzung, weshalb Unsloth das einzige RL-Framework ist, das O(N) Speichernutzung für gpt-oss über unsere Flex-Attention-Implementierung anbietet.

⚡Inferenz viel schneller machen

Inferenz ist im RL-Training entscheidend, da wir sie benötigen, um Kandidatenlösungen zu generieren, bevor eine Belohnungsfunktion maximiert wird (siehe hier für eine detailliertere Erklärung). Um die schnellste Inferenzgeschwindigkeit für gpt-oss ohne vLLM zu erreichen, haben wir den Transformers-Inferenzcode neu geschrieben und viele Innovationen integriert, einschließlich benutzerdefinierter Algorithmen wie Unsloth Flex Attention, wobei wir spezielle Flags innerhalb von torch.compile (wie Kombi-Kernel) verwenden. Unser neuer Inferenzcode für gpt-oss wurde gegen eine bereits optimierte Basislinie bewertet (2x schneller als native Transformers).

vLLM unterstützt RL für gpt-oss nicht, da ihm BF16-Training und LoRA-Unterstützung für gpt-oss fehlen. Ohne Unsloth funktioniert Training nur in voller BF16-Präzision, wodurch die Speichernutzung um über 800% höherwird. Die meisten Frameworks aktivieren FA3 (Flash Attention 3) standardmäßig (was den VRAM-Verbrauch reduziert und die Geschwindigkeit erhöht) aber das verursacht falsche Trainingsverluste. Siehe Issue 1797 im FA3-Repo. Sie müssen FA3 deaktivieren, da es sonst das Training mit langem Kontext verhindert, weil FA3 O(N) Speicher verwendet, während naive Attention mit O(N^2) anwächst. Um Attention Sinks differenzierbar zu machen, haben wir Unsloth Flex Attention.

Wir haben die gpt-oss RL-Inferenz durch Benchmarking von BitsandBytes 4-Bit evaluiert und außerdem separate Tests für BF16 durchgeführt. Unsloths 4-Bit-Inferenz ist etwa 4x schneller, und BF16 ist ebenfalls effizienter, besonders beim VRAM-Verbrauch.

Das Beste an Unsloths gpt-oss RL ist, dass es auf jeder GPU funktionieren kann, sogar auf solchen, die BF16 nicht unterstützen. Unsere kostenlosen gpt-oss-20b Colab-Notebooks nutzen ältere 15GB T4-GPUs, daher funktionieren die Inferenzbeispiele gut!

🛠️ gpt-oss Flex Attention Probleme und Eigenheiten

Wir mussten unsere Implementierung für Attention Sinks ändern, wie hier beschrieben um die Generierung mit Left Padding zu ermöglichen. Wir mussten das LogSumExp berechnen und die Sigmoid-Aktivierung anwenden, um die Aufmerksamkeitsgewichte wie unten zu ändern:

$$A(X) = \sigma \bigg( \frac{1}{\sqrt{d}}QK^T \bigg)V \\ A(X) = \frac{\exp{\frac{1}{\sqrt{d}}QK^T}}{\sum{\exp{\frac{1}{\sqrt{d}}QK^T}}}V \\ \text{LSE} = \log{\sum{\exp{\frac{1}{\sqrt{d}}QK^T}}} \\ A_{sinks}(X) = A(X) \odot \sigma (\text{LSE} - \text{sinks})$$

Left-padded Masking während der Inferenz war ebenfalls ein kniffliges Problem in gpt-oss. Wir stellten fest, dass wir nicht nur die KV-Cache-Vorbefüllung während der Token-Generierung berücksichtigen mussten, sondern auch die unterschiedliche Anzahl von Pad-Tokens in jedem Prompt bei Batch-Generierungen, was die Art und Weise verändert hätte, wie wir die Block-Maske speichern müssen. Ein Beispiel dafür ist unten zu sehen:

Normale kausale Maske:

   k0 k1 k2 k3 k4   <-- Schlüssel
q0  X
q1  X  X
q2  X  X  X
q3  X  X  X  X
q4  X  X  X  X  X   <-- letzte Query-Zeile (am wichtigsten für das Decoding)

Für Inferenz im allgemeinen Fall (Decoding)

    k0 k1 k2 k3 k4
q0
q1
q2
q3
q4   X  X  X  X  X

Wenn wir naiv dieselbe Maskierungsstrategie verwenden, schlägt das fehl:

    k0 k1 k2 k3 k4
q0
q1
q2
q3
q4   X   (beachte, dass q4 den q_idx=0 hat, da dies die erste Query in der aktuellen Konfiguration ist)

Für die Generierung (Decoding-Phase) interessiert uns normalerweise nur die letzte Zeile der Attention-Matrix, da es nur ein Query-Token gibt, das auf alle vorherigen Key-Tokens achtet. Wenn wir naiv die kausale Maske anwenden (q_idx ≥ k_idx), versagt dies, da unsere einzelne Query den Index 0 hat, während es n_k Key-Tokens gibt. Um dies zu beheben, benötigen wir einen Offset bei der Maskenerstellung, um zu entscheiden, welche Tokens beachtet werden. Ein naiver Ansatz ist jedoch langsam, da sich Offsets bei jedem Schritt ändern und Masken- sowie Kernel-Neugenerierung erzwingen. Wir haben dies mit Cache- und Compile-Optimierungen gelöst.

Der schwierigere Teil ist die Batch-Generierung. Sequenzen unterscheiden sich in der Länge, sodass Padding die Maskenerstellung verkompliziert. Flex Attention hatte viele Herausforderungen und dynamische Masken sind knifflig. Schlimmer noch, wenn nicht kompiliert, fällt es auf eine eager Attention zurück, die langsam und speicherintensiv ist (quadratisch vs. linear in der Sequenzlänge).

Zitat von https://github.com/meta-pytorch/attention-gym/issues/15#issuecomment-2284148665

Sie müssen dies mit _compile=True aufrufen. Wir bilden im Wesentlichen Ihre Block-Maske über eine vollständige Q_LEN x KV_LEN-Matrix ab, um die Block-Maske zu erzeugen. Ohne Compile müssen wir dieses vollständige Gebilde materialisieren, und es kann bei langen Sequenzen zu OOMs führen.

Außerdem müssen Sie flex_attention = torch.compile(flex_attention)ausführen. Ohne Compile fällt flex auf eine nicht-fused eager-Implementierung zurück, die großartig zum Debuggen ist, aber viel langsamer ist und die vollständige Scores-Matrix materialisiert.

Letztendlich muss die Maske dynamisch Prefill vs. Decode mit dem KV-Cache, Batch- und Padding-Tokens pro Sequenz handhaben, torch.compile freundlich bleiben und gleitende Fenster unterstützen.

🔍 Untersuchung von Flash Attention

Eine weitere interessante Richtung, die wir untersucht haben, war der Versuch, Flash Attention zu integrieren. Ihre Vorteile sind weithin anerkannt, aber eine Einschränkung ist, dass sie während des Backward-Passes keine Attention Sinks für gpt-oss unterstützt. Um dies zu umgehen, haben wir den Attention-Mechanismus so umstrukturiert, dass er ausschließlich auf dem Attention-Output und den LogSumExp-Werten arbeitet, die FlashAttention bereitstellt. Angesichts dieser Vorteile schien es eine naheliegende Wahl zu sein, es zu versuchen.

Wir bemerkten jedoch bald Probleme. Während die ersten Schichten wie erwartet verhielten, produzierten die späteren Schichten, insbesondere Schichten 18 bis 24, Ausgaben, die deutlich von der Eager-Mode-Implementierung in transformers abwichen. Wichtig ist, dass diese Diskrepanz nicht auf Fehlerakkumulation zurückgeführt werden kann, da die Eingaben für jede Methode in jeder Schicht identisch sind. Zur weiteren Validierung verglichen wir die Ergebnisse auch mit Unsloth FlexAttention.

Dies erfordert weitere Untersuchungen, warum nur die letzten Schichten einen so drastischen Unterschied zwischen der Flash-Attention-Implementierung und den anderen zeigen.

Flash Attention 3 unterstützt den Backward-Pass für Attention Sinks nicht

FA3 ist bei den meisten Trainingspaketen (nicht Unsloth) oft standardmäßig aktiviert, aber das ist für gpt-oss inkorrekt. Die Verwendung von FA3 macht den Trainingsverlust völlig falsch, da FA3 die Backward-Pässe für Attention Sinks in gpt-oss nicht unterstützt. Viele Leute sind sich dessen noch nicht bewusst, daher bitte vorsichtig sein!

⚠️ Können wir Reward Hacking entgegenwirken?

Das ultimative Ziel von RL ist es, eine Belohnung (z. B. Geschwindigkeit, Umsatz, eine Metrik) zu maximieren. Aber RL kann betrügen. Wenn der RL-Algorithmus einen Trick lernt oder etwas ausnutzt, um die Belohnung zu erhöhen, ohne die Aufgabe tatsächlich zu erfüllen, nennt man das "Belohnungsmanipulation".

Es ist der Grund, warum Modelle lernen, Unit-Tests zu verändern, um Programmieraufgaben zu bestehen, und dies sind kritische Blocker für den Einsatz in der realen Welt. Einige weitere gute Beispiele stammen aus Wikipedia.

In unserem kostenlosen gpt-oss RL-Notebook untersuchen wir, wie man Belohnungsmanipulation in einem Code-Generierungs-Setting entgegenwirken kann, und zeigen greifbare Lösungen für häufige Fehlermodi. Wir sahen, wie das Modell die Zeitmessfunktion bearbeitete, Aufgaben an andere Bibliotheken auslagerte, Ergebnisse zwischenspeicherte und offen betrog. Nach Gegenmaßnahmen erzeugt unser Modell wirklich optimierte Matrixmultiplikations-Kernel, keine cleveren Tricks.

🏆Belohnungsmanipulation

Einige häufige Beispiele für Belohnungsmanipulation während RL sind:

Faulheit

RL lernt, Numpy, Torch und andere Bibliotheken zu verwenden, die optimierte CUDA-Kernel aufrufen. Wir können den RL-Algorithmus daran hindern, optimierten Code aufzurufen, indem wir prüfen, ob der generierte Code andere nicht standardmäßige Python-Bibliotheken importiert.

Zwischenspeicherung & Betrug

RL lernt, das Ergebnis der Ausgabe zwischenzuspeichern, und RL lernt, die tatsächliche Ausgabe zu finden, indem es Python-Globale überprüft.

Wir können den RL-Algorithmus daran hindern, zwischengespeicherte Daten zu verwenden, indem wir den Cache mit einer großen falschen Matrix löschen. Wir müssen auch sorgfältig mit mehreren Schleifen und Durchläufen benchmarken.

Betrug

RL lernt, die Zeitmessfunktion zu bearbeiten, damit sie 0 Zeit als bestanden ausgibt. Wir können den RL-Algorithmus daran hindern, globale oder zwischengespeicherte Variablen zu verwenden, indem wir seine lokalen Variablen und globalen Variablen. Wir werden außerdem exec verwenden, um die Funktion zu erstellen, daher müssen wir die Ausgabe in einem leeren Dict speichern. Wir untersagen außerdem den Zugriff auf globale Variablen über types.FunctionType(f.__code__, {})\

Tutorial: Wie man gpt-oss mit RL trainiert

LLMs haben oft Schwierigkeiten mit Aufgaben, die komplexe Umgebungen beinhalten. Durch die Anwendung von Verstärkendem Lernen (RL) und dem Entwurf einer benutzerdefinierten Belohnungsfunktion, können diese Herausforderungen jedoch überwunden werden.

RL kann für Aufgaben wie automatische Kernel- oder Strategiegenerierung angepasst werden. Dieses Tutorial zeigt, wie man gpt-oss mit GRPO und Unsloth trainiert, um autonom 2048 zu schlagen.

Unsere Notebooks enthalten bereits Schritt-für-Schritt-Anleitungen, wie man den gesamten Prozess durchführt.

2048-Notebook (Offizielles OpenAI-Beispiel)	Kernel-Generierungs-Notebook

Was du bauen wirst:

• Trainiere gpt-oss-20b, sodass das Modell 2048 automatisch gewinnen kann

• Erstelle eine minimale 2048-Umgebung, mit der das Modell interagieren kann

• Definiere Belohnungsfunktionen die:

  1. Überprüfe, ob die generierte Strategie kompiliert und ausgeführt wird,

  2. Verhindere Belohnungsmanipulation (externen Imports nicht erlauben), und

  3. Belohne tatsächlichen Spiel-Erfolg

• Führe Inferenz aus und exportiere das Modell (MXFP4 4‑Bit oder zusammengeführtes FP16)

Hardware: Das 2048-Beispiel läuft auf einem kostenlosen Colab T4, aber das Training wird langsam sein. A100/H100 ist viel schneller. 4-Bit-Loading + LoRA ermöglicht es Ihnen, ein 20B-Modell in moderatem VRAM unterzubringen