Jak przekształcić znajomy NumPy w kompilowalny kod GPU

Jeśli kiedykolwiek próbowałeś przyspieszyć intensywne obliczenia w Pythonie, prawdopodobnie natknąłeś się na ograniczenia standardowego stosu. NumPy jest szybki dzięki swojemu backendowi w C, ale nie może domyślnie pracować z GPU. PyTorch i TensorFlow rozwiązują ten problem, ale wprowadzają rozbudowane abstrakcje, klasy warstw i własną semantykę grafów obliczeniowych.

W 2018 roku inżynierowie z Google udostępnili bibliotekę JAX jako open source. Główna koncepcja jest prosta: zapewnić programistom znajomą składnię NumPy, wzbogaconą o automatyczne różniczkowanie i kompilator XLA. Efektem jest czysty, funkcyjny kod Python, który jest kompilowany w locie na zoptymalizowany kod maszynowy dla GPU lub TPU.

Czym JAX tak naprawdę jest

Wiele osób postrzega JAX jako kolejny framework ML, ale to nie do końca oddaje istoty. Twórcy sami określają go w repozytorium jako system komponowalnych transformacji funkcyjnych dla tablic.

Zamiast budować złożone modele obiektowe, JAX zachęca do pracy z czystymi funkcjami. Piszesz zwykły kod Python, a następnie stosujesz do niego funkcje transformujące. Bez ukrytych globalnych stanów czy mutacji danych w miejscu. Jeśli potrzebujesz obliczyć gradient, skompilować fragment programu lub rozproszyć obliczenia między wsadami, wystarczy opakować funkcję odpowiednim dekoratorem.

Przyjrzyjmy się teraz czterem głównym transformacjom, na których opiera się cała biblioteka.

Automatyczne różniczkowanie za pomocą grad

Funkcja jax.grad przyjmuje twoją funkcję i zwraca nową, która oblicza gradient oryginału:

import jax
import jax.numpy as jnp

def tanh(x):
    y = jnp.exp(-2.0 * x)
    return (1.0 - y) / (1.0 + y)

grad_tanh = jax.grad(tanh)
print(grad_tanh(1.0))  # 0.4199743

Możesz obliczać gradienty dowolnego rzędu, po prostu zagnieżdżając wywołania jax.grad. Algorytm obsługuje standardowe warunki Python if/else, pętle i rekursję bez problemów.

Kompilacja przez jit

Zwykły Python wykonuje każdą operację tablicową sekwencyjnie, z narzutem wywołań funkcji i pośredniej alokacji pamięci. Dekorator jax.jit wysyła ciało funkcji do kompilatora XLA:

def calculate(x):
    return x * x + x * 2.0

x = jnp.ones((5000, 5000))

# Обычный запуск против скомпилированного
fast_calculate = jax.jit(calculate)

Kompilator łączy podstawowe operacje w pojedynczy kernel obliczeniowy. W rezultacie dane nie przemieszczają się tam i z powrotem między pamięcią podręczną a pamięcią GPU, a wydajność wzrasta o rzędy wielkości.

WektorYZacja kodu przez vmap

Każdy, kto pisał algorytmy uczenia maszynowego, spędzał godziny dopasowując wymiary tensorów do rozmiaru wsadu. jax.vmap rozwiązuje ten problem: piszesz logikę dla pojedynczego elementu lub wektora, a JAX sam wektoryzuje operację:

def l1_distance(x, y):
    return jnp.sum(jnp.abs(x - y))

# Превращаем функцию для векторов в функцию для матриц
def pairwise_distances(xs):
    return jax.vmap(jax.vmap(l1_distance, (0, None)), (None, 0))(xs, xs)

xs = jax.random.normal(jax.random.key(0), (100, 3))
matrix = pairwise_distances(xs)  # форма (100, 100)

Zamiast wolnej pętli Python, wektoryzator przenosi pętlę do niskopoziomowych operacji, zamieniając mnożenia macierzowo-wektorowe w pełne mnożenia macierzowe.

Równoległość i dzielenie danych

Gdy model nie mieści się już w pamięci pojedynczego akceleratora, JAX oferuje deklaratywne podejście do równoległości. Definiujesz siatkę urządzeń i reguły partycjonowania tablicy (specyfikację partycji), a kompilator samodzielnie dystrybuuje obliczenia i konfiguruje wymianę danych między kartami:

from jax.sharding import set_mesh, AxisType, PartitionSpec as P

# Создаем сетку из 8 ускорителей
mesh = jax.make_mesh((8,), ('data',), axis_types=(AxisType.Explicit,))
set_mesh(mesh)

# Шардируем входные данные
inputs, targets = jax.device_put((inputs, targets), P('data'))

# Обычная функция градиента теперь выполняется параллельно
grad_fn = jax.jit(jax.grad(loss_fn))
grads = grad_fn(params, (inputs, targets))

Zastrzeżenia i specyfika

JAX ma drugą stronę, do której trzeba się przyzwyczaić.

Po pierwsze, paradygmat funkcyjny wymaga braku efektów ubocznych. Nie możesz po prostu zmienić elementu tablicy przez indeks (arr[0] = 5), ponieważ tablice w JAX są niezmienne. W tym celu użyj metody arr.at[0].set(5).

Po drugie, generowanie liczb losowych wymaga jawnego przekazywania kluczy stanu (jax.random.key), ponieważ globalne ziarno zakłóciłoby powtarzalność podczas równoległej kompilacji.

Po trzecie, debugowanie kodu JIT może być nietypowe: przy pierwszym wywołaniu funkcja przechodzi przez fazę śledzenia, a zwykłe instrukcje print w Pythonie wewnątrz niej wykonają się tylko raz.

Instalacja i platformy

Biblioteka oficjalnie obsługuje Linuksa i macOS, a także działa na Windowsie przez podsystem WSL2.

Do uruchamiania na zwykłym CPU:

pip install -U jax

Do budowania z obsługą NVIDIA CUDA:

pip install -U "jax[cuda13]"

Jest też wsparcie dla akceleratorów Google TPU i AMD ROCm.

Czy warto spróbować

JAX świetnie sprawdza się w zadaniach badawczych, modelowaniu fizyki, niestandardowych optymalizacjach i obliczeniach naukowych, gdzie PyTorch wydaje się zbyt rozbudowany, a zwykły NumPy niewystarczająco szybki. Jeśli twój projekt napotyka limity wydajności na operacjach matematycznych lub wymaga obliczania złożonych pochodnych, zdecydowanie warto przyjrzeć się jax-ml/jax.