Les chercheurs de Nvidia, géant des puces d'intelligence artificielle, ont récemment publié une technique d'optimisation architecturale innovante appelée « FFN Fusion ». Cette technique vise à résoudre les goulots d'étranglement de calculs séquentiels inhérents à l'architecture Transformer, améliorant ainsi considérablement l'efficacité de l'inférence des grands modèles de langage (LLM) et ouvrant la voie à un déploiement plus large d'applications IA hautes performances.
Ces dernières années, les grands modèles de langage ont démontré de puissantes capacités dans les domaines du traitement du langage naturel, de la recherche scientifique et des agents conversationnels. Cependant, avec l'augmentation constante de la taille et de la complexité des modèles, les ressources de calcul nécessaires à leur processus d'inférence ont également considérablement augmenté, entraînant des goulots d'étranglement d'efficacité. L'architecture Transformer est à la base des LLM, et son mécanisme d'attention alterné et ses couches de réseaux de transmission en avant (FFN) nécessitent un traitement séquentiel des entrées. Cette structure séquentielle inhérente, lorsque la taille du modèle augmente, augmente considérablement les coûts de calcul et de communication entre les GPU, réduisant l'efficacité et augmentant les coûts de déploiement. Ce problème est particulièrement marqué dans les scénarios nécessitant une génération rapide de plusieurs jetons (comme les assistants IA en temps réel).
Pour relever ce défi, les chercheurs de Nvidia ont proposé la technique FFN Fusion. L'idée centrale de cette méthode consiste à fusionner les couches FFN consécutives et faiblement dépendantes en un FFN plus large. Les chercheurs ont observé qu'après la suppression des couches d'attention, il existe généralement de longues séquences de FFN consécutives dans les LLM. En analysant ces séquences, ils ont constaté que la dépendance entre ces couches FFN était faible, et qu'elles pouvaient donc être exécutées en parallèle.
Les bases mathématiques de FFN Fusion reposent sur la concaténation des poids de plusieurs FFN en série pour créer un seul module équivalent calculable en parallèle. Par exemple, si trois FFN sont empilés séquentiellement, la sortie de chaque FFN servant d'entrée au suivant, FFN Fusion supprime cette dépendance, permettant aux trois FFN de traiter simultanément la même entrée et d'agréger leurs sorties. L'analyse théorique montre que le FFN fusionné conserve la même capacité de représentation que le FFN d'origine.
Ultra-253B-Base : double amélioration des performances et de l'efficacité
Les chercheurs de Nvidia ont appliqué la technique FFN Fusion au modèle Llama-3.1-405B-Instruct de Meta, créant un nouveau modèle appelé Ultra-253B-Base par élagage et reconstruction. Les résultats expérimentaux montrent que Ultra-253B-Base a considérablement amélioré la vitesse d'inférence et l'efficacité des ressources. Plus précisément, ce modèle a permis de réduire le délai d'inférence de 1,71 fois et le coût de calcul par jeton de 35 fois pour une taille de lot de 32.
Plus impressionnant encore, cette amélioration de l'efficacité ne s'est pas faite au détriment des capacités du modèle. Ultra-253B-Base a obtenu d'excellents résultats sur plusieurs benchmarks d'évaluation reconnus, tels que : MMLU 85,17 %, MMLU-Pro 72,25 %, HumanEval 86,58 %, Arena Hard 84,92 %, MT-Bench 9,19. Ces résultats sont généralement comparables, voire supérieurs, à ceux du modèle d'origine à 405 milliards de paramètres, alors qu'Ultra-253B-Base ne contient que 253 milliards de paramètres. De plus, la consommation mémoire du modèle a été réduite de moitié, grâce à l'optimisation du kv-cache.
Les chercheurs ont utilisé la distance cosinus pour analyser les sorties entre les couches FFN afin d'identifier les zones à faible interdépendance, qui sont les meilleures candidates à la fusion. La technique FFN Fusion a été validée sur des modèles de différentes tailles (y compris 49 milliards, 700 milliards et 253 milliards de paramètres), démontrant sa bonne généralisation.
Cette étude montre que, grâce à une analyse approfondie et à une conception architecturale astucieuse, il est possible d'améliorer considérablement l'efficacité des LLM. FFN Fusion jette les bases de la conception de LLM plus parallélisés et mieux adaptés au matériel. Bien que la parallélisation des modules Transformer complets soit confrontée à des défis plus importants en raison d'une dépendance inter-couches plus forte, le succès de FFN Fusion indique sans aucun doute une voie importante pour l'optimisation de l'efficacité des futurs LLM.
Article : https://arxiv.org/abs/2503.18908
