Vertrautes NumPy in kompilierbaren GPU-Code umwandeln

Wenn Sie schon einmal versucht haben, rechenintensive Operationen in Python zu beschleunigen, sind Sie wahrscheinlich an die Grenzen des Standard-Stacks gestoßen. NumPy ist dank seines C-Backends schnell, kann aber nicht nativ mit GPUs arbeiten. PyTorch und TensorFlow lösen dieses Problem, bringen aber sperrige Abstraktionen, Layer-Klassen und ihre eigene Semantik für Berechnungsgraphen mit.

Im Jahr 2018 haben Ingenieure bei Google die JAX-Bibliothek als Open Source veröffentlicht. Die Idee dahinter ist einfach: Entwicklern eine vertraute NumPy-Syntax geben, aber automatische Differentiation und den XLA-Compiler oben drauf setzen. Das Ergebnis ist sauberer funktionaler Python-Code, der on-the-fly in optimierten Maschinencode für GPUs oder TPUs kompiliert wird.

Was JAX eigentlich ist

Viele betrachten JAX als nur ein weiteres ML-Framework, aber das stimmt nicht ganz. Die Autoren selbst beschreiben es im Repository als ein System aus zusammensetzbaren funktionalen Transformationen für Arrays.

Statt komplexer Objektmodelle fördert JAX die Arbeit mit reinen Funktionen. Sie schreiben normalen Python-Code und wenden dann Transformationsfunktionen darauf an. Keine versteckten globalen Zustände oder In-Place-Datenmutationen. Wenn Sie einen Gradienten berechnen, ein Programmstück kompilieren oder Berechnungen über Batches verteilen müssen, wickeln Sie die Funktion einfach in den passenden Decorator.

Schauen wir uns die vier Haupttransformationen an, auf denen die gesamte Bibliothek aufbaut.

Automatische Differentiation über grad

Die jax.grad Funktion nimmt Ihre Funktion und gibt eine neue zurück, die den Gradienten der ursprünglichen berechnet:

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

Sie können Gradienten beliebiger Ordnung berechnen, indem Sie einfach jax.grad Aufrufe verschachteln. Der Algorithmus bewältigt Standard-Python-Bedingungen if/else, Schleifen und Rekursion problemlos.

Kompilierung über jit

Normales Python führt jeden Array-Operation sequenziell aus, mit Overhead durch Funktionsaufrufe und Zwischenallokationen. Der jax.jit Decorator sendet den Funktionskörper an den XLA-Compiler:

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

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

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

Der Compiler verschmilzt elementare Operationen zu einem einzigen Compute-Kernel. Dadurch springt data nicht mehr zwischen Cache und GPU-Speicher hin und her, und die Performance steigt um Größenordnungen.

Code-Vektorisierung über vmap

Jeder, der Machine-Learning-Algorithmen geschrieben hat, hat Stunden damit verbracht, Tensor-Dimensionen an die Batch-Größe anzupassen. jax.vmap löst dieses Problem: Sie schreiben Logik für ein einzelnes Element oder einen Vektor, und JAX vektorisiert die Operation selbst:

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)

Statt einer langsamen Python-Schleife schiebt der Vektorisierer die Schleife in Low-Level-Operationen hinein und verwandelt Matrix-Vektor-Multiplikationen in vollständige Matrix-Multiplikationen.

Parallelität und Data Sharding

Wenn ein Modell nicht mehr in den Speicher eines einzelnen Accelerators passt, bietet JAX einen deklarativen Ansatz für Parallelität. Sie definieren ein Device-Mesh und Array-Partitionierungsregeln (Partition Spec), und der Compiler selbst verteilt Berechnungen und konfiguriert den Datenaustausch zwischen den Karten:

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))

Fallstricke und Besonderheiten

JAX hat eine Kehrseite, an die man sich gewöhnen muss.

Erstens erfordert das funktionale Paradigma keine Nebeneffekte. Sie können nicht einfach ein Array-Element per Index ändern (arr[0] = 5), weil Arrays in JAX immutable sind. Verwenden Sie dafür die arr.at[0].set(5) Methode.

Zweitens erfordert die Generierung von Zufallszahlen explizites Übergeben von State-Keys (jax.random.key), weil ein globaler Seed die Reproduzierbarkeit während paralleler Kompilierung brechen würde.

Drittens kann das Debugging von JIT-Code ungewohnt sein: Beim ersten Aufruf durchläuft die Funktion eine Tracing-Phase, und reguläre Python print-Anweisungen darin werden nur einmal ausgeführt.

Installation und Plattformen

Die Bibliothek unterstützt offiziell Linux und macOS und läuft auch unter Windows über das WSL2-Subsystem.

Für die Ausführung auf einer normalen CPU:

pip install -U jax

Für den Build mit NVIDIA CUDA-Unterstützung:

pip install -U "jax[cuda13]"

Es gibt auch Unterstützung für Google TPU-Acceleratoren und AMD ROCm.

Lohnt es sich, es zu probieren

JAX eignet sich hervorragend für Forschung, Physik-Modellierung, nicht-standard Optimierungen und wissenschaftliches Rechnen, wo sich PyTorch zu sperrig anfühlt und plain NumPy nicht schnell genug ist. Wenn Ihr Projekt bei mathematischen Operationen an Leistungsgrenzen stößt oder komplexe Ableitungen berechnen muss, lohnt es sich definitiv, einen Blick auf jax-ml/jax zu werfen.