ANE — Training von neuronalen Netzen auf der Apple Neural Engine

GitHub: https://github.com/maderix/ANE

🔍 Übersicht

ANE ist ein Forschungsprojekt, das erstmals öffentlich demonstriert, dass das vollständige Training neuronaler Netze (Forward + Backward Pass) direkt auf Apples Neural Engine (ANE) — dem dedizierten KI-Beschleuniger, der in Apple Silicon Chips (M1/M2/M4-Serie) integriert ist — ausgeführt werden kann.

Apple beschränkt die ANE offiziell auf die ausschließliche Verwendung für Inferenz über das CoreML-Framework. Dieses Projekt reverse-engineert Apples private APIs (_ANEClient, _ANECompiler, _ANEInMemoryModelDescriptor), um diese Einschränkung zu umgehen und benutzerdefinierte Berechnungsdiagramme — einschließlich Backpropagation — nativ auf der ANE-Hardware auszuführen.

Es handelt sich nicht um ein Produktions-Framework. Es ist ein Proof of Concept und eine Benchmark-Referenz, die beweist, dass die Einschränkung softwareseitig auferlegt ist und keine Hardwarebeschränkung darstellt.

🎯 Warum das wichtig ist

Apple Silicon Chips (insbesondere M4) enthalten eine ANE mit einer Leistung von 15,8 TFLOPS — eine riesige Menge an dedizierter KI-Rechenleistung, die Apple auf Inferenz beschränkt. Jedes Mal, wenn Sie CoreML ausführen, erhalten Sie nur Inferenz. Das Training fällt immer auf die GPU oder CPU zurück.

Dieses Projekt beantwortet die Frage: "Kann man überhaupt auf der ANE trainieren?"

Die Antwort ist ja — und dieses Repository ist der Beweis.

⚙️ Funktionsweise

Das Projekt implementiert eine vollständige Trainingsschleife für Transformer-Schichten durch:

MIL-Programmgenerierung — Erstellt Apples Model Intermediate Language (MIL)-Programme zur Laufzeit in Objective-C, die Konvolutionen (für lineare Schichten), Matrixmultiplikationen (für Attention), Softmax und elementweise Operationen definieren.
In-Memory-Kompilierung — Verwendet _ANEInMemoryModelDescriptor, um MIL-Text + Gewichts-Blobs direkt in ANE-Programme zu kompilieren, ohne .mlmodelc-Dateien auf die Festplatte schreiben zu müssen.
IOSurface I/O — Übergibt Eingabe-/Ausgabe-Tensoren über IOSurface Shared Memory im fp16-Format [1, channels, 1, spatial] — dem Format, das die ANE-Hardware nativ erwartet.
Gewichtseinbettung — Gewichte werden als BLOBFILE-Konstanten in ANE-Programme eingebettet und bei Gewichtsaktualisierungen für jeden Batch neu kompiliert.
Gradientenfluss — Forward-"Taps" legen Zwischenaktivierungen frei, die für Backward-Pässe benötigt werden. Backward-Kernel berechnen Eingabegradienten (dx) auf der ANE; Gewichtsgradienten (dW) werden auf der CPU über cblas_sgemm berechnet.

🧱 Architektur: 6 ANE-Kernel pro Trainingsschritt

Kernel	Funktion
`kFwdAttn`	RMSNorm + QKV-Projektion + SDPA + Ausgabe-Projektion
`kFwdFFN`	RMSNorm + SwiGLU FFN (W1, W3, SiLU, W2)
`kFFNBwd`	FFN-Backward (W2ᵀ + SiLU_bwd + W1ᵀ + W3ᵀ)
`kSdpaBwd1`	Woᵀ + SDPA-Backward Teil 1 (dV, probs, dp)
`kSdpaBwd2`	SDPA-Backward Teil 2 (softmax grad, dQ, dK)
`kQKVb`	QKV-Backward (Wqᵀ + Wkᵀ + Wvᵀ → dx)

CPU behandelt: RMSNorm-Backward, Residualverbindungen, Verlustberechnung, dW-Gradientenakkumulation (cblas_sgemm) und Adam-Optimierer-Updates.

📊 Benchmark-Ergebnisse (M4 Mac, einzelne Transformer-Schicht, dim=768, seq=512)

Optimierung	ms/Schritt	ANE-Auslastung
Baseline (vDSP Transpose)	33.5	3.1%
Channel-First-Layout	20.3	5.2%
vDSP vektorisierte RMSNorm	14.2	7.4%
GCD asynchrone cblas-Überlappung	11.4	9.2%
ANE RMSNorm-Fusion	11.4	9.2%
Wo^T-Fusion (7→6 Kernel)	11.4	9.2%
Verzögerte cblas-Wartezeit	9.3	11.2%

Bestes Ergebnis: 9,3 ms/Schritt, mit einer nachhaltigen Leistung von 1,78 TFLOPS (11,2 % der 15,8 TFLOPS Spitzenleistung der M4 ANE).

🔑 Wichtige Optimierungen

Channel-First CPU-Layout — Entspricht dem nativen [1,C,1,S] IOSurface-Format der ANE und eliminiert jeglichen Transpose-Overhead.
vDSP vektorisierte RMSNorm — 10× schneller als die naive Implementierung (6,7 ms → 0,7 ms).
GCD asynchrone cblas-Überlappung — Gewicht-Gradienten sgemm-Operationen laufen parallel zu ANE-Auswertungen auf einer seriellen Dispatch-Queue.
Verzögerte cblas-Wartezeit — Das Warten auf dW sgemms wird in den Forward-Pass des nächsten Schritts verschoben, um eine maximale Rechenüberlappung zu erzielen.
ANE RMSNorm-Fusion — RMSNorm wird direkt in die Forward-ANE-Kernel als MIL-Operationen integriert.
Forward-Taps — Q, K, V, Attention-Scores und Hidden States werden über Konkatenationsausgaben freigelegt, wodurch eine Neuberechnung auf der CPU während der Backward-Pässe vermieden wird.
exec()-Neustart — Umgeht das Limit von ~119 ANE-Kompilierungen pro Prozess durch Checkpoint/Resume.

📁 Repository-Struktur

├── api_exploration.m       # Erste Experimente zur Entdeckung von ANE-APIs
├── inmem_basic.m           # Proof-of-Concept für In-Memory MIL-Kompilierung
├── inmem_bench.m           # Benchmarks für ANE-Dispatch-Latenz
├── inmem_peak.m            # Messung der Spitzen-TFLOPS (2048×2048 Matmul)
├── sram_bench.m            # Untersuchung der ANE SRAM-Bandbreite
├── sram_probe.m            # Untersuchung von SRAM-Größe/-Layout
└── training/
    ├── ane_runtime.h       # Wrapper für private ANE-APIs (Kompilieren, Auswerten, IOSurface)
    ├── ane_mil_gen.h       # Hilfsfunktionen zur Generierung von MIL-Programmen
    ├── model.h             # Modellgewichtsinitialisierung und Blob-Builder
    ├── forward.h           # MIL-Generatoren für den Forward-Pass
    ├── backward.h          # MIL-Generatoren für den Backward-Pass
    ├── train.m             # Minimale Trainingsschleife (frühes Prototyp)
    ├── tiny_train.m        # Training eines 2-Schichten-Tiny-Modells
    ├── train_large.m       # Haupt: Training einer einzelnen Schicht dim=768 (optimiert)
    ├── test_*.m            # Unit-Tests für einzelne Kernel
    └── Makefile

🛠️ Bauen und Ausführen

Voraussetzungen: macOS 15+ auf Apple Silicon (getestet auf M4).

# Bauen
xcrun clang -O2 -framework Foundation -framework IOSurface \
  -framework CoreML -framework Accelerate -ldl -lobjc \
  -o train_large training/train_large.m

# Ausführen
./train_large

Keine externen Abhängigkeiten — verwendet nur System-Frameworks plus private ANE-APIs, die zur Laufzeit über objc_msgSend aufgelöst werden.

⚠️ Bekannte Einschränkungen

Einschränkung	Detail
Geringe ANE-Auslastung	~11,2 % der Spitzenleistung; viele elementweise Operationen fallen immer noch auf die CPU zurück
~119 Kompilierungslimit	ANE-Compiler leckt Ressourcen; umgangen durch `exec()`-Neustart
Einzelne Transformer-Schicht	Multi-Layer-Pipeline-Scheduling noch nicht implementiert
Nur synthetische Daten	Unterstützung für echte tokenisierte Daten ist in Arbeit
Kausale Maskierung von SDPA	ANE ignoriert `attn_mask` in SDPA; Umgehung durch manuelle Zerlegung
Private APIs	Verwendet undokumentierte APIs, die bei jedem macOS-Update brechen können

📚 Verwandte Forschungsartikel

🏷️ Tech Stack

Sprache: Objective-C, C
Plattform: macOS 15+ / Apple Silicon (M1/M2/M4)
Frameworks: Foundation, IOSurface, CoreML, Accelerate
Private APIs: _ANEClient, _ANECompiler, _ANEInMemoryModelDescriptor
Format: MIL (Model Intermediate Language), fp16

⚖️ Rechtlicher Haftungsausschluss

Dieses Projekt verwendet Apples private, undokumentierte APIs. Diese APIs bieten keine Stabilitätsgarantie und können bei jedem macOS-Update brechen. Das Projekt ist unabhängige Forschung unter Fair Use und Interoperabilitätsbestimmungen (Sega v. Accolade, 1992; DMCA §1201(f)). Es werden keine proprietären Apple-Codes oder Binärdateien enthalten. Nicht mit Apple Inc. verbunden oder von Apple Inc. unterstützt.