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

Le papier présente Orion, le premier système open source permettant l'entraînement et l'inférence de grands modèles de langage directement sur le Neural Engine d'Apple en contournant CoreML, en caractérisant ses contraintes matérielles et en optimisant le temps de compilation pour accélérer l'entraînement de 3,8 fois.

L'essentiel

🍎 Orion : Réveiller le "Super-Cerveau" caché dans votre iPhone

Imaginez que vous possédez une voiture de sport de dernière génération (un iPhone ou un Mac récent). Sous le capot, il y a un moteur électrique ultra-puissant et très économe en énergie, capable de faire des miracles. Mais jusqu'à présent, les conducteurs (les développeurs) ne savaient pas comment l'allumer. Ils utilisaient toujours le vieux moteur à essence (le processeur principal) pour tout faire, même les tâches simples, gaspillant ainsi le potentiel de la voiture.

Ce moteur caché s'appelle le Neural Engine (ANE) d'Apple. Il est présent dans plus de 2 milliards d'appareils, mais personne ne l'utilisait vraiment pour entraîner les intelligences artificielles modernes (les "LLM" comme ChatGPT).

Orion est le projet qui a réussi à trouver la clé pour démarrer ce moteur et l'utiliser à plein régime.

🔑 Le problème : La boîte noire

Apple a construit une barrière autour de ce moteur. L'outil officiel qu'ils donnent aux développeurs (CoreML) est comme une boîte noire.

• Vous mettez une tâche dedans.
• La boîte décide elle-même si elle l'exécute sur le moteur principal, la carte graphique ou le Neural Engine.
• Vous ne pouvez pas forcer l'utilisation du Neural Engine.
• Surtout, vous ne pouvez pas entraîner l'IA directement dessus (apprendre de nouvelles choses).

C'est comme si vous aviez un Ferrari, mais que le concessionnaire vous disait : "Vous pouvez conduire, mais vous ne pouvez jamais changer l'huile, et je décide seul de la vitesse."

🛠️ La solution : Orion, le mécanicien de génie

Les chercheurs ont créé Orion. C'est un système complet qui contourne la boîte noire d'Apple pour parler directement au moteur.

Voici comment ils ont fait, avec des analogies simples :

1. Le Dictionnaire Secret (La Caractérisation)

Avant de pouvoir conduire, il faut connaître les règles de la route. Le Neural Engine a des règles très strictes et secrètes (par exemple : "ne faites pas ceci", "les données doivent être rangées dans cet ordre").

• L'analogie : C'est comme si le moteur avait un manuel d'utilisation écrit dans une langue inconnue. Orion a traduit ce manuel. Les chercheurs ont découvert 20 règles secrètes (comme "ne mélangez pas trop de pièces d'un coup" ou "les données doivent être empilées d'une manière très précise"). Sans ces règles, le moteur s'arrête ou fait des erreurs silencieuses.

2. Le Traducteur Instantané (Le Compilateur)

L'IA parle un langage complexe. Orion agit comme un traducteur ultra-rapide qui transforme les instructions de l'IA en un langage que le moteur Apple comprend parfaitement.

• L'analogie : C'est un traducteur qui ne se contente pas de traduire mot à mot, mais qui réorganise toute la phrase pour qu'elle soit plus fluide et plus rapide pour le moteur.

3. La Magie du "Changement à la Volée" (Delta Compilation)

C'est la partie la plus brillante.

• Le problème : Pour apprendre, l'IA doit changer ses "poids" (ses connaissances) à chaque étape. Normalement, sur ce moteur, changer un poids oblige à tout reconstruire (recompiler) le programme. C'est comme si, pour changer une roue de votre voiture, vous deviez démonter tout le moteur, le reconstruire, et le remonter. Cela prendrait 4 secondes à chaque fois, ce qui rend l'apprentissage très lent.
• La solution d'Orion : Ils ont découvert une astuce. Au lieu de reconstruire tout le moteur, ils peuvent éteindre le programme, changer les fichiers de poids sur le disque dur, et le rallumer instantanément.
• L'analogie : Imaginez un chef cuisinier qui prépare un plat. Au lieu de réécrire tout le livre de recettes à chaque fois qu'il change une pincée de sel, il ouvre juste le livre, efface la ligne du sel, écrit la nouvelle quantité, et continue. Orion a réduit le temps de "changement" de 4 secondes à 0,5 seconde. C'est un gain de vitesse énorme !

4. Les Adaptateurs Chauds (LoRA)

Parfois, on veut adapter l'IA à une tâche spécifique (par exemple, faire de la poésie ou du code) sans tout réapprendre.

• L'analogie : Orion permet de brancher un "accessoire" (un adaptateur) sur le moteur sans le démonter. Vous pouvez changer cet accessoire en quelques secondes, comme changer de lunettes, sans avoir à redémarrer la voiture.

🏁 Les Résultats : La voiture de course enfin lancée

Grâce à Orion, les chercheurs ont pu :

1. Entraîner une IA directement sur un Mac (M4 Max) sans utiliser la carte graphique.
2. Le faire 22 minutes pour 1 000 étapes d'apprentissage (au lieu de 85 minutes avec l'ancienne méthode).
3. Obtenir des résultats stables, sans que l'IA ne devienne folle (pas d'erreurs mathématiques).
4. Faire tourner l'IA très vite pour répondre à des questions (170 mots par seconde).

💡 Pourquoi c'est important pour vous ?

Aujourd'hui, votre iPhone ou votre Mac a un super-moteur qui dort. Orion prouve qu'on peut le réveiller pour faire tourner des intelligences artificielles puissantes, directement sur votre appareil, sans avoir besoin d'envoyer vos données dans le cloud.

C'est comme si on avait trouvé le moyen de transformer votre voiture de ville en une voiture de course capable de faire des courses d'endurance, tout en consommant très peu d'énergie. Et le meilleur ? Orion est gratuit et ouvert à tout le monde pour que d'autres puissent construire dessus.

En résumé : Orion est la clé qui déverrouille le potentiel caché de vos appareils Apple pour l'intelligence artificielle, en apprenant à parler directement au langage secret du matériel.

Résumé technique

1. Problématique

Bien que plus de deux milliards d'appareils Apple soient équipés du Neural Engine (ANE), un processeur dédié aux tâches d'intelligence artificielle (NPU), cette ressource reste largement sous-exploitée pour les modèles de langage (LLM).

• Absence de support public : Le framework public de ML d'Apple, CoreML, agit comme une boîte noire. Il ne permet pas aux développeurs de forcer l'exécution sur l'ANE, d'inspecter les programmes générés ou de réaliser un apprentissage (training) directement sur l'appareil.
• Limitations de l'API privée : L'API native de l'ANE (via AppleNeuralEngine.framework) est privée, non documentée et repose sur un modèle de compilation où les poids sont "cuits" (baked) au moment de la compilation. Cela signifie que toute mise à jour des poids (nécessaire pour l'entraînement) exigeait auparavant une recompilation complète du programme, rendant l'entraînement itératif extrêmement lent et inefficace.
• État de l'art incomplet : Des travaux antérieurs (comme maderix et ANEgpt) ont démontré la faisabilité de l'accès direct à l'ANE, mais aucun système complet n'offrait à la fois un pipeline de compilation, une exécution stable sur plusieurs étapes d'entraînement avec reprise de checkpoint, et une gestion efficace des mises à jour de poids.

2. Méthodologie et Architecture

Orion est présenté comme le premier système open-source complet permettant d'exécuter directement l'ANE pour l'inférence et l'entraînement des LLM, en contournant totalement CoreML.

A. Caractérisation de l'ANE

Les auteurs ont étendu les travaux de maderix pour établir un catalogue de 20 contraintes de programmation de l'ANE. Ces contraintes couvrent :

• Restrictions MIL IR : Certaines opérations (comme concat) sont rejetées, tandis que d'autres (comme gelu) doivent être remplacées par des approximations (tanh).
• Contraintes mémoire et I/O : Les buffers de sortie doivent avoir des tailles uniformes, les entrées/sorties doivent être ordonnées alphabétiquement, et les surfaces IOSurface doivent respecter une taille minimale (~49 Ko).
• Limites de compilation : Un processus ne peut compiler qu'environ 119 programmes avant de planter silencieusement.

B. Conception du Système

Orion est structuré en cinq couches :

1. API Privées : Utilisation directe de ANEClient, ANECompiler et ANEModel via Objective-C.
2. Runtime : Gestion du cycle de vie des programmes, mise en cache, et I/O via IOSurface (transfert de données sans copie entre CPU et ANE).
3. Compilateur : Transforme un IR de graphe (27 opérations) en code MIL natif via 5 passes d'optimisation (élimination de code mort, fusion de cast, annotation SRAM, etc.).
4. Couche Modèle : Gestion des configurations, des tokenizers (BPE/SentencePiece) et du chargement des poids.
5. Applications CLI : Pour l'inférence, l'entraînement et les benchmarks.

C. Innovation Clé : La Compilation Delta (Delta Compilation)

C'est la contribution majeure pour l'entraînement.

• Problème initial : Comme l'ANE cuit les poids à la compilation, chaque étape de mise à jour des poids (gradient descent) nécessitait une recompilation complète (~4,2 secondes par étape), consommant 84 % du temps total.
• Solution Orion (v2.0) : Au lieu de recompiler, le système :
  1. Décharge le programme compilé de l'ANE.
  2. Met à jour les fichiers de poids sur le disque (BLOBFILE).
  3. Recharge le programme existant.
  • Puisque le texte MIL et les clés du dictionnaire de poids restent identiques, l'ANE conserve l'identité du programme et accepte les nouveaux poids sans passer par le compilateur (ANECCompile).
  • Cela réduit le temps de mise à jour de 4 200 ms à 494 ms (accélération de 8,5x).

D. Adaptateurs LoRA en tant qu'Entrées

Pour contourner la limitation des poids figés, Orion implémente LoRA (Low-Rank Adaptation) en passant les matrices d'adaptation (A et B) comme entrées IOSurface plutôt que comme poids cuits. Cela permet de changer d'adaptateur à chaud ("hot-swap") sans aucune recompilation.

3. Résultats Expérimentaux

Les tests ont été réalisés sur un Mac Studio M4 Max (16 cœurs ANE, 64 Go de mémoire unifiée).

A. Performance d'Inférence (GPT-2 124M)

• Débit : Orion atteint >170 tokens/seconde en mode décodage sur l'ANE.
• Précision : Accord à 100 % avec une référence CPU (fp32) sur les séquences générées.
• Latence : Bien que le CPU soit plus rapide pour le décodage token par token (283 tok/s) en raison de la surcharge des allers-retours IOSurface, l'ANE offre un avantage pour les longues séquences (prefill) et consomme zéro puissance au repos.

B. Performance d'Entraînement (Stories110M)

• Vitesse : L'entraînement de 1 000 étapes sur le jeu de données TinyStories prend 22,4 minutes avec la compilation delta (v2.0), contre ~85 minutes avec la recompilation complète (v1.0). C'est une accélération globale de 3,8x.
• Stabilité : Le système a réussi 1 000 étapes d'entraînement sans aucune occurrence de NaN (Not a Number), résolvant trois bugs critiques (désynchronisation des poids, débordement fp16, corruption des fichiers de poids) qui faisaient diverger les systèmes précédents.
• Efficacité : La compilation delta réduit la part du temps de recompilation dans une étape d'entraînement de 83,9 % à 36,8 %.

C. Accélérations Spécifiques

L'ANE montre des accélérations massives pour certaines opérations :

• Softmax (vocabulaire 32k) : 33,8x plus rapide que le CPU.
• Convolution 1x1 : 3x plus rapide que les multiplications matricielles équivalentes.

4. Contributions Clés

1. Catalogue de contraintes : Identification de 20 limitations techniques de l'ANE (14 nouvelles découvertes), couvrant le IR MIL, la disposition mémoire et le comportement numérique.
2. Compilateur et Runtime : Un compilateur de graphe IR vers MIL natif et un runtime en Objective-C gérant les I/O sans copie et la mise en cache.
3. Compilation Delta : Une technique ingénieuse contournant la limite de recompilation de l'ANE, rendant l'entraînement itératif viable.
4. LoRA Hot-Swap : Mécanisme permettant de changer d'adaptateurs LoRA dynamiquement sans recompilation.
5. Stabilité d'entraînement : Résolution des problèmes de divergence numérique (NaN) pour permettre un entraînement résilient avec reprise de checkpoint.
6. Open Source : Publication complète du code sous licence MIT, incluant le runtime, les convertisseurs de poids et les benchmarks.

5. Signification et Impact

• Démocratisation de l'IA sur appareil : Orion prouve que le Neural Engine d'Apple, présent sur des milliards d'appareils, peut être utilisé pour l'entraînement et l'inférence de LLM, et non seulement pour l'inférence de modèles pré-entraînés.
• Nouveau paradigme pour les NPU : La méthode de caractérisation (sondage systématique des API privées) et la technique de compilation delta peuvent inspirer le développement de systèmes similaires pour d'autres NPU (Qualcomm, Samsung, Google).
• Économie d'énergie : L'ANE étant "hard power-gated" (éteint physiquement) lorsqu'il n'est pas utilisé, son utilisation pour l'inférence "toujours active" est idéale pour les appareils mobiles, offrant une efficacité énergétique supérieure aux GPU/CPU.
• Futur de la recherche : En libérant Orion, les auteurs permettent à la communauté de construire sur cette base pour développer la prochaine génération de systèmes d'IA natifs pour les NPU, exploitant pleinement le matériel existant qui était auparavant inexploité.

En résumé, Orion transforme l'ANE d'un "accélérateur sombre" en une plateforme programmable et performante pour l'entraînement et l'inférence de modèles de langage, en surmontant les obstacles logiciels et matériels grâce à une ingénierie logicielle avancée et une compréhension profonde du matériel.