Orion : Contourner CoreML pour exécuter et entraîner des LLM directement sur le Neural Engine d'Apple

Accès direct à l'ANE pour les charges de travail LLM

Orion fournit un système complet en Objective-C qui contourne entièrement CoreML pour exécuter et entraîner des LLM directement sur le moteur neuronal Apple (ANE). Cette approche donne aux développeurs un contrôle direct sur l'ANE, qui était auparavant traité comme un planificateur boîte noire par CoreML, leur retirant tout contrôle direct ou capacité d'entraînement.

Implémentation technique et contraintes

Le projet s'appuie sur des travaux de rétro-ingénierie qui ont cartographié les API privées ANEClient et ANECompiler. L'ANE présente ce que le développeur appelle un "désaccord d'impédance matérielle" avec 17 contraintes de programmation au total, dont 11 étaient complètement non documentées. Les contraintes clés incluent :

• L'opération concat provoque un échec immédiat et silencieux du compilateur
• Les poids BLOBFILE nécessitent un décalage de 64 octets depuis l'en-tête du bloc, sinon vous obtenez une corruption numérique silencieuse
• L'ANE maintient un état interne qui plafonne à environ 119 compilations par processus avant d'échouer silencieusement

Solutions aux défis d'entraînement

Les tentatives précédentes d'entraînement sur l'ANE rencontraient une divergence NaN après une seule étape. Orion résout ce problème en :

• Mettant en place un pipeline de compilation différée
• Implémentant un clampage strict des activations pour arrêter la cascade de débordement fp16 (clampage des activations entre -65504 et +65504)
• Utilisant une boucle de redémarrage de processus exec() après chaque étape d'entraînement pour contourner la limite de 119 compilations

Résultats de performance

Le compilateur réduit un graphe IR de 27 opérations à travers cinq passes d'optimisation jusqu'au MIL natif de l'ANE. Les performances actuelles incluent :

• Plus de 170 tokens/s pour le décodage GPT-2 124M
• Un entraînement multi-étapes mécaniquement stable sur un transformateur de 110 millions de paramètres (le "plafond de cohérence" du matériel)
• Sur plus de 1 000 étapes, la perte est passée de 12,3 à 6,2 sans aucun NaN

Limitations actuelles

L'ANE incorpore les poids au moment de la compilation, ce qui signifie que chaque mise à jour d'entraînement nécessite une pénalité de recompilation d'environ 4,2 s. L'ANE tire environ 19 TFLOPS en fp16, mais la contrainte fondamentale pour son utilisation n'a pas été la puissance de calcul—c'a été l'absence totale d'une couche d'orchestration native.