OpenXLA to wspólnie opracowany ekosystem kompilatora ML o otwartym kodzie źródłowym
Ostatnio zaprezentowano największe firmy zajmujące się rozwojem w obszarze uczenia maszynowego projekt OtwórzXLA, przeznaczone do wspólnego opracowywania narzędzi kompilować i optymalizować modele dla systemów uczenia maszynowego.
Projekt zajął się opracowaniem narzędzi pozwalających na ujednolicenie kompilacji modeli przygotowanych we frameworkach TensorFlow, PyTorch i JAX w celu efektywnego uczenia i wykonywania na różnych GPU i wyspecjalizowanych akceleratorach. Do wspólnych prac nad projektem włączyły się takie firmy jak Google, NVIDIA, AMD, Intel, Meta, Apple, Arm, Alibaba czy Amazon.
Projekt OpenXLA zapewnia najnowocześniejszy kompilator ML, który można skalować pomimo złożoności infrastruktury ML. Jego podstawowymi filarami są wydajność, skalowalność, przenośność, elastyczność i rozszerzalność dla użytkowników. Wraz z OpenXLA dążymy do uwolnienia prawdziwego potencjału sztucznej inteligencji poprzez przyspieszenie jej rozwoju i dostarczania.
OpenXLA umożliwia programistom kompilację i optymalizację modeli ze wszystkich wiodących platform ML w celu wydajnego szkolenia i obsługi szerokiej gamy sprzętu. Deweloperzy korzystający z OpenXLA zauważą znaczną poprawę w zakresie czasu szkolenia, wydajności, opóźnień usług, a ostatecznie czasu wprowadzenia na rynek i kosztów obliczeniowych.
Mamy nadzieję, że łącząc wysiłki z głównych zespołów badawczych i przedstawicieli społeczności, możliwe będzie stymulowanie rozwoju systemów uczenia maszynowego i rozwiązać problem fragmentacji infrastruktury dla różnych frameworków i zespołów.
OpenXLA pozwala zaimplementować efektywną obsługę różnych urządzeń, niezależnie od frameworka, na którym oparty jest model uczenia maszynowego. Oczekuje się, że OpenXLA skróci czas szkolenia modeli, poprawi wydajność, zmniejszy opóźnienia, zmniejszy obciążenie obliczeniowe i skróci czas wprowadzania na rynek.
OtwórzXLA składa się z trzech głównych elementów, którego kod rozpowszechniany jest na licencji Apache 2.0:
- XLA (accelerated linear algebra) to kompilator, który umożliwia optymalizację modeli uczenia maszynowego pod kątem wydajnego wykonywania na różnych platformach sprzętowych, w tym na procesorach graficznych, procesorach i wyspecjalizowanych akceleratorach różnych producentów.
- StableHLO to podstawowa specyfikacja i implementacja zestawu operacji wysokiego poziomu (HLO) do użytku w modelach systemów uczenia maszynowego. Działa jako warstwa między platformami uczenia maszynowego a kompilatorami, które przekształcają model tak, aby działał na określonym sprzęcie. Warstwy przygotowane są do generowania modeli w formacie StableHLO dla frameworków PyTorch, TensorFlow oraz JAX. Pakiet MHLO jest używany jako podstawa dla StableHLO, które jest rozszerzone o obsługę serializacji i kontroli wersji.
- IREE (Intermediate Representation Execution Environment) to kompilator i środowisko wykonawcze, które konwertuje modele uczenia maszynowego na uniwersalną reprezentację pośrednią opartą na formacie MLIR (Intermediate Multi-Level Representation) projektu LLVM. Spośród funkcji wyróżniono możliwość prekompilacji (z wyprzedzeniem), obsługę kontroli przepływu, możliwość wykorzystania elementów dynamicznych w modelach, optymalizację pod kątem różnych procesorów i procesorów graficznych oraz niski narzut.
Jeśli chodzi o główne zalety OpenXLA, wspomniano, że optymalna wydajność została osiągnięta bez konieczności zagłębiania się w pisanie kodu specyficzne dla urządzenia, oprócz dostarczać gotowe optymalizacje, w tym upraszczanie wyrażeń algebraicznych, wydajną alokację pamięci, szeregowanie wykonywania z uwzględnieniem redukcji maksymalnego zużycia pamięci i narzutów.
Kolejną zaletą jest uproszczenie skalowania i zrównoleglenie obliczeń. Wystarczy, że programista doda adnotacje dla podzbioru krytycznych tensorów, na podstawie których kompilator może automatycznie wygenerować kod do obliczeń równoległych.
Podkreśla się również, że mobilność zapewnia wsparcie dla wielu platform sprzętowych, takie jak procesory graficzne AMD i NVIDIA, procesory x86 i ARM, akceleratory Google TPU ML, jednostki IPU AWS Trainium Inferentia, Graphcore i mózgi silników w skali płytki.
Obsługa łączenia rozszerzeń z realizacją dodatkowych funkcji, jako wsparcie dla pisania prymitywów głębokiego uczenia maszynowego przy użyciu CUDA, HIP, SYCL, Triton i innych języków do obliczeń równoległych, a także możliwość ręcznej regulacji wąskich gardeł w modelach.
Wreszcie, jeśli chcesz dowiedzieć się więcej na ten temat, możesz zapoznać się z szczegóły w poniższym linku.