gpt-oss Reinforcement Learning

Sie können jetzt OpenAI gpt-oss mit RL und GRPO über Unsloth. Unsloth bietet jetzt das schnellste Inferenz- (3x schneller), geringste VRAM-Nutzung (50 % weniger) und längste Kontext- (8x länger) für gpt-oss RL im Vergleich zu jeder Implementierung – ohne Genauigkeitsverlust. Da Reinforcement Learning (RL) auf gpt-oss noch nicht mit vLLM kompatibel ist, mussten wir den Inferenzcode von Transformers-Code neu schreiben, um eine 3x schnellere Inferenz für gpt-oss mit ~21 Token/s zu liefern. Für BF16 erreicht Unsloth ebenfalls die schnellste Inferenz (~30 Token/s), insbesondere im Verhältnis zur VRAM-Nutzung, mit 50 % weniger VRAM als jede andere RL-Implementierung. Wir planen, unsere 50 %-Gewichtsteilungsfunktion sobald vLLM mit RL kompatibel wird, zu unterstützen.

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

  

Mit Unsloth können Sie gpt-oss-20b mit GRPO auf 15 GB VRAM und kostenlos in Colab trainieren. Wir haben Embedding-Offloading eingeführt, wodurch der Verbrauch ebenfalls um 1 GB reduziert wird, über offload_embeddings. Unsloths neue Inferenz läuft schneller auf jede GPU, einschließlich A100, H100 und alten T4s. gpt-oss-120b passt problemlos auf eine GPU mit 120 GB VRAM.

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

Hinweis: Flash Attention 3 (FA3) ist für das gpt-oss- Training ungeeignet da es derzeit den Backward-Pass für Attention Sinks nicht unterstützt, was zu falschen Trainingsverlustenführt. Wenn Sie nicht Unsloth verwenden, kann FA3 standardmäßigaktiviert sein, also bitte prüfen Sie genau, dass es nicht verwendet wird! Das Deaktivieren von FA3 führt ebenfalls zu O(N^2) Speicherverbrauch, daher ist Unsloth das einzige RL-Framework, das O(N) Speicherverbrauch für gpt-oss über unsere Flex-Attention-Implementierung bietet.

⚡Inferenz deutlich schneller machen

Inference ist im RL-Training entscheidend, da wir sie benötigen, um Kandidatenlösungen zu erzeugen, bevor wir eine Belohnungsfunktion maximieren (sehen Sie 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, unter Verwendung spezieller Flags innerhalb von torch.compile (wie Combo-Kernel). Unser neuer Inferenzcode für gpt-oss wurde mit einer bereits optimierten Baseline verglichen (2x schneller als native Transformers).

vLLM unterstützt kein RL für gpt-oss, da es kein BF16-Training und keine LoRA-Unterstützung für gpt-oss bietet. Ohne Unsloth funktioniert nur das Training über volle Präzision BF16, was den Speicherverbrauch um 800 %+ erhöht. Die meisten Frameworks aktivieren FA3 (Flash Attention 3) standardmäßig (was den VRAM-Verbrauch senkt und die Geschwindigkeit erhöht), aber dies verursacht falschen Trainingsverlust. Siehe Issue 1797 im FA3-Repo. Sie müssen FA3 dennoch deaktivieren, da es sonst Long-Context-Training verhindert, weil FA3 O(N)-Speicherverbrauch verwendet, während naive Attention mit O(N^2)-Nutzung explodieren würde. 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 auch separate Tests für BF16 durchgeführt. Unsloths 4-Bit-Inferenz ist ~4x schneller, und BF16 ist ebenfalls effizienter, insbesondere 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 verwenden ältere 15-GB-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 Linkspadding zu ermöglichen. Wir mussten die logsumexp berechnen und die Sigmoid-Aktivierung anwenden, um die Attention-Gewichte wie unten zu verä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})$$

Linkspadded Masking während der Inferenz war ebenfalls ein kniffliges Problem bei gpt-oss. Wir stellten fest, dass wir nicht nur das KV-Cache-Prefill während der Token-Generierung berücksichtigen mussten, sondern auch eine einzigartige Anzahl von Pad-Tokens in jedem Prompt für Batch-Generierungen, was die Art und Weise verändern würde, wie wir die Blockmaske speichern müssten. Ein Beispiel dafür ist unten zu sehen:

Normale kausale Maske:

   k0 k1 k2 k3 k4   <-- Keys
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   (beachten Sie, dass q4 q_idx=0 hat, da dies die erste Query im aktuellen Setup 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), scheitert das, da unsere einzelne Query den Index 0 hat, während es n_k Key-Tokens gibt. Um das zu beheben, brauchen wir einen Offset bei der Maskenerstellung, um zu entscheiden, auf welche Tokens geachtet wird. Aber ein naiver Ansatz ist langsam, da sich Offsets bei jedem Schritt ändern und dadurch Masken- und Kernel-Neugenerierung erzwingen. Wir haben das 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 es nicht kompiliert ist, fällt es auf Eager Attention zurück, was langsam und speicherintensiv ist (quadratisch statt 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 Ihre Blockmaske im Wesentlichen auf eine vollständige Q_LEN x KV_LEN-Matrix ab, um die Blockmaske zu erzeugen. Ohne Compile müssen wir das Ganze vollständig materialisieren, und das 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-fusionierte Eager-Implementierung zurück, die sich gut zum Debuggen eignet, aber viel langsamer ist und die vollständige Scores-Matrix materialisiert.

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

🔍 Untersuchung von Flash Attention

Ein weiterer interessanter Ansatz, den wir untersucht haben, war die Integration von Flash Attention. Seine Vorteile sind weithin anerkannt, aber eine Einschränkung besteht darin, dass Attention Sinks während des Backward-Passes für gpt-oss nicht unterstützt werden. Um dies zu umgehen, haben wir den Attention-Mechanismus so umstrukturiert, dass er ausschließlich auf dem Attention-Output und den logsumexp-Werten basiert, die FlashAttention problemlos bereitstellt. Angesichts dieser Vorteile schien es naheliegend, es zu versuchen.

Wir begannen jedoch bald, Probleme zu bemerken. Während sich die ersten paar Schichten wie erwartet verhielten, lieferten die späteren Schichten, insbesondere die Schichten 18 bis 24, Ausgaben, die erheblich 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 auf jeder Schicht identisch sind. Zur weiteren Validierung haben wir die Ergebnisse auch mit Unsloth verglichen FlexAttention.

Dies erfordert weitere Untersuchungen, warum nur die letzten wenigen 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 falsch. Die Verwendung von FA3 macht den Trainingsverlust völlig falsch, da FA3 keine gpt-oss-Backward-Pässe für Attention Sinks unterstützt. Vielen Menschen ist das immer noch nicht bewusst, also bitte Vorsicht!

⚠️ Können wir Reward-Hacking entgegenwirken?

Das ultimative Ziel von RL ist es, eine Belohnung zu maximieren (z. B. Geschwindigkeit, Umsatz, eine Metrik). Aber RL kann schummeln. Wenn der RL-Algorithmus einen Trick lernt oder etwas ausnutzt, um die Belohnung zu erhöhen, ohne die Aufgabe am Ende tatsächlich zu erledigen, nennt man das „Reward Hacking".

Es ist der Grund, warum Modelle lernen, Unit-Tests zu ändern, um Coding-Herausforderungen zu bestehen, und das 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 Reward Hacking in einem Codegenerierungs-Setting entgegenwirken kann, und zeigen greifbare Lösungen für häufige Fehlerarten. Wir sahen, wie das Modell die Timing-Funktion bearbeitete, auf andere Bibliotheken auslagerte, die Ergebnisse zwischenspeicherte und offen schummelte. Nach dem Gegensteuern ist das Ergebnis, dass unser Modell wirklich optimierte Matrixmultiplikations-Kernel generiert und keine cleveren Cheats.

🏆Reward Hacking

Einige häufige Beispiele für Reward Hacking 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.

Caching & Schummeln

RL lernt, das Ergebnis der Ausgabe zwischenzuspeichern, und RL lernt, die tatsächliche Ausgabe zu finden, indem es globale Python-Variablen inspiziert.

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

Schummeln

RL lernt, die Timing-Funktion zu bearbeiten, sodass sie wie übergeben 0 Zeit ausgibt. Wir können den RL-Algorithmus daran hindern, globale oder zwischengespeicherte Variablen zu verwenden, indem wir seine lokalen und globalenEinschränkungen. Wir werden auch exec verwenden, um die Funktion zu erstellen, daher müssen wir die Ausgabe in einem leeren Dict speichern. Wir verbieten 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 das Anwenden von bestärkendem Lernen (RL) und das Entwerfen einer benutzerdefinierten Belohnungsfunktionkönnen diese Herausforderungen jedoch überwunden werden.

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

Unsere Notebooks enthalten bereits Schritt-für-Schritt-Anleitungen, wie man den gesamten Prozess durchläuft.

2048-Notebook (Offizielles OpenAI-Beispiel)	Notebook zur Kernel-Generierung

Was Sie erstellen werden:

• Trainiere gpt-oss-20b so, dass das Modell 2048 automatisch gewinnt

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

• Definiere Belohnungsfunktionen die:

  1. Überprüfen, ob die generierte Strategie kompiliert und ausgeführt werden kann,

  2. Reward-Hacking verhindern (externe Imports verbieten) und

  3. den tatsächlichen Spielerfolg belohnen

• Inference ausführen und das Modell exportieren (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-Laden + LoRA ermöglicht es Ihnen, ein 20B-Modell in moderaten VRAM zu packen