Co naprawdę napędza sieć neuronową: kompilatory AI w Pythonie

Każdy, kto ma doświadczenie w kompilowaniu i podstawową wiedzę z zakresu uczenia maszynowego, być może nigdy nie powiązał tych dwóch dziedzin. Lexery, parserzy, drzewa atrybutów oraz LLVM IR należą do jednego świata, natomiast sieci neuronowe do innego. Spotykają się one w kompilatorach do AI/ML, a pierwszym zaskoczeniem jest to, że model uczenia głębokiego tak naprawdę nie działa na Pythonie.

To brzmi dziwnie, ponieważ linia kodu taka jak poniżej sprawia wrażenie, jakby Python robił wszystko:

output = model(input)

Python sam w sobie nie wykonywa miliardów mnożeń macierzy. Pod spodem odbywa się znacznie bardziej złożony proces.

Od wzoru do instrukcji sprzętowych

Wyobraź sobie małą sieć, która przepuszcza dane od wejścia, przez mnożenie macierzy, a dalej:

Input
  ↓
Matrix Multiplication
  ↓
ReLU
  ↓
Matrix Multiplication
  ↓
Output

W Pythonie jego rdzeń może wyglądać w ten sposób:

y = relu(x @ W + b)

Wygląda to prosto, ale pomyśl o tym, czego potrzebuje maszyna. CPU lub GPU nie rozumie tej funkcji:

relu()

lub tego operatora:

x @ W

jak występują w kodzie. Sprzęt wykonywa instrukcje. Gdzieś opis modelu na wysokim poziomie musi zostać przekształcony w coś, co chip może efektywnie uruchomić, a tą transformacją zajmują się kompilatory AI/ML.

Uproszczony obraz łańcucha, od modelu aż po procesor:

AI Model
   ↓
PyTorch / TensorFlow
   ↓
Computational Graph
   ↓
Intermediate Representation (IR)
   ↓
Optimisations
   ↓
Lowering
   ↓
Hardware-specific Code
   ↓
Machine Instructions
   ↓
CPU / GPU / NPU

To dużo etapów, więc warto podejmować się ich po jednym.

Co robi ta ramowa struktura

W PyTorch ta sama idea jest zapisywana w ten sposób:

y = torch.relu(x @ W + b)

Rozwijający nie chce zajmować się rejestrami GPU, układem pamięci, planowaniem instrukcji ani tysiącami wątków sprzętowych. Określają jedynie, jakie obliczenia chcą przeprowadzić. Python i PyTorch oferują wygodny sposób na ich sformułowanie, natomiast coś innego musi przenieść tę opis na poziom sprzętu.

Graf obliczeniowy

Sieć neuronową można postrzegać jako graf obliczeniowy. Wyrażenie:

y = relu(x @ W + b)

Można to przedstawić jako graf połączonych operacji:

Teraz kompilator ma coś konkretnego do opracowania. Gdy operacje stają się widoczne, może sprawdzić, czy obliczenia można przyspieszyć. W tym momencie optymalizacja staje się kluczowa.

Fuzja operatorów i ruch pamięci

Weźmy model zawierający łańcuch w takim kształcie:

MatMul
   ↓
Add
   ↓
ReLU

Kompilator może zauważyć, że te kroki można połączyć lub dostosować do celowej architektury sprzętowej. Kluczową techniką jest fuzja operatorów: zamiast traktować każdą operację jako oddzielny krok, niektóre z nich są łączone, co redukuje niepotrzebną pracę, szczególnie przesyłanie danych pomiędzy pamięcią a jednostkami obliczeniowymi.

To ma znaczenie, ponieważ przesyłanie danych jest kosztowne w nowoczesnych zadaniami AI. Wąskim gardłem nie jest tylko szybkość obliczania wyniku, ale także efektywność, z jaką dane trafiają do obliczeń. To jeden z powodów, dla których ta dziedzina jest interesująca.

Reprezentacje pośrednie

Następnie pojawia się koncepcja znana z klasycznych kursów z zakresu kompilatorów – reprezentacja pośrednia (IR). Tradycyjny kompilator rzadko przechodzi bezpośrednio od tego etapu:

Source Code

do tego:

Machine Code

Przechodzi przez IR-y. Kompilatory AI stosują tę samą zasadę, ale ich reprezentacje muszą odzwierciedlać koncepcje ważne dla uczenia maszynowego. Gdzie zwykły kompilator myśli o:

variables
functions
loops

kompilator ML musi również reprezentować takie elementy jak:

tensors
shapes
matrix operations
convolutions
attention
memory layouts

To właśnie sprawia, że projekty takie jak MLIR są godne uwagi. MLIR dostarcza infrastrukturę do reprezentowania i przekształcania obliczeń na różnych poziomach abstrakcji.

Stopniowe obniżanie poziomu abstrakcji

Kompilator stopniowo przechodzi od form wysokiego poziomu do tych, które są bliższe docelowemu sprzętowi. Wyobraź sobie schodzenie po schodach:

High-level ML operation
          ↓
Tensor operation
          ↓
Lower-level operation
          ↓
Hardware-oriented operation
          ↓
Machine instructions

Na każdym poziomie abstrakcyjne idee stają się nieco bardziej konkretne, aż kompilator wygeneruje coś, co może być wykorzystane przez docelowe sprzęty.

Sprzęt u podstaw

Teraz docieramy do tej części, o której nikt nie myśli podczas pisania:

model(input)

Ostateczne obliczenia mogą być wykonywane na CPU, GPU, NPU, TPU lub innym akceleratorze AI. Żaden z nich nie przejmuje się tym, co stwierdzał oryginalny kod:

y = relu(x @ W + b)

Wymagane są wykonywalne instrukcje oraz dane uporządkowane w sposób odpowiedni dla danego urządzenia. Karta graficzna może potrzebować kodu specyficznego dla niej oraz modeli wykonywania. Procesor korzysta z instrukcji procesora oraz technik wektorowania. Akcelerator może wymagać zupełnie innego backendu. Dlatego ten sam model wysokiego poziomu może być skompilowany inaczej w zależności od celu jego użycia.

Więc gdzie mieści się Python?

Python nie zniknął i nadal jest bardzo przydatny. Służy jako interfejs do definiowania modeli, kontrolowania ich wykonywania, ładowania danych oraz tworzenia aplikacji. Po prostu nie musi samodzielnie wykonywać każdej operacji numerycznej. Uproszczony model mentalny:

To zmienia obraz sztucznej inteligencji. Patrząc tylko z perspektywy modelu, obraz jest prosty:

Data → Model → Prediction

Dodaj kompilator, a pojawia się kolejna warstwa pomiędzy nimi:

Model
  ↓
Compiler
  ↓
Hardware

Model określa, jakie obliczenia muszą zostać przeprowadzone. Kompilator ustala, jak je przedstawić i zoptymalizować. Sprzęt wykonuje te obliczenia.

Dlaczego potrzebna jest dodatkowa warstwa

Skoro istnieją już frameworki, po co dodawać warstwę kompilatora? Ponieważ sprzęt do sztucznej inteligencji staje się coraz bardziej wyspecjalizowany. Istnieją procesory CPU, karty graficzne GPU, jednostki TPU, NPU, przyspieszacze inferencji oraz specjalistyczne chipy, z których każdy ma inne właściwości. Model, który dobrze działa w jednej architekturze, może wymagać zupełnie innych ustawień w innej. Kompilator łączy te dwa aspekty. Jednym z nich jest pytanie o to, co model chce obliczyć:

WHAT
The ML model wants to compute

Druga kwestia dotyczy tego, w jaki sposób dany urządzenie powinien to robić efektywnie:

HOW
A specific piece of hardware should execute it efficiently

Ten most zyskuje coraz większe znaczenie.

Znana filozofia w nowej dziedzinie

Każdy, kto zna tradycyjne kompilatory, rozpoznaje podstawową architekturę przetwarzania:

Source Code
    ↓
Lexer
    ↓
Parser
    ↓
AST
    ↓
IR
    ↓
Optimisation
    ↓
Machine Code

Machine learning stosuje podobną filozofię w modelach:

ML Model
    ↓
ML Representation
    ↓
ML IR
    ↓
Optimisation
    ↓
Lowering
    ↓
Hardware Code

Domena jest inna, a problemy znacznie trudniejsze, jednak sposób myślenia pozostaje ten sam: reprezentacja, transformacja, optymalizacja, obniżenie kosztów, wykonywanie.

Naturalną następną ideą jest stworzenie małego kompilatora AI, aby zrozumieć, jak on działa. Nie miałby on konkuruje z LLVM, XLA czy MLIR, a celem byłoby stworzenie czegoś na tyle małego, by można go było zrozumieć, być może rozpoczynając od prostego języka ML:

Simple ML Language
        ↓
       Parser
        ↓
      ML AST
        ↓
       ML IR
        ↓
   Simple Optimiser
        ↓
    Lowering
        ↓
   LLVM / Backend

Stos w jednej linii

Sztuczna inteligencja zazwyczaj jest dostępna poprzez bardzo prosty interfejs:

model(input)

Pod tym jednym wywołaniem kryje się cała warstwa oprogramowania i sprzętu: frameworki, grafy, reprezentacje pośrednie, kompilatory, optymalizacje, backendy, sprzęt oraz wreszcie instrukcje maszynowe. Uruchomienie modelu oznacza uruchomienie całego ekosystemu kompilatorów i sprzętu, przy czym Python pełni tu jedynie rolę warstwy powierzchniowej.