AlphaEvolve: Kódovací agent pro vědecké a algoritmické objevy
Optimalizace jádra Gemini
Trénování velkých modelů jako Gemini vyžaduje značné výpočetní zdroje. Gemini je postaven na frameworku JAX, přičemž Pallas představuje jeho rozšíření, které umožňuje psát vysoce specializované programy (tzv. kernely) přizpůsobené pro optimální běh na hardwarových akcelerátorech. Efektivní Pallas kernely jsou proto zásadní pro optimalizaci výkonu tréninku modelu Gemini. Důležitým aspektem optimalizace kernelu je ladění strategie dlaždic pro operace násobení matic. Tato technika zahrnuje rozdělení výpočtu velkého násobení matic na menší podproblémy, aby se lépe vyvážil výpočet s přesunem dat, což je klíčové pro urychlení celkového výpočtu. Tradičně se vývojáři kernelů spoléhají buď na automatické ladění založené na vyhledávání, nebo na ručně vytvořené heuristiky k určení téměř optimálních konfigurací dlaždic pro různé vstupní tvary. Automatické ladění však přerušuje výzkumný pracovní postup a vyžaduje přeladění pro každou změnu vstupního tvaru. Naopak ruční vytváření efektivních heuristik dlaždic představuje hlavní inženýrský problém kvůli své složitosti, která vyžaduje hluboké porozumění jak funkčnosti kernelu, tak hardwarovým specifikům. Hlavní výhodou výkonné heuristiky je její schopnost poskytovat vysoký výkon napříč libovolnými vstupními tvary. Proto je pro urychlení návrhu výkonných kernelů pro nový hardware a zjednodušení jejich využití vývojáři modelů naším cílem usnadnit proces generování heuristik.
Ta výzva byla řešena nasazením AlphaEvolve k optimalizaci heuristik dlaždic pro důležitý kernel násobení matic používaný při tréninku Gemini. Cílem bylo minimalizovat skutečnou dobu běhu kernelu. AlphaEvolve iterativně zkoumá a zdokonaluje heuristiky dlaždic pro tento kernel navrhováním kandidátského kódu, s cílem minimalizovat tuto dobu běhu na různých vstupních tvarech na skutečných akcelerátorech TPU. Korektnost kernelu je zachována, protože AlphaEvolve optimalizuje strategii dlaždic pro tento kernel, místo aby měnil jeho základní matematickou operaci. Pro vytvoření trénovacích a evaluačních datasetů pro AlphaEvolve byly automaticky shromážděny realistické vstupní tvary kernelu od jeho uživatelů. Polovina těchto vstupních tvarů tvoří trénovací sadu, která poskytuje optimalizační cíle během evolučního procesu. Zbývající vstupní tvary tvoří evaluační sadu, která slouží k testování obecné aplikovatelnosti výsledné heuristiky.
Tento automatizovaný přístup umožnil AlphaEvolve objevit heuristiku, která poskytuje průměrné 23% zrychlení kernelu napříč všemi kernely ve srovnání s existující heuristikou navrženou experty, a odpovídající 1% snížení celkové doby tréninku Gemini. Kromě toho použití AlphaEvolve významně zkrátilo dobu optimalizace kernelu, z několika měsíců specializované inženýrské práce na pouhé dny automatizovaných experimentů. Toto zrychlení urychluje nasazení optimalizovaných kernelů, což umožňuje inženýrům kernelů věnovat své odborné znalosti strategičtějším, optimalizačním problémům vyšší úrovně. Navíc AlphaEvolve nabízí cestu k automatizaci procesu manuálního ladění a zlepšení ergonomie používání kernelů Gemini. Heuristika dlaždic objevená pomocí AlphaEvolve byla nasazena do produkce, přímo zvyšuje efektivitu tréninku Gemini a rychlost výzkumu a vývoje týmu Gemini. Toto nasazení také představuje nový případ, kdy Gemini, prostřednictvím schopností AlphaEvolve, optimalizuje svůj vlastní trénovací proces.
Asistence při návrhu hardwarových obvodů
Specializovaný hardware, jako jsou Google's Tensor Processing Units (TPU), je zásadní pro dosažení efektivity zdrojů potřebné k provozování moderních AI systémů ve velkém měřítku. Návrh nových počítačových čipů je však komplexní a časově náročný proces, který často trvá roky. Optimalizace na úrovni Register-Transfer Level (RTL), která je kritickým krokem v tomto procesu, zahrnuje ruční přepisování hardwarových popisů za účelem zlepšení metrik jako jsou výkon, spotřeba energie a plocha, což vyžaduje měsíce iterací vysoce kvalifikovaných inženýrů.
V této práci byl AlphaEvolve vyzván, aby optimalizoval již vysoce optimalizovanou implementaci klíčového aritmetického obvodu TPU ve Verilogu v rámci jednotky násobení matic. Cílem optimalizace bylo snížit jak plochu, tak spotřebu energie při zachování základní funkčnosti komponenty. Zásadní bylo, že konečný návrh musel projít důkladnými verifikačními metodami potvrzujícími, že upravený obvod zachovává funkční správnost. AlphaEvolve dokázal najít jednoduché přepsání kódu, které odstranilo nepotřebné bity - změnu, jejíž správnost byla ověřena návrháři TPU. Ačkoli toto konkrétní vylepšení bylo později nezávisle zachyceno i nástroji pro syntézu, příspěvek AlphaEvolve ve fázi RTL ukazuje jeho schopnost vylepšovat zdrojové RTL a poskytovat optimalizace v rané fázi návrhového procesu.
Toto vylepšení, integrované do připravovaného TPU, představuje první přímý příspěvek Gemini k aritmetickým obvodům TPU, dosažený prostřednictvím AlphaEvolve, a otevírá cestu pro budoucí podobné příspěvky. Klíčovou výhodou AlphaEvolve je, že komunikuje navrhované změny přímo ve Verilogu, standardním jazyce používaném hardwarovými inženýry, což podporuje důvěru a zjednodušuje adopci. Tato prvotní explorace demonstruje nový přístup, kde kódová evoluce poháněná jazykovými modely asistuje při návrhu hardwaru a potenciálně zkracuje dobu uvedení na trh.
Přímá optimalizace kódu generovaného kompilátorem
Architektura transformeru je používána ve většině moderních neuronových sítí, od velkých jazykových modelů až po AlphaFold. Základním výpočetním mechanismem transformerů je tzv. attention (mechanismus pozornosti), který je nejčastěji implementován pomocí FlashAttention. V našem výpočetním stacku je FlashAttention implementován jako akcelerační kernel v Pallasu, obalený vyšší úrovní kódu v JAXu, který se stará o přípravu vstupu a následné zpracování výstupu. Kompilátor pro strojové učení (XLA) pak překládá tuto implementaci do sekvence mezilehlých reprezentací (IR), přičemž každá z nich přidává více detailů pro spuštění na konkrétním hardwaru. V těchto fázích mohou vylepšená rozhodnutí o orchestraci přístupu do paměti nebo plánování výpočtů významně snížit dobu běhu na specifickém hardwaru.
AlphaEvolve byl vyzván, aby přímo optimalizoval IR generované XLA kompilátorem, které zapouzdřuje kernel FlashAttention spolu s kódem pro předzpracování a dodatečné zpracování. Optimalizovali jsme konfiguraci odpovídající vysoce důležitému transformerovému modelu používanému pro inference ve velkém měřítku na GPU, s cílem minimalizovat celkovou dobu provádění modulu. Byl to mimořádně náročný úkol, protože (1) IR je navržena primárně pro účely ladění, nikoli pro přímou editaci vývojáři, a (2) je generována kompilátorem a již vysoce optimalizována. Každá modifikace navržená AlphaEvolve byla kontrolována proti referenčnímu (nemodifikovanému) kódu na náhodných vstupech, aby byla zajištěna numerická správnost během celého procesu optimalizace. Konečná verze kódu byla důkladně ověřena lidskými experty a potvrzena jako správná pro všechny možné vstupy.
AlphaEvolve dokázal poskytnout smysluplné optimalizace pro obě úrovně abstrakce vystavené IR. Zaprvé, kernel FlashAttention pro příslušnou konfiguraci byl zrychlen o 32%. Zadruhé, AlphaEvolve našel vylepšení v předzpracování a následném zpracování vstupů a výstupů kernelu, což vedlo k 15% zrychlení v této části. Tyto výsledky demonstrují schopnost AlphaEvolve optimalizovat kód generovaný kompilátorem, což nabízí potenciál pro začlenění objevených optimalizací do existujících kompilátorů pro specifické případy použití, nebo z dlouhodobějšího hlediska, zahrnutí AlphaEvolve přímo do kompilačního workflow.
Přestože je AlphaEvolve relativně novým nástrojem, již nyní přináší významné přínosy ve třech klíčových oblastech vývoje a optimalizace umělé inteligence a hardwaru. Schopnost automaticky optimalizovat složité kernely, přispívat k návrhu hardwarových obvodů a vylepšovat kompilátor generovaný kód otevírá nové možnosti pro efektivnější využití výpočetních zdrojů a urychlení výzkumu v oblasti AI. S pokračujícím vývojem tohoto nástroje lze očekávat, že jeho dopady na vědecké a algoritmické objevy budou ještě výraznější.
