Was wirklich ein neuronales Netz antreibt: KI-Kompilatoren hinter Python

Jeder mit Erfahrung in der Programmierung mit Compilern und grundlegendem Wissen über maschinelles Lernen hat diese beiden Bereiche vermutlich noch nie miteinander in Verbindung gebracht. Lexer, Parser, ASTs und LLVM IR gehören zur einen Welt, während neuronale Netze zur anderen gehören. Sie treffen in Compilern für KI/maschinelles Lernen aufeinander – und die erste Überraschung ist folgende: Ein tiefes Lernmodell läuft eigentlich nicht unter Python.

Das klingt seltsam, denn eine Zeile wie die folgende scheint zu zeigen, dass Python alles erledigt:

output = model(input)

Python führt jedoch nicht selbst Milliarden von Matrixmultiplikationen durch. Im Hintergrund findet ein weitaus komplexerer Prozess statt.

Von einer Formel zu Hardware-Anweisungen

Stellen Sie sich ein kleines Netzwerk vor, das von der Eingabe über eine Matrixmultiplikation weitergeleitet wird:

Input
  ↓
Matrix Multiplication
  ↓
ReLU
  ↓
Matrix Multiplication
  ↓
Output

In Python könnte der Kern wie folgt aussehen:

y = relu(x @ W + b)

Es sieht einfach aus, aber man sollte darüber nachdenken, was die Maschine benötigt. Eine CPU oder GPU hat kein Verständnis für diese Funktion:

relu()

oder dieses Operators:

x @ W

so wie sie im Code erscheinen. Die Hardware führt die Anweisungen aus. Irgendwo muss die hochwertige Beschreibung des Modells in etwas umgewandelt werden, das der Chip effizient ausführen kann, und diese Umwandlung ist die Aufgabe von KI/ML-Kompilatoren.

Eine vereinfachte Darstellung der Kette, vom Modell bis zum Prozessor:

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

Das sind viele Schritte, daher lohnt es sich, sie nacheinander durchzugehen.

Was das Framework leistet

Mit PyTorch wird die gleiche Idee so formuliert:

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

Ein Entwickler möchte sich nicht mit GPU-Registern, Speicherlayout, Instruktionsplanung oder Tausenden von Hardware-Threads beschäftigen. Er gibt an, welche Berechnung er durchführen möchte. Python und PyTorch bieten eine einfache Möglichkeit, dies auszudrücken, während etwas anderes notwendig ist, um diese Beschreibung dem Hardware-Setup anzunähern.

Der Berechnungsgraph

Ein neuronales Netzwerk kann als Berechnungsgraph betrachtet werden. Die Darstellung lautet:

y = relu(x @ W + b)

Es kann als Graph verbundener Operationen dargestellt werden:

Nun hat der Compiler etwas Konkretes, an dem er arbeiten kann. Sobald die Operationen sichtbar sind, kann er prüfen, ob die Berechnung beschleunigt werden kann. Hier wird die Optimierung zentral.

Operatorfusion und Speicherdurchsatz

Nehmen wir ein Modell mit einer solchen Kette:

MatMul
   ↓
Add
   ↓
ReLU

Ein Compiler kann feststellen, dass diese Schritte zusammengeführt oder an die Zielhardware angepasst werden können. Eine wichtige Technik ist die Operatorfusion: Anstatt jede Operation als eigenständigen Schritt zu behandeln, werden einige miteinander kombiniert, wodurch unnötige Arbeiten, insbesondere der Datentransport zwischen Speicher und Recheneinheiten, eingespart werden.

Das ist wichtig, weil der Datentransport in modernen KI-Arbeitslasten aufwändig ist. Das Engpassproblem liegt nicht nur darin, wie schnell ein Ergebnis berechnet werden kann, sondern auch darin, wie effizient die Eingaben zur Berechnung gelangen. Das ist einer der Gründe, warum dieses Feld interessant ist.

Intermediäre Darstellungen

Danach kommt ein Konzept, das aus klassischen Compilerkursen bekannt ist: die intermediale Darstellung (IR). Ein herkömmlicher Compiler geht selten direkt davon aus:

Source Code

zu diesem:

Machine Code

Es durchläuft IRs. KI-Kompilatoren folgen demselben Prinzip, doch ihre Darstellungen müssen Ideen erfassen, die für das maschinelle Lernen wichtig sind. Während ein herkömmlicher Kompilator darüber nachdenkt:

variables
functions
loops

Ein ML-Kompilator muss auch Elemente wie folgende darstellen:

tensors
shapes
matrix operations
convolutions
attention
memory layouts

Das ist es, was Projekte wie MLIR bemerkenswert macht. MLIR stellt Infrastruktur zur Darstellung und Umwandlung von Berechnungen auf verschiedenen Abstraktionsebenen bereit.

Schritt für Schritt nach unten

Der Compiler wechselt allmählich von einer Hochsprache zu Formen, die der Zielhardware näher kommen. Stellen Sie sich vor, Sie steigen eine Treppe hinunter:

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

Auf jeder Ebene werden abstrakte Ideen etwas konkreter, bis der Compiler etwas erzeugt, das die Zielhardware verwenden kann.

Die Hardware am Grund

Nun kommen wir zu dem Teil, an den niemand beim Schreiben denkt:

model(input)

Die endgültige Berechnung kann auf einer CPU, GPU, NPU, TPU oder einem anderen KI-Accelerator durchgeführt werden. Keiner von ihnen kümmert sich darum, was im ursprünglichen Code stand:

y = relu(x @ W + b)

Dafür benötigen sie ausführbare Anweisungen sowie Daten, die für das jeweilige Gerät sinnvoll strukturiert sind. Eine GPU kann Code sowie Ausführungsmodelle benötigen, die speziell für sie entwickelt wurden. Eine CPU verwendet CPU-Anweisungen sowie Vektorisierung. Ein Accelerator wiederum benötigt möglicherweise einen völlig anderen Backend. Derselbe Hochleistungsmodell kann daher je nach Zielgerät unterschiedlich kompiliert werden.

Wo passt Python da hinein?

Python ist nicht verschwunden und bleibt weiterhin sehr nützlich. Es dient als Schnittstelle zur Definition von Modellen, zur Steuerung der Ausführung, zum Laden von Daten sowie zum Erstellen von Anwendungen. Es führt jedoch nicht unbedingt jede numerische Operation selbst aus. Ein vereinfachtes mentales Modell:

Das verändert das Bild der KI. Betrachtet man nur die Modellseite, ist das Bild einfach:

Data → Model → Prediction

Fügen Sie den Compiler hinzu, dann gibt es eine weitere Schicht dazwischen:

Model
  ↓
Compiler
  ↓
Hardware

Das Modell gibt an, welche Berechnungen durchgeführt werden müssen. Der Compiler ermittelt, wie diese dargestellt und optimiert werden können. Die Hardware führt die Ausführung durch.

Warum eine weitere Schicht notwendig ist

Wenn bereits Frameworks vorhanden sind, warum dann eine Compiler-Schicht hinzufügen? Weil KI-Hardware immer spezialisierter wird. Es gibt CPUs, GPUs, TPUs, NPUs, Inferenzbeschleuniger sowie maßgeschneiderte Chips – jeweils mit unterschiedlichen Eigenschaften. Ein Modell, das auf einer Architektur gut läuft, benötigt möglicherweise eine völlig andere Anpassung auf einer anderen Architektur. Der Compiler verbindet diese beiden Seiten. Eine davon ist die Frage, was das Modell berechnen möchte:

WHAT
The ML model wants to compute

Der andere Aspekt betrifft die Art und Weise, wie ein bestimmtes Gerät dies effizient bewältigen sollte:

HOW
A specific piece of hardware should execute it efficiently

Diese Brücke gewinnt ständig an Bedeutung.

Eine vertraute Philosophie in einem neuen Bereich

Jeder, der sich mit traditionellen Compilern auskennt, erkennt den grundlegenden Pipeline-Ablauf:

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

Das maschinelle Lernen wendet eine ähnliche Philosophie auf Modelle an:

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

Die Domäne unterscheidet sich und die Probleme sind weitaus schwieriger, doch die Denkweise bleibt dieselbe: darstellen, transformieren, optimieren, verringern, ausführen.

Eine natürliche nächste Idee ist es, einen kleinen AI-Compiler zu entwickeln, um zu verstehen, wie er funktioniert. Er würde nicht mit LLVM, XLA oder MLIR konkurrieren, und das Ziel wäre lediglich etwas Kleines genug, um es zu verstehen – vielleicht beginnend mit einer einfachen ML-Sprache:

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

Der Stack in einer Zeile

KI zeigt sich in der Regel über eine sehr einfache Benutzeroberfläche:

model(input)

Unter diesem einzigen Aufruf verbirgt sich eine ganze Software- und Hardware-Stack-Struktur: Frameworks, Graphen, Zwischenrepräsentationen, Compiler, Optimierungen, Backend-Systeme, Hardware sowie schließlich Maschinenanweisungen. Das Ausführen eines Modells bedeutet das Betreiben eines vollständigen Compiler- und Hardware-Ökosystems, wobei Python lediglich die Oberfläche bildet.