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

Résumé Technique : NeuronFabric

Problématique

Les accélérateurs neuronaux actuellement documentés publiquement séparent généralement la phase de calcul de la mise à jour de l'état de l'optimiseur. Les puces dédiées uniquement à l'inférence (ex. : Groq LPU, Apple ANE, IBM NorthPole) manquent de calcul de gradient sur puce, tandis que les accélérateurs d'entraînement (ex. : Cerebras, Tenstorrent) délèguent typiquement les mises à jour de poids à un CPU hôte ou à un GPU externe. Cette séparation architecturale crée des goulots d'étranglement significatifs :

• Trafic de poids : L'inférence nécessite le chargement de tous les poids depuis une mémoire externe (HBM) pour chaque passe avant.
• Latence d'entraînement : La puce qui calcule les gradients ne les applique pas ; les mises à jour sont orchestrées séparément, nécessitant souvent des mouvements de données hors de la puce.
• Limitations de l'apprentissage continu : Les puces d'inférence ne peuvent pas mettre à jour les poids à partir de nouvelles données sans recharger l'intégralité du modèle depuis un cluster externe, un processus prenant des minutes ou des heures.

L'article postule qu'une puce de silicium où chaque neurone détient ses propres poids et exécute l'optimiseur Adam localement pourrait minimiser le trafic hors puce et permettre un véritable entraînement sur puce. Cependant, avant de s'engager dans le silicium, la validité mathématique et la justesse numérique d'une telle architecture « Adam locale » doivent être établies.

Méthodologie

Les auteurs ont développé NeuronFabric, une architecture de référence logicielle complète implémentée en C# sans frameworks d'apprentissage automatique externes. Le système est conçu de telle sorte que chaque abstraction logicielle correspond directement à un bloc matériel (ex. : NeuronCore vers BRAM, Forward() vers des chaînes DSP48).

Architecture Centrale :

• Adam Local : Chaque NeuronCore stocke son propre vecteur de poids $w$ et son biais $b$. Lors de la passe arrière, il calcule les gradients et met à jour les poids sur place en utilisant l'algorithme Adam. Il n'y a pas de bus de poids entre les cœurs.
• Configuration du Modèle : Un transformeur autoregressif de 334K paramètres (Pre-LN) entraîné sur le corpus de Shakespeare.
  • Dimensions : Embedding $d=88$, 4 têtes, Feedforward $f=264$, 4 couches.
  • Vocabulaire : 256 (niveau octet).
  • Techniques : Partage de poids (la projection de sortie réutilise la matrice d'embedding) pour économiser les paramètres.

Schéma de Précision BF16W :
Pour répondre aux contraintes strictes de la SRAM pour l'entraînement sur FPGA, les auteurs ont introduit le BF16W (Brain Floating Point 16 Weights).

• Stockage : Les poids sont stockés en BF16 (2 octets).
• Calcul : Les poids sont convertis en FP32 pour le calcul (forward/backward) et les mises à jour des moments Adam.
• Moments : Les premiers ($m$) et seconds ($v$) moments d'Adam restent en FP32 (4 octets chacun) pour garantir la stabilité.
• Mise à jour : Après l'étape Adam, le poids FP32 mis à jour est arrondi en BF16 et réécrit dans le stockage.

Contrainte de Budget de Vocabulaire :
L'article identifie une contrainte critique pour les modèles à faibles paramètres : la table d'embedding ($|V| \times d$) consomme une part disproportionnée du budget de paramètres (« taxe de vocabulaire »), laissant une capacité insuffisante pour que les couches du transformeur apprennent des structures de raisonnement. Cela a été quantifié à travers trois domaines, montrant que lorsque la capacité de raisonnement restante ($P_{reason}$) descend sous ~20K, la génération de texte cohérent échoue.

Contributions Clés

1. Implémentation de Référence Logicielle : Une implémentation C# complète et sans framework d'un transformeur avec rétropropagation complète et mises à jour Adam locales, exécutant l'entraînement, l'évaluation et la sauvegarde des points de contrôle (checkpointing) dans un seul binaire.
2. Validation BF16W : Démonstration que le stockage des poids en BF16 tout en conservant les moments en FP32 réduit l'utilisation de la SRAM de 12 octets/paramètre à 10 octets/paramètre avec une pénalité de convergence négligeable (+0,020 de perte de validation sur Shakespeare).
3. Quantification du Budget de Vocabulaire : Preuve empirique montrant que pour un matériel à budget fixe, la table d'embedding domine le modèle, nécessitant soit des vocabulaires réduits, soit des architectures multi-puces spécialisées (ex. : puces d'embedding séparées).
4. Cible de Faisabilité FPGA : Une configuration concrète (modèle 334K BF16W) qui tient entièrement dans les 4,0 Mo de BRAM sur puce d'un FPGA Xilinx ZCU102, laissant 660 Ko de marge pour les buffers d'activation, sans nécessiter d'accès à la DDR pendant l'entraînement.

Résultats

Le système a été validé sur une machine Windows 11 (AMD Ryzen 9 9900X CPU, NVIDIA RTX 4090 GPU) utilisant le corpus de Shakespeare (80K échantillons).

• Exactitude : Les vérifications de gradient numériques (différences finies vs gradients analytiques) ont réussi pour tous les composants, confirmant la justesse mathématique de l'implémentation Adam locale.
• Convergence :
  • GPU FP32 (Oracle) : A atteint une perte de validation de 1,5224.
  • CPU BF16W : A atteint une perte de validation de 1,5426.
  • L'écart de +0,020 est attribué à la différence de précision entre le stockage des poids en BF16 et le calcul en FP32.
• Qualité Générative : Le modèle a produit un texte cohérent au niveau des caractères avec des structures de dialogue et des motifs de noms de personnages reconnaissables de Shakespeare, confirmant que la boucle d'entraînement converge vers une solution fonctionnelle.
• Efficacité de la SRAM :
  • FP32 Adam (12 octets/param) : 4,00 Mo au total (remplit exactement la BRAM du ZCU102, ne laissant aucune place pour les activations).
  • BF16W Adam (10 octets/param) : 3,34 Mo au total (laisse 660 Ko libres pour les activations et l'état de contrôle).

Signification et Revendications

L'article précise explicitement sa portée modeste : NeuronFabric est un prototype de recherche, pas un accélérateur de production. Il ne contient pas de mesures FPGA ni de résultats de clôture de timing.

La principale signification de l'article réside dans :

1. Établir l'Antériorité : Il sert de divulgation architecturale publique et de référence logicielle pour établir l'architecture de « mise à jour Adam locale » comme antériorité pour de futures implémentations FPGA/ASIC.
2. Preuve de Concept : Il démontre que les mathématiques fonctionnent et que les chiffres sont corrects pour un transformeur qui s'entraîne entièrement sur puce avec un état d'optimiseur local, sans CPU hôte dans la boucle de mise à jour.
3. Définir la Voie à Suivre : Il établit le modèle 334K BF16W de Shakespeare comme une cible falsifiable pour une future implémentation FPGA. Si la même courbe de perte peut être atteinte sur le FPGA sans intervention de l'hôte, la viabilité de l'entraînement de transformeurs sur puce devient un fait mesurable plutôt qu'une projection logicielle.

Les auteurs concluent que la prochaine étape immédiate est d'implémenter cette boucle d'entraînement spécifique sur le FPGA ZCU102 pour vérifier que les résultats logiciels se traduisent par du matériel, validant ainsi le paradigme d'entraînement « sans CPU hôte ».