Equivalence Checking of ML GPU Kernels
Dieses Paper stellt Volta vor, den ersten korrekten und vollständigen Äquivalenzprüfer für GPU-Kernel, der die Korrektheit von maschinellen Lernberechnungen verifiziert, die von Hand, durch Compiler oder durch LLMs optimiert wurden.
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
In der gewaltigen, unsichtbaren Maschinerie der modernen künstlichen Intelligenz findet die wichtigste Arbeit nicht in der Cloud statt, sondern auf spezialisierten Computerchips, den sogenannten GPUs. Diese Chips sind darauf ausgelegt, Millionen von winzigen Berechnungen gleichzeitig durchzuführen – eine Notwendigkeit für das Training der großen Sprachmodelle, die heute Code schreiben, Sprachen übersetzen und Kunst generieren. Um diese Modelle schnell genug laufen zu lassen, damit sie nützlich sind, müssen Ingeniere hochspezialisierte Anweisungen schreiben, die als „Kernels“ bekannt sind und der GPU genau vorschreiben, wie sie Daten bewegen und mathematische Operationen ausführen soll. In den letzten Jahren haben Unternehmen begonnen, selbst künstliche Intelligenz einzusetzen, um diese Kernels zu schreiben, in der Hoffnung, schnellere Wege für die Arbeit zu finden, als es menschliche Ingenieure könnten. Diese Geschwindigkeit bringt jedoch ein Risiko mit sich: Wenn eine KI oder ein Compiler ein Stück Code umschreibt, um es schneller zu machen, kann dies versehentlich subtile Fehler einführen. Diese Fehler können dazu führen, dass der Computer das falsche Ergebnis liefert oder, schlimmer noch, auf eine Weise lautlos abstürzt, die durch Standardtests kaum zu finden ist. Die zentrale Herausforderung besteht darin, dass diese Chips tausende von Threads (Arbeitsfäden) gleichzeitig ausführen, und wenn diese nicht perfekt koordiniert sind, können sie sich gegenseitig in die Quere kommen, was eine sogenannte „Race Condition“ (Wettlaufsituation) erzeugt, bei der das Endergebnis von der unvorhersehbaren Reihenfolge abhängt, in der die Ereignisse ablaufen.
Ein Team von Forschern hat ein neues Werkzeug namens Volta entwickelt, um dieses Problem zu lösen. Anstatt zu raten, ob eine neue, schnellere Version eines Kernels korrekt ist, fungiert Volta als formaler Verifizierer, der mathematisch beweist, dass zwei Versionen identische Ergebnisse liefern. Die Forscher bauten ein System, das die Low-Level-Anweisungen eines Referenz-Kernels – der ursprünglichen, vertrauenswürdigen Version – und des optimierten Kernels – der neuen, schnelleren Version – nimmt und sie durch eine symbolische Engine laufen lässt. Anstatt dem Code spezifische Zahlen zuzuführen und zu sehen, was herauskommt, behandelt die Engine die Eingaben als abstrakte Symbole. Sie verfolgt jeden möglichen Pfad, den der Code nehmen könnte, und zeichnet nach, wie Daten durch die tausenden parallelen Threads fließen und wie diese miteinander synchronisieren. Wenn der Code versucht, auf Speicher zuzugreifen, der einen Konflikt verursachen könnte, oder wenn die Threads ewig aufeinander warten, markiert das Tool den Fehler sofort. Wenn der Code fehlerfrei läuft, übersetzt das Tool die Endausgabe beider Kernels in komplexe mathematische Ausdrücke und prüft, ob diese Ausdrücke fundamental gleich sind, unabhängig von den spezifischen Zahlen, die eingespeist werden.
Die Forscher testeten Volta bei einer Vielzahl von realen Aufgaben des maschinellen Lernens, einschließlich Matrixmultiplikationen, Faltungen (Convolutions) und der Attention-Mechanismen, die große Sprachmodelle antreiben. Sie fanden heraus, dass das Tool in der Lage war, Kernels erfolgreich zu verifizieren, die von Hand, durch Compiler und sogar durch große Sprachmodelle optimiert worden waren. In einem Fall untersuchten sie einen KI-generierten Kernel, der durch dreizehn Runden automatisierter Verbesserungen optimiert worden war. Volta bestätigte, dass dieser KI-generierte Code mathematisch äquivalent zum ursprünglichen, von Menschen geschriebenen Referenzcode war, was bewies, dass die aggressiven Optimierungen die Logik nicht zerstört hatten. Das Tool bewies auch seinen Wert, indem es Fehler entdeckte, die andere Methoden übersahen. Beispielsweise detektierte es Data Races in einem populären, viel zitierten Tutorial für GPU-Programmierung, das seit Jahren von tausenden Entwicklern genutzt wurde. Diese Fehler blieben verborgen, weil sie nur unter sehr spezifischen Zeitbedingungen auftraten, die Standardtests selten erfassen. Das Tool identifizierte zudem einen Bug in einem KI-generierten Kernel, bei dem der Code versuchte, Daten von einem Speicherort zu lesen, der nicht existierte; während die aktuelle Hardware diesen Fehler gerade noch ignorierte, zeigten die Forscher auf, dass der Code fundamental unsicher war und auf zukünftigen Maschinen scheitern könnte.
Die Stärke dieses Ansatzes liegt in seiner Fähigkeit, die einzigartige Komplexität der GPU-Programmierung zu bewältigen, bei der tausende Threads ihre Aktionen koordinieren müssen. Frühere Werkzeuge konnten Single-Threaded-Programme oder hochgradige mathematische Operationen prüfen, hatten aber Schwierigkeiten, die massive Parallelität einer GPU in handhabbare Teile zu zerlegen. Volta überwindet dies durch die Annahme, dass die analysierten Kernels einem spezifischen, strukturierten Muster folgen, das im maschinellen Lernen üblich ist, wobei die Anzahl der Threads und die Größe der Daten im Voraus bekannt sind. Innerhalb dieses Rahmens kann das Tool mit Sicherheit beweisen, ob eine Race Condition vorliegt, falls die Threads nicht ordnungsgemäß synchronisiert sind, und es kann beweisen, dass zwei verschiedene Versionen eines Programms äquivalent sind, wenn sie das gleiche symbolische Ergebnis liefern. Die Forscher demonstrierten, dass ihr Tool diese Eigenschaften innerhalb von Sekunden oder Minuten verifizieren konnte, selbst bei Kernels, die Hunderttausende von Instruktionen enthalten. Sie bewiesen auch, dass die mathematische Logik hinter ihrem Tool fundiert ist, was bedeutet: Wenn das Tool sagt, dass zwei Programme gleich sind, dann sind sie es auch für alle möglichen Eingaben.
Diese Arbeit stellt einen bedeutenden Schritt dar, um die Entwicklung künstlicher Intelligenz sicherer und zuverlässiger zu machen. Da Unternehmen zunehmend auf automatisierte Systeme setzen, um den Code zu generieren, der ihre Modelle antreibt, wird der Bedarf an einer rigorosen Methode zur Überprüfung dieses Codes entscheidend. Die Forscher zeigten, dass es möglich ist, über einfaches Testen hinauszugehen, welches nur eine begrenzte Anzahl von Szenarien prüfen kann, hin zu einer Methode, die eine formale Garantie der Korrektheit bietet. Durch die Verifizierung der Äquivalenz optimierter Kernels gibt Volta Entwicklern das Vertrauen, schnellere und aggressivere Optimierungen zu nutzen, ohne Angst vor der Einführung stiller Bugs haben zu müssen. Das Tool steht derzeit zur Nutzung zur Verfügung, und die Forscher haben ihren Code sowie die Beweise dahinter der Öffentlichkeit zugänglich gemacht, um anderen die Möglichkeit zu geben, auf diesem Fundament aufzubauen. Obwohl das Tool noch nicht jeden denkbaren Typ von GPU-Code abdeckt, bearbeitet es erfolgreich die überwiegende Mehrheit der Kernels, die das moderne maschinelle Lernen antreiben, und setzt damit einen neuen Standard für Vertrauen in die automatisierte Generierung von Hochleistungsrechen-Code.
Ertrinken Sie in Arbeiten in Ihrem Fachgebiet?
Erhalten Sie tägliche Digests der neuesten Arbeiten passend zu Ihren Forschungsbegriffen — mit technischen Zusammenfassungen, in Ihrer Sprache.