Neurale netwerken direct trainen op de Apple Neural Engine zonder CoreML en GPU

Iedere Apple Silicon-processor bevat een speciaal blok voor neurale netwerkoperaties — de Apple Neural Engine (ANE). De fabrikant claimt tientallen teraflops aan prestaties, maar officiële tools beperken het gebruik tot inferentie via CoreML. Het trainen van modellen wordt aangeboden via de CPU of de grafische chip.

Een ontwikkelaar met de nickname maderix besloot te testen of ANE-hardware fysiek niet in staat is tot backpropagation, of dat het probleem uitsluitend ligt aan Apples softwarebeperkingen. In een weekend ontleedde hij macOS-privéframeworks en kreeg ANE zover om niet alleen forward maar ook backward pass voor transformers uit te voeren.

Het project verzamelde meer dan 7 duizend sterren op GitHub. Laten we uiteenzetten hoe deze hack werkt, welke trucs er moesten worden toegepast en hoe de werkelijke prestaties eruitzien.

Waarom graven in gesloten frameworks

De officiële stack zoals CoreML creëert de indruk dat ANE een gesloten black box is. Je geeft het een kant-en-klaar model, het geeft het resultaat terug. Als je zelfs een klein model direct op de client wilt trainen, moest je toevlucht nemen tot Metal of het MLX-framework, waarbij de GPU wordt belast.

De auteur van de ANE-repository demonstreerde dat de chip perfect in staat is om willekeurige computationele grafieken uit te voeren. Om dit te doen, reverse-engineerde hij de privé-bibliotheken _ANEClient en _ANECompiler, evenals de interne modelbeschrijvingstaal MIL (Model Intermediate Language).

Het resultaat was interessant: volwaardig trainen van transformers zonder een enkele regel CoreML of Metal. De modeltekst wordt direct in RAM samengesteld, on-the-fly gecompileerd en naar de neurale processor gestuurd.

Beperkingen en harde realiteit

Voordat je je trainingsscripts op Mac gaat herschrijven, is het de moeite waard om naar eerlijke cijfers te kijken. De auteur waarschuwt zelf vooraf: dit is een academisch experiment, geen productieklare bibliotheek.

Het bereiken van 100% benutting van ANE-bronnen is tot nu toe niet mogelijk geweest. Realistische chipbenutting ligt rond de 5–9% van de piek. Software- en hardwarebeperkingen maken zich kenbaar:

• Sommige wiskundige bewerkingen worden niet ondersteund door de chip in de vereiste vorm en vallen terug naar CPU.
• De backward pass voor gewichten (dW) moet nog steeds door de processor worden berekend.
• ANE-compiler heeft geheugenlekken, dus na ongeveer honderd iteraties zijn workarounds nodig.

Desondanks levert het project in deze modus een degelijke snelheid op basisarchitecturen.

Hoe de trainingspijplijn werkt

De projectarchitectuur vertrouwt op het verdelen van taken tussen ANE en CPU. De neurale netwerkaccelerator verwerkt de zwaarste matrixvermenigvuldigingen, terwijl de processor zorgt voor de omliggende logica en gradiëntaccumulatie.

De forward pass en berekening van input-gradiënten (dx) worden volledig uitgevoerd op ANE. Gewichtsgradiënten (dW) worden berekend door CPU via geoptimaliseerde Accelerate- en cblas_sgemm-bibliotheken. De Adam-optimizer en RMSNorm-laag draaien ook op de processor.

Om herberekening van de modelgrafiek bij elke stap te voorkomen wanneer gewichten veranderen, paste de auteur een truc toe: gewichten en activaties worden verpakt in één tensor via ruimtelijke dimensies, en binnen de MIL-kernel worden ze eenvoudig weer uit elkaar gehaald.

IOSurface-geheugen wordt gebruikt voor gegevensuitwisseling tussen CPU en ANE. Dit maakt het mogelijk om tensors over te dragen zonder onnodig kopiëren tussen adresruimten. Gegevens worden verpakt in het ANE-specifieke formaat [1, C, 1, S], waarbij kanalen eerst komen. Deze aanpak elimineerde de overhead van transponeren van matrices.

Niet voor de hand liggende problemen en workarounds

Tijdens het reverse-engineeren kwamen veel Apple-hardwarevalkuilen aan het licht.

Ten eerste negeert de SDPA (Scaled Dot-Product Attention)-operatie in ANE de causale maskering attn_mask op hardwareniveau. Het attention-mechanisme moest worden opgesplitst in drie fasen: vermenigvuldigen van Q en K op ANE, maskeren met softmax op CPU, en definitieve vermenigvuldiging met V terug op ANE.

Ten tweede bevat de ingebouwde compiler _ANECompiler een geheugenlek. Na ongeveer 119 compilaties crasht het proces door uitputting van bronnen. De auteur loste het probleem radicaal op: wanneer de teller de limiet nadert, slaat het programma een checkpoint op en doet exec() — start zichzelf opnieuw op met behoud van status.

Ten derde leiden FP16-berekeningen tijdens backward pass snel tot underflow, waardoor gradiënten nullen worden. Het probleem werd opgelost door het verlies te schalen met coëfficiënt 256 * NLAYERS.

Prestaties op M4

Op de Apple M4-chip bleken de resultaten vrij illustratief. Testen werd uitgevoerd op twee architecturen:

Voor het Stories110M-model met 109 miljoen parameters (12 lagen, klassieke Multi-Head Attention) was de tijd voor één trainingsstap 91 milliseconden.

Het grotere Qwen3-0.6B met 596 miljoen parameters en Grouped-Query Attention verwerkt één stap in 412 milliseconden.

De auteur testte ook INT8 W8A8-kwantisatie. Het gebruik van 8-bit gewichten en activaties vermindert de belasting op het L2 SRAM-geheugen van de chip en verhoogt de doorvoer van 18,6 TOPS naar 35,1 TOPS op M4 — een versnelling van bijna 1,88x ten opzichte van FP16.

Hoe het project uit te voeren

Het project vereist geen externe afhankelijkheden zoals PyTorch of Conda. Alles wat je nodig hebt is de nieuwste macOS 15 op een machine met Apple Silicon en de Clang-compiler. Privé-API's worden aangetrokken tijdens runtime via objc_msgSend.

Om de dynamische pijplijn te bouwen, navigeer je eenvoudig naar de projectmap en voer je het make-commando uit:

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

Als je INT8-kwantisatie wilt uitproberen of de piek-TOPS van je chip wilt benchmarken, zijn er aparte benchmarks in de hoofdmap van de repository.

Afsluitende gedachten

Het maderix-project is een voorbeeld van kwaliteitsonderzoek "onder de motorkap" van Apple-hardware. Het demonstreert dat ANE-beperkingen uitsluitend in het softwarevlak liggen en de gesloten aard van het ecosysteem.

De repository gebruiken voor het trainen van grote taalmodellen in productie is momenteel zinloos — GPU en het MLX-framework bestaan daarvoor. Maar als je bestudeert hoe neurale accelerators werken, je eigen compilers voor Edge AI schrijft, of wilt begrijpen hoe je rechtstreeks met macOS-privé-API's kunt werken vanuit C en Objective-C, zal deze code een uitstekende leermiddel zijn.