Neuronale Netze direkt auf dem Apple Neural Engine trainieren ohne CoreML und GPU

Jeder Apple-Silicon-Prozessor enthält einen dedizierten Block für neuronale Netzwerkoperationen – den Apple Neural Engine (ANE). Der Hersteller verspricht Tens of Teraflops an Leistung, aber offizielle Tools beschränken dessen Nutzung auf Inferenz über CoreML. Das Training von Modellen wird entweder auf der CPU oder dem Grafikchip angeboten.

Ein Entwickler mit dem Nickname maderix beschloss zu testen, ob die ANE-Hardware physisch nicht zu Backpropagation fähig ist, oder ob das Problem nur an den Software-Einschränkungen von Apple liegt. Innerhalb eines Wochenendes sezierte er die privaten macOS-Frameworks und brachte den ANE dazu, nicht nur Forward-, sondern auch Backward-Pass für Transformer auszuführen.

Das Projekt sammelte über 7.000 Sterne auf GitHub. Lassen Sie uns aufschlüsseln, wie dieser Hack funktioniert, welche Tricks angewendet werden mussten und wie die tatsächliche Leistung aussieht.

Warum in geschlossene Frameworks eintauchen

Der offizielle Stack wie CoreML erzeugt den Eindruck, dass ANE eine geschlossene Blackbox ist. Man gibt ihm ein fertiges Modell, es gibt das Ergebnis aus. Wenn man auch nur ein kleines Modell direkt auf dem Client trainieren musste, musste man auf Metal oder das MLX-Framework zurückgreifen und die GPU laden.

Der Autor des ANE-Repositorys demonstrierte, dass der Chip durchaus in der Lage ist, beliebige Rechengraphen auszuführen. Dazu reverse-engineerte er die privaten Bibliotheken _ANEClient und _ANECompiler sowie die interne Modellbeschreibungssprache MIL (Model Intermediate Language).

Das Ergebnis war interessant: vollwertiges Training von Transformern ohne eine einzige Zeile CoreML oder Metal. Der Modelltext wird direkt im RAM assembliert, zur Laufzeit kompiliert und an den neuronalen Prozessor gesendet.

Einschränkungen und rauhe Realität

Bevor Sie Ihre Trainingsskripte auf dem Mac umschreiben, lohnt es sich, ehrliche Zahlen zu betrachten. Der Autor warnt selbst vorab: Dies ist ein akademisches Experiment, keine produktionsreife Bibliothek.

Eine 100-prozentige Auslastung der ANE-Ressourcen war noch nicht möglich. Die tatsächliche Chip-Auslastung liegt bei etwa 5–9 % des Peaks. Software- und Hardware-Einschränkungen machen sich bemerkbar:

• Einige mathematische Operationen werden vom Chip nicht in der benötigten Form unterstützt und fallen auf die CPU zurück.
• Der Backward-Pass für Gewichte (dW) muss weiterhin vom Prozessor berechnet werden.
• Der ANE-Compiler hat Speicherlecks, daher sind nach etwa hundert Iterationen Workarounds erforderlich.

Nichtsdestotrotz liefert das Projekt in diesem Modus eine ordentliche Geschwindigkeit bei grundlegenden Architekturen.

So funktioniert die Training-Pipeline

Die Projektarchitektur basiert auf der Verteilung von Aufgaben zwischen ANE und CPU. Der neuronale Netzwerkbeschleuniger übernimmt die schwersten Matrixmultiplikationen, während der Prozessor sich um die umgebende Logik und Gradientenakkumulation kümmert.

Der Forward-Pass und die Berechnung der Eingabegradienten (dx) werden vollständig auf dem ANE ausgeführt. Gewichtsgradienten (dW) werden von der CPU durch optimierte Accelerate- und cblas_sgemm-Bibliotheken berechnet. Der Adam-Optimizer und die RMSNorm-Schicht laufen ebenfalls auf dem Prozessor.

Um eine Neukompilierung des Modellgraphen bei jedem Schritt zu vermeiden, wenn sich Gewichte ändern, wandte der Autor einen Trick an: Gewichte und Aktivierungen werden durch räumliche Dimensionen zu einem einzigen Tensor zusammengepackt, und innerhalb des MIL-Kernels werden sie einfach wieder aufgeteilt.

IOSurface-Speicher wird für den Datenaustausch zwischen CPU und ANE verwendet. Dies ermöglicht die Übertragung von Tensoren ohne unnötiges Kopieren zwischen Adressräumen. Daten werden in das ANE-spezifische Format [1, C, 1, S] gepackt, wo Kanäle zuerst kommen. Dieser Ansatz eliminierte den Transpose-Matrix-Overhead.

Nicht-offensichtliche Probleme und Workarounds

Viele Apple-Hardware-Fallstricke tauchten beim Reverse Engineering auf.

Erstens ignoriert die SDPA-Operation (Scaled Dot-Product Attention) im ANE die kausale Maske attn_mask auf Hardwareebene. Der Attention-Mechanismus musste in drei Stufen aufgeteilt werden: Multiplikation von Q und K auf ANE, Maskierung mit Softmax auf CPU und finale Multiplikation mit V zurück auf ANE.

Zweitens enthält der eingebaute Compiler _ANECompiler ein Speicherleck. Nach ungefähr 119 Kompilierungen stürzt der Prozess wegen Ressourcenerschöpfung ab. Der Autor löste das Problem radikal: wenn der Zähler sich der Grenze nähert, speichert das Programm einen Checkpoint und macht exec() — startet sich selbst mit Zustandserhaltung neu.

Drittens führen FP16-Berechnungen während des Backward-Passes schnell zu Underflow, was dazu führt, dass Gradienten zu Nullen werden. Das Problem wurde durch Skalierung des Losses mit Koeffizient 256 * NLAYERS gelöst.

Leistung auf dem M4

Auf dem Apple M4 Chip waren die Ergebnisse ziemlich aufschlussreich. Tests wurden auf zwei Architekturen durchgeführt:

Für das Stories110M-Modell mit 109 Millionen Parametern (12 Schichten, klassische Multi-Head Attention) betrug die Zeit für einen Trainingsschritt 91 Millisekunden.

Das größere Qwen3-0.6B mit 596 Millionen Parametern und Grouped-Query Attention verarbeitet einen Schritt in 412 Millisekunden.

Der Autor testete auch INT8 W8A8-Quantisierung. Die Verwendung von 8-Bit-Gewichten und -Aktivierungen reduziert die Belastung des L2-SRAM-Speichers des Chips und erhöht den Durchsatz von 18,6 TOPS auf 35,1 TOPS auf dem M4 – eine Beschleunigung von fast 1,88x im Vergleich zu FP16.

So führen Sie das Projekt aus

Das Projekt erfordert keine externen Abhängigkeiten wie PyTorch oder Conda. Alles, was Sie brauchen, ist das neueste macOS 15 auf einer Maschine mit Apple Silicon und dem Clang-Compiler. Private APIs werden zur Laufzeit über objc_msgSend eingebunden.

Um die dynamische Pipeline zu erstellen, navigieren Sie einfach zum Projektordner und führen Sie den Make-Befehl aus:

cd training/training_dynamic
make MODEL=stories110m
./train --scratch

Wenn Sie INT8-Quantisierung ausprobieren oder die Peak-TOPS Ihres Chips benchmarken möchten, gibt es separate Benchmarks im Stammverzeichnis des Repositorys.

Abschließende Gedanken

Das Projekt von maderix ist ein Beispiel für qualitativ hochwertige Recherche „unter der Haube