NeuronFabric: A Software Reference Architecture for On-Chip Transformer Training with Local Adam

Riepilogo Tecnico: NeuronFabric

Problematica

Gli attuali acceleratori neurali pubblicamente documentati generalmente separano la fase di calcolo dall'aggiornamento dello stato dell'ottimizzatore. I chip dedicati solo all'inferenza (es. Groq LPU, Apple ANE, IBM NorthPole) mancano del calcolo dei gradienti on-chip, mentre gli acceleratori per l'addestramento (es. Cerebras, Tenstorrent) tipicamente delegano l'aggiornamento dei pesi a una CPU host o a una GPU esterna. Questa separazione architettonica crea colli di bottiglia significativi:

• Traffico dei Pesi: L'inferenza richiede il caricamento di tutti i pesi dalla memoria esterna (HBM) per ogni passaggio in avanti (forward pass).
• Latenza di Addestramento: Il chip che calcola i gradienti non li applica; gli aggiornamenti sono orchestrati separatamente, richiedendo spesso il movimento di dati fuori dal chip.
• Limitazioni dell'Apprendimento Continuo: I chip di inferenza non possono aggiornare i pesi da nuovi dati senza ricaricare l'intero modello da un cluster esterno, un processo che richiede minuti o ore.

Il documento postula che un chip al silicio dove ogni neurone detiene i propri pesi ed esegue localmente l'ottimizzatore Adam potrebbe minimizzare il traffico off-chip e abilitare un vero addestramento on-chip. Tuttavia, prima di impegnarsi nel silicio, la validità matematica e la correttezza numerica di tale architettura "Local Adam" devono essere stabilite.

Metodologia

Gli autori hanno sviluppato NeuronFabric, un'architettura di riferimento software completa implementata in C# senza framework di machine learning esterni. Il sistema è progettato affinché ogni astrazione software mappi direttamente un blocco hardware (es. NeuronCore a BRAM, Forward() a catene DSP48).

Architettura Core:

• Local Adam: Ogni NeuronCore memorizza il proprio vettore di pesi $w$ e bias $b$. Durante il passaggio all'indietro (backward pass), calcola i gradienti e aggiorna i pesi in loco utilizzando l'algoritmo Adam. Non esiste un bus dei pesi tra i core.
• Configurazione del Modello: Un transformer autoregressivo da 334K parametri (Pre-LN) addestrato sul corpus di Shakespeare.
  • Dimensioni: Embedding $d=88$, 4 teste, Feedforward $f=264$, 4 layer.
  • Vocabolario: 256 (livello byte).
  • Tecniche: Weight tying (la proiezione di output riutilizza la matrice di embedding) per risparmiare parametri.

Schema di Precisione BF16W:
Per affrontare i severi vincoli di SRAM per l'addestramento su FPGA, gli autori hanno introdotto BF16W (Brain Floating Point 16 Weights).

• Memorizzazione: I pesi sono memorizzati come BF16 (2 byte).
• Calcolo: I pesi vengono convertiti in FP32 per il calcolo forward/backward e per gli aggiornamenti dei momenti Adam.
• Momenti: I primi ($m$) e secondi ($v$) momenti di Adam rimangono in FP32 (4 byte ciascuno) per garantire la stabilità.
• Aggiornamento: Dopo il passo Adam, il peso aggiornato in FP32 viene arrotondato nuovamente a BF16 e scritto nella memoria.

Vincolo del Budget del Vocabolario:
Il documento identifica un vincolo critico per i modelli a parametri ridotti: la tabella di embedding ($|V| \times d$) consuma una quota sproporzionata del budget dei parametri ("tassa del vocabolario"), lasciando capacità insufficiente per i layer transformer per apprendere pattern strutturali. Ciò è stato quantificato attraverso tre domini, mostrando che quando la capacità di ragionamento rimanente ($P_{reason}$) scende sotto i ~20K, la generazione coerente di testo fallisce.

Contributi Chiave

1. Implementazione Software di Riferimento: Un'implementazione completa in C# di un transformer senza framework, con backpropagation completa e aggiornamenti Local Adam, che esegue addestramento, valutazione e checkpointing in un unico binario.
2. Validazione BF16W: Dimostrazione che memorizzare i pesi in BF16 mantenendo i momenti in FP32 riduce l'uso di SRAM da 12 byte/parametro a 10 byte/parametro con una penalità di convergenza trascurabile (+0.020 di perdita di validazione su Shakespeare).
3. Quantificazione del Budget del Vocabolario: Evidenza empirica che mostra come, per un hardware a budget fisso, la tabella di embedding domini il modello, rendendo necessari o vocabolari piccoli o architetture multi-chip specializzate (es. chip di embedding separati).
4. Target di Fattibilità FPGA: Una configurazione concreta (modello 334K BF16W) che entra interamente nei 4.0 MB di BRAM on-chip di una FPGA Xilinx ZCU102, lasciando 660 KB di margine per i buffer di attivazione, senza richiedere l'accesso alla DDR durante l'addestramento.

Risultati

Il sistema è stato validato su una macchina Windows 11 (CPU AMD Ryzen 9 9900X, GPU NVIDIA RTX 4090) utilizzando il corpus di Shakespeare (80K campioni).

• Correttezza: I controlli numerici dei gradienti (differenze finite vs gradienti analitici) sono passati per tutti i componenti, confermando la correttezza matematica dell'implementazione Local Adam.
• Convergenza:
  • GPU FP32 (Oracle): Ha raggiunto una perdita di validazione di 1.5224.
  • CPU BF16W: Ha raggiunto una perdita di validazione di 1.5426.
  • Il divario di +0.020 è attribuito alla differenza di precisione tra la memorizzazione dei pesi in BF16 e il calcolo in FP32.
• Qualità Generativa: Il modello ha prodotto un testo coerente a livello di carattere con strutture di dialogo shakespeariane riconoscibili e pattern di personaggi nominati, confermando che il ciclo di addestramento converge verso una soluzione funzionale.
• Efficienza SRAM:
  • FP32 Adam (12 byte/param): 4.00 MB totali (riempie esattamente la BRAM della ZCU102, lasciando zero spazio per le attivazioni).
  • BF16W Adam (10 byte/param): 3.34 MB totali (lascia 660 KB liberi per attivazioni e stato di controllo).

Significato e Dichiarazioni

Il documento dichiara esplicitamente il suo ambito modesto: NeuronFabric è un prototipo di ricerca, non un acceleratore di produzione. Non contiene misurazioni FPGA o risultati di chiusura temporale (timing closure).

Il significato primario del documento risiede in:

1. Stabilire lo Stato dell'Arte: Serve come divulgazione architettonica pubblica e riferimento software per stabilire l'architettura "local Adam update" come stato dell'arte per future implementazioni FPGA/ASIC.
2. Prova di Concetto: Dimostra che la matematica funziona e che i numeri sono corretti per un transformer che si addestra interamente on-chip con lo stato dell'ottimizzatore locale, senza una CPU host nel loop di aggiornamento.
3. Definire la Strada da Seguire: Stabilisce il modello 334K BF16W di Shakespeare come un target falsificabile per la futura implementazione FPGA. Se la stessa curva di perdita può essere raggiunta sulla FPGA senza l'intervento dell'host, la fattibilità dell'addestramento transformer on-chip diventa un fatto misurabile piuttosto che una proiezione software.

Gli autori concludono che il prossimo passo immediato è implementare questo specifico ciclo di addestramento su una FPGA ZCU102 per verificare che i risultati software si traducano in hardware, validando così il paradigma di addestramento "senza CPU host".