A RightNow AI AutoKernelje 2-szeres gyorsulást ígér PyTorch modelleknél
Az AutoKernel egy nyílt forráskódú keretrendszer, amely egy autonóm LLM ágenst használ a GPU-kernel-optimalizáláshoz, és a PyTorch modellek esetében automatizálja a feladatot.
A RightNow AI kutatócsapata egy nyílt forráskódú keretrendszert mutatott be AutoKernel néven, amely autonóm LLM ágensek segítségével optimalizálja a GPU-kerneleket PyTorch modellekhez — írja a MarkTechPost. Ez a megoldás nagymértékben automatizálja a feladatot, így a gépi tanulási mérnököknek nem kell manuálisan finomhangolniuk a kódot.
Az AutoKernel célja, hogy a GPU-kód írásának rendkívül bonyolult feladatát teljes mértékben automatizálja. A fejlesztők szerint a felhasználók egyszerűen betölthetnek bármilyen modellt, és reggelre már gyorsabb Triton kernelekkel dolgozhatnak, anélkül, hogy mélyreható GPU-szakértelemmel rendelkeznének.
A GPU-kernelek optimalizálása rendkívül összetett, mivel egyszerre kell figyelembe venni az aritmetikai intenzitást, a memória-összefésülést, a regiszterek terhelését és a csempe méretét. Egyetlen nagy teljesítményű mátrixszorzó kernel akár 200 sor feletti CUDA vagy Triton kódot is tartalmazhat, több tucat kölcsönösen függő paraméterrel.
Az AutoKernel működési elve
Az AutoKernel alapvető működési elve egy egyszerű ciklusra épül: a LLM ágens módosít egy „kernel.py” fájlt, egy fix teljesítményteszt ellenőrzi a helyességet és méri az átviteli sebességet, majd az eredmény alapján eldől, hogy a változás megmarad-e. Minden kísérlet egy git commitnak felel meg, a sikeres kísérletek előreviszik az ágat, a sikertelenek pedig visszaállíthatók.
Minden iteráció körülbelül 90 másodpercet vesz igénybe: 30 másodperc a helyesség ellenőrzésére, 30 másodperc a teljesítmény mérésére a Triton „do_bench” funkciójával, és 30 másodperc az ágens gondolkodására és kódmódosítására. Egy éjszakai, 10 órás futás során 300-400 kísérlet is lefuthat több kernelen.
Optimalizálási stratégiák
Az ágens egy 909 soros „program.md” utasításdokumentumot olvas be, amely szakértői tudást kódol egy hatlépcsős optimalizálási tervbe. Az optimalizálási szintek a blokkméret-hangolástól a memória-hozzáférési mintákon és számítási optimalizációkon át az architektúra-specifikus stratégiákig terjednek, mint például a TMA a Hopper GPU-kon vagy a cp.async az Ampere architektúrán.
Az AutoKernel a teljes PyTorch modellből indul ki, és a „torch.profiler” segítségével rögzíti a GPU-időt kernelenként. Az optimalizálási célpontokat Amdahl törvénye alapján rangsorolja, ami azt jelenti, hogy a teljesítményjavulás mértéke attól függ, hogy az adott komponens mennyit tesz ki a teljes futási időből.
A profiler ismeri a NVIDIA (H100, A100, L40S, L4, A10, RTX 4090/4080/3090/3080) és AMD (MI300X, MI325X, MI350X, MI355X) gyorsítók specifikációit. Ismeretlen GPU-k esetén a SM-szám, az órajel és a számítási képesség alapján becsüli a FP16 átviteli sebességet.
Az orchestrator (orchestrate.py) négy feltétel teljesülése esetén vált át a következő kernelre: öt egymást követő visszaállítás, a GPU csúcskihasználtságának 90%-a, két óra eltelt időkeret, vagy az adott kernelen már elért 2-szeres sebességjavulás. 2024. március 15-én az AutoKernel fejlesztői bejelentették, hogy a technológia a következő hetekben elérhető lesz a nyilvánosság számára.