Najwięksi liderzy zajmujący się rozwojem w dziedzinie uczenia maszynowego przedstawili projekt OpenXLA, mający na celu wspólne rozwijanie narzędzi do kompilacji i optymalizacji modeli dla systemów uczenia maszynowego. W ramach projektu opracowywane są narzędzia, które mają na celu ujednolicenie kompilacji modeli stworzonych w ramach frameworków TensorFlow, PyTorch i JAX, by zapewnić efektywne uczenie i wykonanie na różnych GPU oraz specjalizowanych akceleratorach. Do wspólnej pracy nad projektem dołączyły takie firmy jak Google, NVIDIA, AMD, Intel, Meta, Apple, Arm, Alibaba i Amazon.
Oczekuje się, że dzięki połączeniu wysiłków wiodących zespołów badawczych oraz przedstawicieli społeczności uda się stymulować rozwój systemów uczenia maszynowego oraz rozwiązać problem fragmentacji infrastruktury dla różnych frameworków i sprzętu. OpenXLA umożliwia efektywne wsparcie różnorodnych urządzeń, niezależnie od tego, jaki framework był użyty do stworzenia modelu uczenia maszynowego. Przewiduje się, że dzięki OpenXLA czas uczenia modeli zostanie skrócony, wydajność wzrośnie, opóźnienia się zmniejszą, a koszty zasobów obliczeniowych zostaną zredukowane, co przyspieszy wprowadzenie produktów na rynek.
OpenXLA składa się z trzech podstawowych komponentów, których kod jest rozpowszechniany na licencji Apache 2.0:
- XLA (Accelerated Linear Algebra) — kompilator, który umożliwia optymalizację modeli uczenia maszynowego do wysokowydajnego wykonywania na różnych platformach sprzętowych, w tym GPU, CPU oraz specjalizowanych akceleratorach różnych producentów.
- StableHLO — specyfikacja i podstawowa implementacja zestawu wysokopoziomowych operacji (HLO, High-Level Operations) do wykorzystania w modelach systemów uczenia maszynowego. Działa jako warstwa pośrednia między frameworkami uczenia maszynowego a kompilatorami, które przekształcają model do wykonania na określonym sprzęcie. Warstwy do generowania modeli w formacie StableHLO przygotowano dla frameworków PyTorch, TensorFlow i JAX. Jako podstawę dla StableHLO wykorzystano zestaw MHLO, który został rozszerzony o wsparcie dla serializacji i wersjonowania.
- IREE (Środowisko Wykonawcze Reprezentacji Pośredniej) — kompilator i runtime, przekształcający modele uczenia maszynowego w uniwersalną reprezentację pośrednią, opartą na formacie MLIR (Wielopoziomowa Reprezentacja Pośrednia) projektu LLVM. Do zalet należy możliwość kompilacji wstępnej (ahead-of-time), wsparcie dla zarządzania przepływem, możliwość użycia dynamicznych elementów w modelach oraz optymalizacja dla różnych CPU i GPU przy niskich kosztach ogólnych.
Główne zalety narzędzi OpenXLA:
- Osiągnięcie optymalnej wydajności bez konieczności zagłębiania się w pisanie kodu specyficznego dla określonych urządzeń. Dostarczanie gotowych optymalizacji, w tym uproszczenie wyrażeń algebraicznych, efektywne rozmieszczenie w pamięci oraz planowanie wykonania z uwzględnieniem redukcji szczytowego zużycia pamięci i przeciążeń.
- Uproszczenie skalowania i równoleglenia obliczeń. Programista musi jedynie dodać adnotacje dla podzbioru krytycznych tensorów, na podstawie których kompilator może automatycznie wygenerować kod dla obliczeń równoległych.
- Zagwarantowanie przenośności dzięki wsparciu różnych platform sprzętowych, takich jak GPU AMD i NVIDIA, CPU o architekturze x86 i ARM, akceleratory ML TPU Google, IPU AWS Trainium Inferentia, Graphcore oraz Cerebras Wafer-Scale Engine.
- Wsparcie dla dodawania rozszerzeń z realizacją dodatkowych możliwości, takich jak wsparcie pisania prymitywów głębokiego uczenia maszynowego z użyciem CUDA, HIP, SYCL, Triton i innych języków do obliczeń równoległych. Możliwość ręcznego dostrajania wąskich gardeł w modelach.
Źródło: opennet.ru
