AlphaEvolve: Ein Coding-Agent für wissenschaftliche und algorithmische Entdeckungen
Optimierung des Gemini-Kernels
Das Training großer Modelle wie Gemini erfordert erhebliche Rechenressourcen. Gemini basiert auf dem JAX-Framework, wobei Pallas eine Erweiterung darstellt, die das Schreiben hochspezialisierter Programme (sogenannter Kernel) ermöglicht, die für eine optimale Ausführung auf Hardwarebeschleunigern angepasst sind. Effiziente Pallas-Kernel sind daher entscheidend für die Optimierung der Trainingsleistung des Gemini-Modells. Ein wichtiger Aspekt der Kernel-Optimierung ist die Feinabstimmung der Kachelungsstrategie für Matrixmultiplikationsoperationen. Diese Technik umfasst die Aufteilung der Berechnung einer großen Matrixmultiplikation in kleinere Teilprobleme, um die Berechnung besser mit dem Datentransfer abzustimmen, was für die Beschleunigung der gesamten Berechnung entscheidend ist. Traditionell verlassen sich Kernel-Entwickler entweder auf suchbasierte automatische Abstimmung oder auf manuell erstellte Heuristiken, um nahezu optimale Kachelkonfigurationen für verschiedene Eingabeformen zu bestimmen. Die automatische Abstimmung unterbricht jedoch den Forschungsworkflow und erfordert für jede Änderung der Eingabeform eine erneute Abstimmung. Die manuelle Erstellung effizienter Kachelheuristiken stellt dagegen aufgrund ihrer Komplexität eine große technische Herausforderung dar, da sie ein tiefes Verständnis sowohl der Funktionalität des Kernels als auch der hardwarespezifischen Eigenschaften erfordert. Der Hauptvorteil einer leistungsfähigen Heuristik besteht in ihrer Fähigkeit, über beliebige Eingabeformen hinweg eine hohe Leistung zu liefern. Um daher den Entwurf leistungsfähiger Kernel für neue Hardware zu beschleunigen und ihre Nutzung durch Modellentwickler zu vereinfachen, wollen wir den Prozess der Heuristikgenerierung erleichtern.
Diese Herausforderung wurde durch den Einsatz von AlphaEvolve zur Optimierung der Kachelheuristiken für einen wichtigen Matrixmultiplikations-Kernel gelöst, der beim Training von Gemini verwendet wird. Ziel war es, die tatsächliche Laufzeit des Kernels zu minimieren. AlphaEvolve untersucht und verfeinert iterativ die Kachelheuristiken für diesen Kernel, indem es Kandidatencode entwirft, mit dem Ziel, diese Laufzeit für verschiedene Eingabeformen auf realen TPU-Beschleunigern zu minimieren. Die Korrektheit des Kernels bleibt erhalten, da AlphaEvolve die Kachelungsstrategie für diesen Kernel optimiert, anstatt seine zugrunde liegende mathematische Operation zu verändern. Zur Erstellung der Trainings- und Evaluationsdatensätze für AlphaEvolve wurden realistische Eingabeformen des Kernels automatisch von seinen Nutzern gesammelt. Die Hälfte dieser Eingabeformen bildet die Trainingsmenge, die während des Evolutionsprozesses Optimierungsziele liefert. Die übrigen Eingabeformen bilden die Evaluationsmenge, die zum Testen der allgemeinen Anwendbarkeit der resultierenden Heuristik dient.
Dieser automatisierte Ansatz ermöglichte es AlphaEvolve, eine Heuristik zu entdecken, die im Vergleich zu einer von Experten entwickelten bestehenden Heuristik eine durchschnittliche Beschleunigung des Kernels um 23 % über alle Kernel hinweg erzielt und die gesamte Trainingszeit von Gemini entsprechend um 1 % reduziert. Darüber hinaus verkürzte der Einsatz von AlphaEvolve die Dauer der Kernel-Optimierung erheblich – von mehreren Monaten spezialisierter Ingenieursarbeit auf wenige Tage automatisierter Experimente. Diese Beschleunigung ermöglicht eine schnellere Bereitstellung optimierter Kernel, sodass Kernel-Ingenieure ihr Fachwissen strategischeren Optimierungsproblemen auf höherer Ebene widmen können. Darüber hinaus bietet AlphaEvolve einen Weg zur Automatisierung des manuellen Abstimmungsprozesses und zur Verbesserung der Benutzerfreundlichkeit von Gemini-Kerneln. Die mit AlphaEvolve entdeckte Kachelheuristik wurde in die Produktion übernommen und steigert direkt die Effizienz des Gemini-Trainings sowie die Forschungs- und Entwicklungsgeschwindigkeit des Gemini-Teams. Diese Bereitstellung stellt außerdem einen neuen Fall dar, in dem Gemini durch die Fähigkeiten von AlphaEvolve seinen eigenen Trainingsprozess optimiert.
Unterstützung beim Entwurf von Hardwareschaltungen
Spezialisierte Hardware wie Googles Tensor Processing Units (TPUs) ist entscheidend, um die für den Betrieb moderner KI-Systeme im großen Maßstab erforderliche Ressourceneffizienz zu erreichen. Der Entwurf neuer Computerchips ist jedoch ein komplexer und zeitaufwendiger Prozess, der oft Jahre dauert. Die Optimierung auf Register-Transfer-Ebene (Register-Transfer Level, RTL), ein kritischer Schritt in diesem Prozess, umfasst das manuelle Umschreiben von Hardwarebeschreibungen, um Kennzahlen wie Leistung, Energieverbrauch und Chipfläche zu verbessern, und erfordert monatelange Iterationen hochqualifizierter Ingenieure.
In dieser Arbeit wurde AlphaEvolve aufgefordert, eine bereits hochoptimierte Implementierung einer zentralen arithmetischen TPU-Schaltung in Verilog innerhalb der Matrixmultiplikationseinheit zu optimieren. Ziel der Optimierung war es, sowohl die Chipfläche als auch den Energieverbrauch zu reduzieren und dabei die grundlegende Funktionalität der Komponente beizubehalten. Entscheidend war, dass der endgültige Entwurf gründliche Verifikationsverfahren bestehen musste, die bestätigten, dass die modifizierte Schaltung ihre funktionale Korrektheit bewahrt. AlphaEvolve konnte eine einfache Umschreibung des Codes finden, durch die überflüssige Bits entfernt wurden – eine Änderung, deren Korrektheit von den TPU-Entwicklern überprüft wurde. Obwohl diese konkrete Verbesserung später auch unabhängig von Synthesewerkzeugen erfasst wurde, zeigt der Beitrag von AlphaEvolve in der RTL-Phase seine Fähigkeit, RTL-Quellcode zu verbessern und Optimierungen in einer frühen Phase des Entwurfsprozesses bereitzustellen.
Diese in die kommende TPU integrierte Verbesserung stellt den ersten direkten Beitrag von Gemini zu den arithmetischen TPU-Schaltungen dar, der mithilfe von AlphaEvolve erzielt wurde, und ebnet den Weg für zukünftige ähnliche Beiträge. Ein entscheidender Vorteil von AlphaEvolve besteht darin, dass es vorgeschlagene Änderungen direkt in Verilog formuliert, der Standardsprache für Hardware-Ingenieure, was Vertrauen schafft und die Einführung vereinfacht. Diese erste Erkundung demonstriert einen neuen Ansatz, bei dem die von Sprachmodellen gesteuerte Code-Evolution beim Hardwareentwurf unterstützt und potenziell die Zeit bis zur Markteinführung verkürzt.
Direkte Optimierung vom Compiler generierten Codes
Die Transformer-Architektur wird in den meisten modernen neuronalen Netzen eingesetzt, von großen Sprachmodellen bis hin zu AlphaFold. Der grundlegende Berechnungsmechanismus von Transformern ist die sogenannte Attention (Aufmerksamkeitsmechanismus), die am häufigsten mithilfe von FlashAttention implementiert wird. In unserem Computing-Stack ist FlashAttention als Beschleunigungskernel in Pallas implementiert und von einer höheren Codeebene in JAX umgeben, die sich um die Vorbereitung der Eingabe und die anschließende Verarbeitung der Ausgabe kümmert. Der Compiler für maschinelles Lernen (XLA) übersetzt diese Implementierung anschließend in eine Sequenz von Zwischenrepräsentationen (IR), wobei jede davon weitere Details für die Ausführung auf einer bestimmten Hardware hinzufügt. In diesen Phasen können verbesserte Entscheidungen zur Orchestrierung des Speicherzugriffs oder zur Planung von Berechnungen die Laufzeit auf spezifischer Hardware erheblich reduzieren.
AlphaEvolve wurde aufgefordert, die von XLA erzeugte IR direkt zu optimieren, die den FlashAttention-Kernel zusammen mit dem Code für die Vor- und Nachverarbeitung kapselt. Wir optimierten die Konfiguration eines besonders wichtigen Transformer-Modells, das für groß angelegte Inferenz auf GPUs verwendet wird, mit dem Ziel, die gesamte Ausführungszeit des Moduls zu minimieren. Dies war eine außerordentlich anspruchsvolle Aufgabe, da (1) die IR primär für Debugging-Zwecke und nicht für die direkte Bearbeitung durch Entwickler konzipiert ist und (2) sie vom Compiler erzeugt und bereits hochgradig optimiert wird. Jede von AlphaEvolve vorgeschlagene Änderung wurde anhand von Referenzcode (unmodifiziertem Code) mit zufälligen Eingaben überprüft, um während des gesamten Optimierungsprozesses die numerische Korrektheit sicherzustellen. Die endgültige Codeversion wurde von menschlichen Experten gründlich verifiziert und als für alle möglichen Eingaben korrekt bestätigt.
AlphaEvolve konnte sinnvolle Optimierungen für beide von der IR offengelegten Abstraktionsebenen bereitstellen. Erstens wurde der FlashAttention-Kernel für die entsprechende Konfiguration um 32 % beschleunigt. Zweitens fand AlphaEvolve Verbesserungen bei der Vor- und Nachverarbeitung der Kernel-Ein- und -Ausgaben, was in diesem Bereich zu einer Beschleunigung um 15 % führte. Diese Ergebnisse zeigen die Fähigkeit von AlphaEvolve, vom Compiler erzeugten Code zu optimieren. Dies eröffnet die Möglichkeit, die entdeckten Optimierungen für bestimmte Anwendungsfälle in bestehende Compiler zu integrieren oder AlphaEvolve längerfristig direkt in den Compiler-Workflow einzubinden.
Obwohl AlphaEvolve ein relativ neues Werkzeug ist, bringt es bereits jetzt bedeutende Vorteile in drei zentralen Bereichen der Entwicklung und Optimierung von KI und Hardware. Die Fähigkeit, komplexe Kernel automatisch zu optimieren, zum Entwurf von Hardwareschaltungen beizutragen und vom Compiler erzeugten Code zu verbessern, eröffnet neue Möglichkeiten für eine effizientere Nutzung von Rechenressourcen und eine Beschleunigung der KI-Forschung. Mit der Weiterentwicklung dieses Werkzeugs ist zu erwarten, dass seine Auswirkungen auf wissenschaftliche und algorithmische Entdeckungen noch größer werden.
