Yuqi Pan Yupeng Feng Jinghao Zhuang Siyu Ding et al

Résumé

Les modèles linguistiques à grande échelle basés sur les Transformers, aujourd’hui dominants, rencontrent des goulets d’étranglement majeurs en termes d’efficacité : le coût de calcul durant l’entraînement croît quadratiquement avec la longueur des séquences, tandis que la mémoire nécessaire à l’inférence augmente linéairement, ce qui limite fortement le traitement de contextes longs. De plus, le développement de grands modèles sur des plateformes non NVIDIA soulève des défis importants en matière de stabilité et d’efficacité de l’entraînement. Pour relever ces défis, nous introduisons SpikingBrain, une famille de modèles inspirés du cerveau, conçus pour permettre un entraînement et une inférence efficaces sur des contextes longs. SpikingBrain exploite le cluster GPU MetaX et se concentre sur trois axes : (1) Architecture du modèle : des architectures d’attention linéaire et hybride-linéaire intégrant des neurones à déclenchement adaptatif (spiking neurons) ; (2) Optimisations algorithmiques : une pipeline d’entraînement efficace basée sur la conversion, ainsi qu’un cadre dédié au codage par impulsions (spike coding) ; (3) Ingénierie système : cadres d’entraînement personnalisés, bibliothèques d’opérateurs et stratégies de parallélisme spécifiquement adaptées au matériel MetaX.Grâce à ces techniques, nous avons développé deux modèles : SpikingBrain-7B, un modèle linguistique à grande échelle à attention linéaire, et SpikingBrain-76B, un modèle à mécanisme d’expert mixte (MoE) hybride-linéaire. Ces modèles démontrent la faisabilité du développement à grande échelle de modèles linguistiques sur des plateformes non NVIDIA. SpikingBrain atteint des performances comparables aux modèles de référence open-source basés sur les Transformers, tout en nécessitant uniquement environ 150 milliards de tokens pour un pré-entraînement continu. Nos modèles améliorent considérablement l’efficacité de l’entraînement sur séquences longues, tout en offrant une inférence à mémoire (partiellement) constante et un comportement événementiel basé sur des impulsions. Par exemple, SpikingBrain-7B réalise un gain de vitesse supérieur à 100× sur le temps jusqu’au premier jeton pour des séquences de 4 millions de tokens. L’entraînement reste stable pendant plusieurs semaines sur des centaines de GPU MetaX C550, avec un taux d’utilisation des FLOPs du modèle atteignant 23,4 % pour le modèle 7B. Le schéma d’impulsion proposé atteint une densité de 69,15 %, permettant ainsi un fonctionnement à faible consommation énergétique. Globalement, ce travail démontre le potentiel des mécanismes inspirés du cerveau pour façonner la prochaine génération de modèles à grande échelle, à la fois efficaces et évolutifs.