Addestrare Reti Neurali Direttamente su Apple Neural Engine Senza CoreML e GPU

Ogni processore Apple Silicon contiene un blocco dedicato per le operazioni di reti neurali — l'Apple Neural Engine (ANE). Il produttore dichiara decine di teraflop di prestazioni, ma gli strumenti ufficiali ne limitano l'uso alla sola inferenza tramite CoreML. L'addestramento dei modelli è offerto o sulla CPU o sulla scheda grafica.

Un developer con il nickname maderix ha deciso di verificare se l'hardware ANE fosse fisicamente incapace di eseguire la backpropagation, o se il problema fosse dovuto esclusivamente alle limitazioni software di Apple. In un fine settimana, ha analizzato i framework privati di macOS e ha fatto eseguire ad ANE non solo il forward pass ma anche il backward pass per i transformer.

Il progetto ha raccolto oltre 7 mila stelle su GitHub. Analizziamo come funziona questo hack, quali trucchi sono stati applicati e quali sono le prestazioni reali.

Perché Esplorare i Framework Chiusi

Lo stack ufficiale come CoreML crea l'impressione che ANE sia una scatola nera chiusa. Gli si fornisce un modello già pronto, restituisce il risultato. Se hai bisogno di addestrare anche un piccolo modello direttamente sul client, devi rivolgerti a Metal o al framework MLX, caricando la GPU.

L'autore del repository ANE ha dimostrato che il chip è perfettamente capace di eseguire grafi computazionali arbitrari. Per farlo, ha fatto reverse engineering delle librerie private _ANEClient e _ANECompiler, oltre al linguaggio interno di descrizione dei modelli MIL (Model Intermediate Language).

Il risultato si è rivelato interessante: addestramento completo di transformer senza una singola riga di CoreML o Metal. Il testo del modello viene assemblato direttamente in RAM, compilato al volo e inviato al processore neurale.

Limitazioni e Realtà

Prima di correre a riscrivere i tuoi script di addestramento su Mac, vale la pena guardare i numeri onesti. L'autore stesso avverte fin dall'inizio: questo è un esperimento accademico, non una libreria pronta per la produzione.

Raggiungere il 100% di utilizzo delle risorse ANE non è stato ancora possibile. L'utilizzo reale del chip è circa il 5–9% del picco. Le limitazioni software e hardware si fanno sentire:

• Alcune operazioni matematiche non sono supportate dal chip nella forma richiesta e ricadono sulla CPU.
• Il backward pass per i pesi (dW) deve ancora essere calcolato dal processore.
• Il compilatore ANE ha memory leak, quindi dopo circa un centinaio di iterazioni servono workaround.

Nonostante ciò, anche in questa modalità il progetto offre velocità decenti su architetture basiche.

Come Funziona la Pipeline di Addestramento

L'architettura del progetto si basa sulla distribuzione dei compiti tra ANE e CPU. L'acceleratore di rete neurale gestisce le moltiplicazioni di matrici più pesanti, mentre il processore si occupa della logica circostante e dell'accumulo dei gradienti.

Il forward pass e il calcolo del gradiente di input (dx) vengono eseguiti completamente su ANE. I gradienti dei pesi (dW) vengono calcolati dalla CPU attraverso le librerie ottimizzate Accelerate e cblas_sgemm. Anche l'ottimizzatore Adam e il layer RMSNorm vengono eseguiti sul processore.

Per evitare di ricompilare il grafo del modello ad ogni step quando i pesi cambiano, l'autore ha applicato un trucco: pesi e attivazioni vengono impacchettati in un singolo tensore attraverso le dimensioni spaziali, e all'interno del kernel MIL vengono semplicemente separati di nuovo.

La memoria IOSurface viene usata per lo scambio di dati tra CPU e ANE. Questo rende possibile trasferire tensori senza copie inutili tra spazi di indirizzi. I dati vengono impacchettati nel formato specifico ANE [1, C, 1, S], dove i canali vengono prima. Questo approccio ha eliminato l'overhead della trasposizione delle matrici.

Problemi Non Ovvi e Workaround

Molte insidie dell'hardware Apple sono emerse durante il reverse engineering.

Prima di tutto, l'operazione SDPA (Scaled Dot-Product Attention) in ANE ignora la maschera causale attn_mask a livello hardware. Il meccanismo di attenzione ha dovuto essere suddiviso in tre stadi: moltiplicazione di Q e K su ANE, mascheratura con softmax su CPU, e moltiplicazione finale per V di nuovo su ANE.

In secondo luogo, il compilatore integrato _ANECompiler contiene un memory leak. Dopo circa 119 compilazioni, il processo va in crash per esaurimento delle risorse. L'autore ha risolto il problema in modo radicale: quando il contatore si avvicina al limite, il programma salva un checkpoint e fa exec() — si riavvia mantenendo lo stato.

In terzo luogo, i calcoli in FP16 durante il backward pass portano rapidamente a underflow, causando la trasformazione dei gradienti in zeri. Il problema è stato risolto scalando la loss con coefficiente 256 * NLAYERS.

Prestazioni su M4

Sul chip Apple M4, i risultati si sono rivelati piuttosto illustrativi. I test sono stati condotti su due architetture:

Per il modello Stories110M con 109 milioni di parametri (12 layer, Multi-Head Attention classica), il tempo per uno step di addestramento è stato di 91 millisecondi.

Il più grande Qwen3-0.6B con 596 milioni di parametri e Grouped-Query Attention processa uno step in 412 millisecondi.

L'autore ha anche testato la quantizzazione INT8 W8A8. Usando pesi e attivazioni a 8 bit riduce il carico sulla SRAM L2 del chip e aumenta il throughput da 18.6 TOPS a 35.1 TOPS su M4 — uno speedup di quasi 1.88x rispetto a FP16.

Come Eseguire il Progetto

Il progetto non richiede dipendenze esterne come PyTorch o Conda. Tutto ciò che serve è l'ultimo macOS 15 su una macchina con Apple Silicon e il compilatore Clang. Le API private vengono caricate a runtime tramite objc_msgSend.

Per costruire la pipeline dinamica, basta navigare nella cartella del progetto ed eseguire il comando make:

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

Se vuoi provare la quantizzazione INT8 o benchmarkare i TOPS di picco del tuo chip, ci sono benchmark separati nella cartella root del repository.

Riflessioni Finali

Il progetto di maderix è un esempio di ricerca di qualità "sotto il cofano" dell'hardware Apple. Dimostra che le limitazioni di ANE risiedono esclusivamente nel piano software e nella natura chiusa dell'ecosistema.

Usare il repository per l'addestramento di modelli linguistici di grandi dimensioni in produzione è attualmente inutile — GPU e il framework MLX esistono per quello. Ma se stai studiando come funzionano gli acceleratori neurali, scrivendo compilatori personalizzati per Edge AI, o vuoi capire come lavorare con le API private di macOS direttamente da C e Objective-C, questo codice sarà un'eccellente risorsa di apprendimento.