Entraîner des Transformers sur le Neural Engine d'Apple - ANE GitHub

Révolutionnaire : Entraîner des Transformers directement sur le Neural Engine d'Apple

Le Neural Engine (ANE) d'Apple délivre 15,8 TFLOPS de puissance d'inférence dans les puces M4, mais l'entraînement ? Officiellement impossible. Jusqu'à présent.

La Percée : Entraînement Pur ANE

ANE Training est une implémentation from-scratch exécutant des boucles d'entraînement complètes de transformers - passes forward ET backward - directement sur le matériel ANE en utilisant des API privées reverse-engineered. Pas d'API d'entraînement CoreML. Pas de shaders Metal. Pas de fallback GPU. Calcul pur ANE.

Benchmarks Actuels (M4, dim=768, seq=512) : - 9,3 ms/étape - 11,2 % d'utilisation ANE (1,78 TFLOPS soutenus) - 6 dispatchs de kernels ANE par étape d'entraînement

Plongée Profonde dans l'Architecture

Le système orchestre 6 kernels ANE spécialisés par étape d'entraînement :

Kernel Fonction Innovation Clé
kFwdAttn RMSNorm + QKV + SDPA Les taps forward évitent le recalcul CPU
kFwdFFN SwiGLU FFN Fusion ANE RMSNorm
kFFNBwd Backward FFN Layout channel-first
kSdpaBwd1/2 Backward SDPA Fusion Wo^T réduit le nombre de kernels
kQKVb Backward QKV Chevauchement GCD async cblas

Le CPU gère seulement : Backward RMSNorm, résidus, perte, accumulation dW (Accelerate cblas), mises à jour Adam.

Optimisations Clés Qui Comptent

  1. Layout channel-first - Correspond à ANE IOSurface [1,C,1,S] éliminant la surcharge de transpose
  2. vDSP RMSNorm - Accélération x10 (6,7 ms → 0,7 ms)
  3. Fusion ANE RMSNorm - Intégrée dans les kernels forward
  4. cblas différé - Chevauchement maximum ANE/CPU
  5. Redémarrage exec() - Contourne la limite de 119 compilations

Évolution des performances : 33,5 ms → 9,3 ms grâce à une optimisation systématique.

Commencez en Quelques Minutes

# macOS 15+ Apple Silicon requis
xcrun clang -O2 -framework Foundation -framework IOSurface \
-framework CoreML -framework Accelerate -ldl -lobjc \
-o train_large training/train_large.m

./train_large

Zéro dépendances au-delà des frameworks système.

Points Saillants de la Structure des Fichiers

  • api_exploration.m - Découverte d'API privées
  • inmem_bench.m - Latence de dispatch ANE
  • sram_probe.m - Exploration de bande passante SRAM
  • training/train_large.m - Entraîneur de production mono-couche

Limitations & Roadmap

✅ Masquage causal via décomposition ✅ Accumulation de gradients/checkpointing ✅ Optimiseur Adam

🔄 Pipeline multi-couches 🔄 Datasets tokenisés réels 🔄 Entraînement de modèle complet

Note Légale

Utilise l'introspection runtime d'API non documentées à des fins de recherche/éducatives (DMCA §1201(f)). Aucun code propriétaire Apple inclus.

2,1k étoiles, 362 forks - Rejoignez la révolution ML Apple Silicon : https://github.com/maderix/ANE

Original Article: Voir l’original

Partager cet article