Orion: Characterizing and Programming Apple's Neural Engine for LLM Training and Inference

Orion is het eerste open end-to-end systeem dat de Apple Neural Engine direct aanstuurt via private API's om zowel stabiel on-device training als snelle inferentie van grote taalmodellen mogelijk te maken door CoreML te omzeilen en recompilatie-tijd tijdens training drastisch te reduceren.

De kern

Stel je voor dat je een enorme fabriek hebt met een superkrachtige, speciale robotarm die alleen maar ontworpen is om heel snel ingewikkelde puzzels op te lossen. Deze robotarm is de Neural Engine (ANE) van Apple. Hij zit in meer dan twee miljard iPhones, iPads en Macs.

Het probleem? Tot nu toe was deze robotarm als een gesloten kast. De fabrieksleutel (de software) die je normaal gebruikt, gaf je geen toegang tot de knoppen van deze robot. Je kon hem niet zelf aansturen, en je kon hem niet gebruiken om nieuwe dingen te leren (trainen). De fabrikant (Apple) zei: "Gebruik maar de algemene werkbank (de GPU of CPU), die is wel goed genoeg."

Orion is het project dat de sleutel heeft gevonden om die kast open te maken. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. De "Blindganger" aanpak

Voorheen was het alsof je een blindeman was die probeerde een auto te besturen. Je wist niet welke knop wat deed. Orion is de eerste die de handleiding heeft geschreven. De onderzoekers hebben de geheime taal van de robotarm vertaald en een lijst gemaakt van 20 regels waar je je aan moet houden.

• Voorbeeld: Ze ontdekten dat je bepaalde puzzelstukken (zoals het samenvoegen van data) niet in één keer mag doen, maar in aparte stapjes. Als je dat niet doet, stopt de robot.

2. De "Koffieautomaat" voor training (Delta Compilation)

Dit is misschien wel het coolste stukje.
Stel je voor dat je een koffieautomaat hebt die elke keer dat je een nieuwe smaak (gewicht) wilt toevoegen, de hele machine moet demonteren en opnieuw moet bouwen. Dat duurt 4 seconden per kopje. Als je 1000 kopjes wilt maken, ben je 4000 seconden kwijt aan bouwen, niet aan koffie drinken.

Orion heeft een trucje bedacht: Delta Compilation.
In plaats van de hele machine te slopen, doen ze het volgende:

1. Zet de machine uit (unloading).
2. Vervang alleen het koffiebonen-reservoir met de nieuwe smaak (patching the weights).
3. Zet de machine weer aan (reloading).

Dit duurt nu nog maar een fractie van de tijd. In plaats van 4 seconden per stap, duurt het nu 0,5 seconden. Hierdoor wordt het trainen van een slimme AI 3,8 keer sneller. Het is alsof je van een handmatige koffiezetapparaat overschakelt naar een automaat die alleen de bonen verwisselt.

3. De "Hot-Swap" Truc (LoRA)

Normaal gesproken moet je een robot volledig herbouwen als je hem wilt leren een nieuwe taal. Orion gebruikt een slimme truc: ze steken de nieuwe kennis in een losse zak (een adapter) die ze tijdens het werk kunnen verwisselen.

• Analogie: Stel je voor dat je een chef-kok hebt die altijd dezelfde basisrecepten maakt. Wil je nu Italiaans koken? Je hoeft de hele keuken niet te verbouwen. Je hangt gewoon een "Italiaans-receptenboek" (de adapter) aan de muur, en de chef gebruikt dat. Zodra je weer Frans wilt koken, hang je dat boekje er weer af. Geen herbouwen nodig!

4. Het resultaat: Een stabiele leerling

Voorheen probeerden anderen dit, maar de robotarm werd vaak gek (hij gaf "NaN" fouten, wat betekent dat de cijfers uit elkaar vielen en het resultaat onzin werd). Orion heeft drie specifieke "bugjes" in de software opgelost die ervoor zorgden dat de robotarm stabiel blijft, zelfs als hij urenlang blijft leren.

• Ze hebben bewezen dat je een model van 110 miljoen parameters (een vrij slimme AI) in 22 minuten kunt laten leren op een M4 Max computer, zonder dat het systeem crasht.

Waarom is dit belangrijk?

• Geen wachtlijst meer: Je hoeft niet meer te wachten tot je AI in de "cloud" (op een server ver weg) wordt verwerkt. Je kunt het direct op je eigen apparaat doen.
• Energiezuinig: De Neural Engine verbruikt bijna geen stroom als hij niet gebruikt wordt.
• Open Source: De onderzoekers hebben de sleutels (de code) gratis beschikbaar gesteld. Iedereen kan nu meedoen aan het bouwen van slimme apps die deze speciale robotarm gebruiken.

Kortom: Orion heeft de "verboden zone" van Apple's chip opengebroken, een slimme manier gevonden om de machine snel aan te passen zonder hem te herbouwen, en bewezen dat je er zelfs mee kunt leren. Het is alsof ze een nieuwe, snellere weg hebben gevonden door een berg die voorheen als onbegaanbaar werd beschouwd.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "Orion: Characterizing and Programming Apple's Neural Engine for LLM Training and Inference" in het Nederlands.

Probleemstelling

Hoewel Apple's Neural Engine (ANE) een dedicated Neural Processing Unit (NPU) is die in meer dan twee miljard Apple-apparaten zit (vanaf de A11 Bionic tot de M4-serie), blijft deze accelerator vrijwel ongebruikt voor Large Language Model (LLM) taken.

• Software-beperkingen: Het enige publieke framework, CoreML, fungeert als een "black box" die operaties dynamisch toewijst aan CPU, GPU of ANE. Ontwikkelaars kunnen de ANE niet direct programmeren, geen gradiënten berekenen of training uitvoeren.
• Ontbrekende ondersteuning: Er is geen openbaar framework dat LLM-training op de ANE ondersteunt. Bestaande frameworks zoals MLX en llama.cpp targeten de GPU of CPU, waardoor de enorme rekenkracht van de ANE (tot 38 TOPS INT8, gemeten ~19 TFLOPS FP16) voor LLM's ongebruikt blijft.
• Technische uitdagingen: De ANE vereist dat gewichten tijdens het compileren worden "ingebrand" (baked). Dit betekent dat elke update van gewichten tijdens training normaal gesproken een volledige hercompilatie vereist, wat extreem tijdrovend is. Bovendien is de instructieset (MIL) en de compilatie-API's niet gedocumenteerd.

Methodologie: Orion

Orion is het eerste open-source systeem dat directe ANE-uitvoering combineert met een compiler-pijplijn en stabiele multi-stap training in een native runtime, volledig omzeilend CoreML via private API's (ANEClient, ANECompiler).

1. Karakterisering van de ANE:
De auteurs hebben, voortbouwend op eerder werk van maderix, een catalogus van 20 beperkingen samengesteld. Hierbij zijn 14 nieuwe beperkingen ontdekt tijdens de ontwikkeling, waaronder:

• MIL IR-beperkingen: Operaties zoals concat worden geweigerd; gelu moet worden vervangen door een tanh-benadering.
• Geheugen- en I/O-beperkingen: Strikte vereisten voor tensor-layouts ([1, C, 1, S]), uniforme buffergroottes voor multi-input/output, en alfabetische sortering van variabelen.
• Compilatie-limieten: Een limiet van ~119 compilaties per proces en een prestatie-dip ("cliff") wanneer het SRAM-begroting (32 MB) wordt overschreden.

2. Compiler en Runtime:

• Compiler: Transformeert een graf-IR met 27 operaties via vijf optimalisatiepasses (o.a. DCE, cast fusion, SRAM-annotatie) naar ANE-natieve MIL-code.
• Runtime: Beheert IOSurface-gebaseerde zero-copy I/O tussen CPU en ANE, en implementeert program-caching.

3. Delta Compilation (Kerninnovatie):
Om het probleem van de "ingebrande gewichten" op te lossen, introduceert Orion Delta Compilation.

• In plaats van het hele programma opnieuw te compileren na een gewichtsupdate, wordt het bestaande programma van de ANE geladen (unload), de gewichtsbestanden op de schijf gepatcht, en het programma opnieuw geladen (reload).
• Omdat de MIL-code en de sleutels van het gewichtwoordenboek ongewijzigd blijven, behoudt de ANE dezelfde interne identiteit en wordt de zware compilatie (ANECCompile()) volledig omzeild.

4. LoRA Adapter-as-Input:
Om LoRA (Low-Rank Adaptation) te ondersteunen zonder hercompilatie, worden de adapter-matrices (A en B) niet als gewichten ingebakken, maar als IOSurface-inputs doorgegeven. Dit maakt het mogelijk om adapters "hot-swap" te wisselen tijdens inferentie.

5. Numerieke Stabiliteit:
De auteurs hebben drie kritieke bugs opgelost die eerder leiden tot NaN-divergentie bij training:

• Vertraagde compilatie: Programma's worden pas gecompileerd na het laden van checkpoints om gewichtsmismatches te voorkomen.
• FP16-overflow: Activeringen worden geklemd binnen het FP16-bereik [-65504, +65504] om overlopen te voorkomen.
• Gewichtsvalidatie: Een validatiepass controleert op corrupte gewichten in BLOB-bestanden.

Belangrijkste Resultaten

De prestaties zijn gemeten op een Mac Studio met M4 Max (16 ANE-kernen, 64 GB unified memory).

• Inferentie:

  • GPT-2 124M bereikt 170+ tokens/second in decode-modus op de ANE.
  • De CPU (via Metal/cblas) is sneller voor korte sequenties (283 tok/s) vanwege de overhead van IOSurface-communicatie per dispatch, maar de ANE is superieur voor specifieke operaties zoals softmax over grote vocabulaires (33,8x sneller dan CPU).
  • 100% overeenkomst in gegenereerde tokens vergeleken met een CPU-basislijn.

• Training (Stories110M op TinyStories):

  • Delta Compilation effect: De tijd voor hercompilatie/reload daalde van 4.200 ms (volledige recompilatie) naar 494 ms (delta reload), een snelheidswinst van 8,5x.
  • Totale trainingssnelheid: 1.000 trainingsstappen werden voltooid in 22 minuten (vs. ~85 minuten met volledige recompilatie), een 3,8x versnelling.
  • Stabiliteit: Training liep stabiel voor 1.000 stappen met 0 NaN-incidenten, zelfs bij het hervatten van checkpoints.

Significantie en Impact

• Eerste End-to-End Systeem: Orion is, voor zover bekend, het eerste open systeem dat directe ANE-programmering, een compiler-pijplijn en stabiele training met checkpoint-resume combineert.
• Ontgrendeling van Hardware: Het paper toont aan dat de ANE, ondanks zijn beperkingen, een krachtige accelerator is voor LLM's en dat de "compile-then-dispatch" architectuur niet noodzakelijk een onoverkomelijke barrière is voor training.
• Open Source Bijdrage: De code is beschikbaar onder de MIT-licentie en bevat een complete Objective-C runtime, compilers, en benchmarks. Dit stelt de onderzoeksgemeenschap in staat om verder te bouwen op deze karakterisering van de ANE.
• Toekomstperspectief: De methodologie van het reverse-engineeren van private API's en het documenteren van ongedocumenteerde beperkingen kan worden toegepast op andere NPU's (zoals Qualcomm Hexagon of Google TPU Edge) om hun potentieel voor on-device AI te maximaliseren.

Kortom, Orion bewijst dat het mogelijk is om Apple's Neural Engine effectief te gebruiken voor zowel inferentie als training van LLM's, door slimme software-architectuur (Delta Compilation) te combineren met diepgaande hardware-karakterisering.