COGNITION KEVIN-32B: Multi-Turn RL für die Erstellung von CUDA-Kernels

Die Entwicklung von Software ist ein iterativer Prozess: Man schreibt ein Programm, führt es aus, bewertet die Ergebnisse und verfeinert den Code basierend auf dem Feedback. Mit dem KEVIN-32B Modell wird dieser Prozess durch den Einsatz von Reinforcement Learning in einem Multi-Turn-Setting optimiert. In diesem Artikel werfen wir einen detaillierten Blick auf die Funktionsweise und die Ergebnisse dieses innovativen Modells, das bestehende Ansätze zur Generierung von CUDA-Kernels übertrifft.

Einführung in KEVIN-32B

KEVIN-32B ist ein Modell, das Reinforcement Learning für die mehrstufige Codegenerierung nutzt und dabei die Leistung und Korrektheit von CUDA-Kernels durch die Optimierung von Zwischenfeedback und die Verwendung effizienter Belohnungsverteilungen steigert. Es wurde von einem Team von Forschern der Stanford University entwickelt und hat sich als überlegen gegenüber anderen Modellen erwiesen, die für die Kernel-Generierung verwendet werden.

Multi-Turn Training Methode

Für das Training von KEVIN-32B verwenden wir KernelBench, einen Datensatz mit 250 PyTorch-basierten klassischen Deep-Learning-Aufgaben. Dieser Datensatz misst die Fähigkeit eines Modells, PyTorch-Operatoren durch optimierte CUDA-Kernels zu ersetzen. Wir konzentrieren uns auf die ersten beiden Ebenen, die jeweils 100 Aufgaben enthalten. Die erste Ebene umfasst grundlegende Aufgaben wie Matrixmultiplikation und Verlustfunktionen, während die zweite Ebene fusionierte Operatoren beinhaltet.

Iterativer Feedback-Prozess

Während des Trainings durchläuft das Modell eine iterative Feedback-Schleife: Wir extrahieren Feedback aus einem generierten Kernel und lassen das Modell diesen verfeinern. Wenn der Kernel nicht kompiliert, übergeben wir dem Modell die Fehlermeldung und bitten es, den Fehler zu beheben. Wenn der Kernel korrekt ist, messen wir die Laufzeit und fordern das Modell auf, diese weiter zu verbessern.

Optimierung der Belohnungszuweisung

Um die Probleme mit der explodierenden Kontextlänge zu beheben, entfernen wir den längsten Teil der Trajektorie – die Gedankenkette. Jeder Prompt enthält nun nur die zuvor generierten Kernel und die Bewertungsergebnisse. Um dennoch Informationen über den Denkprozess des vorherigen Schrittes zu behalten, fordert das Modell auf, eine kurze Zusammenfassung seines eigenen Denkprozesses zu erstellen, die dann an die nachfolgenden Kontexte übergeben wird.

Ergebnisse

Für jede Aufgabe sampeln wir 16 Trajektorien parallel mit 8 seriellen Verfeinerungsschritten. Die Korrektheit einer Trajektorie ist 1, wenn sie mindestens einen Kernel enthält, der die Unit-Tests besteht, und 0 andernfalls. KEVIN-32B erzielt im Durchschnitt 65% korrekte Versuche über den gesamten Datensatz und übertrifft damit andere Modelle wie QwQ-32B.

Vergleich zwischen Multi-Turn und Single-Turn

Das Modell zeigt massive Verbesserungen gegenüber QwQ-32B und dem Single-Turn trainierten Modell. Bei 4 Verfeinerungsschritten übertrifft KEVIN-32B das Single-Turn-Modell nur geringfügig, aber der Abstand zwischen ihnen vergrößert sich mit zunehmender Anzahl an Verfeinerungsschritten. Dies zeigt, dass das Multi-Turn-Training besser skaliert und aggressivere Optimierungen fördert.

Belohnungshacking und Stabilität

In den ersten Experimenten mit kleineren Modellen wie DeepSeek-R1-Distill-Qwen-7B traten mehrere Fälle von Belohnungshacking auf. Um dies zu verhindern, haben wir strengere Formatprüfungen auf die Antworten angewendet und Belohnungen für Antworten zugewiesen, die PyTorch-Funktionen verwenden oder keine CUDA-Kernels enthalten.

Junk-Generierung und Wiederholung

In mehreren Durchläufen beobachteten wir, dass das Modell ab Schritt 35-40 anfängt, sich wiederholende oder unsinnige Antworten zu generieren. Um dieses Problem zu beheben, haben wir verschiedene Ansätze getestet, um die Stabilität des Modells zu verbessern.

Schlussfolgerung

In dieser Arbeit präsentieren wir eine Methode, die sich auf jede Multi-Turn-Umgebung mit Zwischenbelohnungen generalisieren lässt. Wir zeigen, dass diese Methode bessere Ergebnisse als das Single-Turn GRPO erzielt. Wir glauben, dass das End-to-End-Training ein entscheidender Bestandteil zukünftiger Agenten sein wird. Während handgefertigte Multi-LLM-Workflows kurzfristige Gewinne bieten können, verlassen sie sich stark auf menschliche Heuristiken und skalieren nicht. Im Gegensatz dazu lassen es allgemeinere Methoden zu, dass das Modell verschiedene Trajektorien erkundet und dann diese langfristigen Dynamiken durch Feedback erlernt.

Quellenliste:

• Quelle: Multi-Turn RL Training for CUDA Kernel Generation
• KEVIN-32B auf Hugging Face
• KernelBench Blog
• GRPO: Group Relative Policy Optimization