Entraîner des réseaux de neurones directement sur l'Apple Neural Engine sans CoreML ni GPU

Chaque processeur Apple Silicon contient un bloc dédié aux opérations de réseaux de neurones — l'Apple Neural Engine (ANE). Le fabricant revendique des dizaines de téraflops de performance, mais les outils officiels limitent son utilisation à l'inférence uniquement via CoreML. L'entraînement des modèles est proposé soit sur CPU, soit sur la puce graphique.

Un développeur sous le pseudonyme maderix a décidé de vérifier si le matériel ANE est physiquement incapable de rétropropagation, ou si le problème est uniquement dû aux limitations logicielles d'Apple. En un week-end, il a dissecté les frameworks privés de macOS et a réussi à faire exécuter par l'ANE non seulement la passe avant mais aussi la passe arrière pour les transformers.

Le projet a recueilli plus de 7 milliers d'étoiles sur GitHub. Décortiquons comment ce hack fonctionne, quelles astuces ont dû être appliquées, et à quoi ressemblent les performances réelles.

Pourquoi explorer les frameworks fermés

La pile officielle comme CoreML donne l'impression que l'ANE est une boîte noire fermée. Vous lui donnez un modèle prêt à l'emploi, il renvoie le résultat. Si vous devez entraîner même un petit modèle directement sur le client, vous deviez vous tourner vers Metal ou le framework MLX, en chargeant le GPU.

L'auteur du dépôt ANE a démontré que la puce est parfaitement capable d'exécuter des graphes de calcul arbitraires. Pour ce faire, il a rétro-ingénieré les bibliothèques privées _ANEClient et _ANECompiler, ainsi que le langage interne de description de modèle MIL (Model Intermediate Language).

Le résultat s'est avéré intéressant : un entraînement complet de transformers sans une seule ligne de CoreML ni de Metal. Le texte du modèle est assemblé directement en RAM, compilé à la volée et envoyé au processeur neuronal.

Limitations et réalité brutale

Avant de vous précipiter pour réécrire vos scripts d'entraînement sur Mac, il vaut la peine de regarder les chiffres honnêtes. L'auteur lui-même met en garde dès le départ : c'est une expérience académique, pas une bibliothèque prête pour la production.

Atteindre 100% d'utilisation des ressources ANE n'a pas encore été possible. L'utilisation réelle du chip est d'environ 5 à 9% du pic. Les limitations logicielles et matérielles se manifestent :

• Certaines opérations mathématiques ne sont pas supportées par la puce sous la forme requise et retombent sur le CPU.
• La passe arrière pour les poids (dW) doit encore être calculée par le processeur.
• Le compilateur ANE présente des fuites mémoire, donc après environ une hundred itérations, des contournements sont nécessaires.

Néanmoins, même dans ce mode, le projet offre une vitesse décente sur des architectures basiques.

Fonctionnement du pipeline d'entraînement

L'architecture du projet repose sur la distribution des tâches entre l'ANE et le CPU. L'accélérateur de réseau neuronal gère les multiplications matricielles les plus lourdes, tandis que le processeur s'occupe de la logique environnante et de l'accumulation des gradients.

La passe avant et le calcul des gradients d'entrée (dx) sont entièrement exécutés sur l'ANE. Les gradients de poids (dW) sont calculés par le CPU via les bibliothèques optimisées Accelerate et cblas_sgemm. L'optimiseur Adam et la couche RMSNorm s'exécutent également sur le processeur.

Pour éviter de recompiler le graphe du modèle à chaque étape lorsque les poids changent, l'auteur a appliqué une astuce : les poids et les activations sont regroupés dans un tenseur unique à travers les dimensions spatiales, et à l'intérieur du noyau MIL, ils sont simplement séparés.

La mémoire IOSurface est utilisée pour l'échange de données entre le CPU et l'ANE. Cela permet de transférer les tenseurs sans copie inutile entre les espaces d'adressage. Les données sont conditionnées au format spécifique ANE [1, C, 1, S], où les canaux arrivent en premier. Cette approche a éliminé la surcharge des transposes matricielles.

Problèmes non évidents et contournements

De nombreux pièges du matériel Apple ont émergé lors de la rétro-ingénierie.

Premièrement, l'opération SDPA (Scaled Dot-Product Attention) dans l'ANE ignore le masque causal attn_mask au niveau matériel. Le mécanisme d'attention a dû être divisé en trois étapes : multiplication de Q et K sur l'ANE, masquage avec softmax sur le CPU, et multiplication finale par V sur l'ANE.

Deuxièmement, le compilateur intégré _ANECompiler contient une fuite mémoire. Après environ 119 compilations, le processus plante en raison de l'épuisement des ressources. L'auteur a résolu le problème de manière radicale : lorsque le compteur approche la limite, le programme sauvegarde un point de contrôle et effectue exec() — se redémarre lui-même avec préservation de l'état.

Troisièmement, les calculs FP16 lors de la passe arrière conduisent rapidement à un underflow, provoquant la transformation des gradients en zéros. Le problème a été résolu en mettant à l'échelle la loss avec le coefficient 256 * NLAYERS.

Performances sur M4

Sur la puce Apple M4, les résultats se sont avérés assez illustratifs. Les tests ont été effectués sur deux architectures :

Pour le modèle Stories110M avec 109 millions de paramètres (12 couches, Multi-Head Attention classique), le temps pour une étape d'entraînement était de 91 millisecondes.

Le plus grand Qwen3-0.6B avec 596 millions de paramètres et Grouped-Query Attention traite une étape en 412 millisecondes.

L'auteur a également testé la quantification INT8 W8A8. L'utilisation de poids et d'activations de 8 bits réduit la charge sur la mémoire SRAM L2 de la puce et augmente le débit de 18,6 TOPS à 35,1 TOPS sur M4 — une accélération de près de 1,88x par rapport au FP16.

Comment exécuter le projet

Le projet ne nécessite pas de dépendances externes comme PyTorch ou Conda. Tout ce dont vous avez besoin est macOS 15 sur une machine avec Apple Silicon et le compilateur Clang. Les API privées sont chargées à l'exécution via objc_msgSend.

Pour construire le pipeline dynamique, naviguez simplement vers le dossier du projet et exécutez la commande make :

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

Si vous souhaitez tester la quantification INT8 ou mesurer les TOPS pic de votre puce, il existe des benchmarks séparés dans le dossier racine du dépôt.

Réflexions finales

Le projet maderix est un exemple de recherche de qualité « sous le capot » du matériel Apple. Il démontre que les limitations de l'ANE résident uniquement dans le plan logiciel et la nature fermée de l'écosystème.

Utiliser le dépôt pour entraîner des grands modèles de langage en production est actuellement inutile — le GPU et le framework MLX existent pour cela. Mais si vous étudiez le fonctionnement des accélérateurs neuronaux, si vous écrivez vos propres compilateurs pour l'Edge AI, ou si vous souhaitez comprendre comment travailler directement avec les API privées macOS depuis C et Objective-C, ce code sera une excellente ressource d'apprentissage.