OpenXLA je společně vyvinutý open source ekosystém kompilátoru ML
Největší společnosti věnující se vývoji v oblasti strojového učení nedávno představily projekt OpenXLA , jehož cílem je společný vývoj nástrojů pro kompilaci a optimalizaci modelů pro systémy strojového učení.
Projekt si vzal na starost vývoj nástrojů, které umožňují sjednotit kompilaci modelů připravených v rámci TensorFlow, PyTorch a JAX pro efektivní školení a spouštění na různých GPU a specializovaných akcelerátorech. Ke společné práci na projektu se připojily společnosti jako Google, NVIDIA, AMD, Intel, Meta, Apple, Arm, Alibaba a Amazon.
Projekt OpenXLA poskytuje nejmodernější kompilátor ML, který lze škálovat v rámci složitosti infrastruktury ML. Jeho základními pilíři jsou výkon, škálovatelnost, přenositelnost, flexibilita a rozšiřitelnost pro uživatele. S OpenXLA se snažíme odemknout skutečný potenciál umělé inteligence urychlením jejího vývoje a poskytování.
OpenXLA umožňuje vývojářům kompilovat a optimalizovat modely ze všech předních ML frameworků pro efektivní školení a servis na široké škále hardwaru. Vývojáři používající OpenXLA zaznamenají výrazné zlepšení v době školení, výkonu, latenci služeb a v konečném důsledku i v době uvedení na trh a výpočtu nákladů.
Doufá se, že spojením úsilí předních výzkumných týmů a zástupců komunity bude možné stimulovat rozvoj systémů strojového učení a vyřešit problém fragmentace infrastruktury pro různé frameworky a týmy.
OpenXLA umožňuje implementaci efektivní podpory pro různý hardware, bez ohledu na framework, na kterém je model strojového učení založen. Očekává se, že OpenXLA zkrátí dobu trénování modelu, zlepší výkon, sníží latenci, sníží výpočetní režii a zkrátí dobu uvedení na trh.
OpenXLA se skládá ze tří hlavních komponent , jejichž kód je distribuován pod licencí Apache 2.0:
- XLA (accelerated linear algebra) je kompilátor, který umožňuje optimalizovat modely strojového učení pro vysoce výkonné provádění na různých hardwarových platformách, včetně GPU, CPU a specializovaných akcelerátorů od různých výrobců.
- StableHLO je základní specifikace a implementace sady operací na vysoké úrovni (HLO) pro použití v modelech systémů strojového učení. Funguje jako vrstva mezi frameworky strojového učení a kompilátory, které transformují model tak, aby běžel na konkrétním hardwaru. Vrstvy jsou připraveny pro generování modelů ve formátu StableHLO pro frameworky PyTorch, TensorFlow a JAX. Jako základ pro StableHLO je použita sada MHLO, která je rozšířena o podporu serializace a správy verzí.
- IREE (Intermediate Representation Execution Environment) je kompilátor a běhové prostředí, které převádí modely strojového učení do univerzální střední reprezentace založené na formátu MLIR (Intermediate Multi-Level Representation) projektu LLVM. Z funkcí je vyzdvihována možnost předkompilace (v předstihu), podpora řízení toku, možnost využití dynamických prvků v modelech, optimalizace pro různá CPU a GPU a nízká režie.
Mezi hlavní výhody OpenXLA patří dosažení optimálního výkonu bez nutnosti zabývat se psaním kódu specifického pro dané zařízení, a to navíc s využitím optimalizací připravených k použití, včetně zjednodušení algebraických výrazů, efektivní alokace paměti, plánování provádění s ohledem na snížení maximální spotřeby paměti a režie.
Další výhodou je zjednodušení škálování a paralelizace výpočtů. Vývojář stačí přidat anotace pro podmnožinu kritických tenzorů a kompilátor pak může automaticky generovat kód pro paralelní výpočty.
Je také zdůrazněno, že přenositelnost je zajištěna podporou pro více hardwarových platforem, jako jsou grafické karty AMD a NVIDIA, procesory x86 a ARM, akcelerátory strojového učení Google TPU, IPU AWS Trainium Inferentia, Graphcore a engine Cerebras Wafer-Scale.
Podpora propojení rozšíření s implementací dalších funkcí, jako je podpora psaní primitiv hlubokého strojového učení pomocí CUDA, HIP, SYCL, Triton a dalších jazyků pro paralelní výpočty, a také možnost ruční úpravy úzkých míst v modelech.
A konečně, pokud máte zájem dozvědět se více, podrobnosti naleznete na následujícím odkazu.