NVIDIA FFN Fusion : Efficacité accrue des LLM

Le défi computationnel de l'IA moderne

Les grands modèles de langage (LLM) sont les piliers de l’intelligence artificielle contemporaine, démontrant des capacités remarquables qui remodèlent les industries et la découverte scientifique. Leur compétence à générer du texte de type humain, à alimenter des agents conversationnels sophistiqués et même à aider dans des tâches de recherche complexes en a fait des outils indispensables. Au cœur de ces modèles puissants bat l’architecture transformer, une conception caractérisée par ses couches alternées. Les données d’entrée, décomposées en tokens, traversent une séquence de mécanismes d’attention, qui pondèrent l’importance des différents tokens, suivis par des réseaux feed-forward (FFN), qui traitent les informations glanées. Ce traitement séquentiel en couches est fondamental pour la manière dont les transformeurs apprennent et génèrent des sorties.

Cependant, cette architecture même, bien qu’efficace, présente un défi croissant à mesure que les modèles augmentent en taille et en complexité. La nature séquentielle signifie que chaque couche doit généralement attendre que la précédente ait terminé son calcul avant de pouvoir commencer. Ce traitement étape par étape crée un goulot d’étranglement inhérent, en particulier pendant la phase d’inférence – l’étape où un modèle entraîné est réellement utilisé pour générer des prédictions ou du texte. Alors que des modèles comme ceux qui alimentent les assistants IA avancés intègrent des centaines de milliards, voire des trillions, de paramètres, les ressources computationnelles et le temps requis pour l’inférence augmentent de façon spectaculaire. Cette demande croissante se traduit par une latence significative (délai de réponse), un débit réduit (nombre de requêtes traitées dans le temps) et des coûts opérationnels croissants, entravant le déploiement généralisé et l’application en temps réel des LLM les plus puissants. Par conséquent, l’amélioration de l’efficacité de l’inférence est devenue une préoccupation primordiale au sein de la communauté de recherche en IA, stimulant une quête de stratégies innovantes capables de rationaliser le calcul sans compromettre les performances remarquables offertes par ces modèles. Le défi central réside dans l’atténuation des contraintes imposées par l’exécution séquentielle, en particulier dans les environnements distribués où les calculs s’étendent sur plusieurs GPU, ajoutant une surcharge de communication au temps de traitement.

Naviguer dans le paysage de l'optimisation :Outils existants et leurs limites

Dans l’effort continu pour rendre les LLM plus légers et plus rapides, les chercheurs ont développé une boîte à outils de techniques d’optimisation. Chacune offre une voie vers l’efficacité, mais s’accompagne souvent de ses propres compromis, empêchant une méthode unique d’être une solution universelle. Comprendre ces compromis est crucial pour apprécier le besoin d’approches novatrices comme FFN Fusion.

Une technique prédominante est la quantification. Elle consiste à réduire la précision numérique utilisée pour représenter les poids et les activations du modèle. Au lieu d’utiliser des nombres à virgule flottante standard de 32 bits, les modèles peuvent utiliser des représentations de 16 bits, 8 bits, voire moins. Cela réduit directement l’empreinte mémoire du modèle et peut accélérer considérablement les calculs, car les opérations sur des nombres de moindre précision sont généralement plus rapides et nécessitent moins d’énergie. Cependant, la quantification n’est pas sans risque. La réduction de la précision peut entraîner une perte d’informations, dégradant potentiellement la précision du modèle. Ce risque devient plus prononcé à de très faibles largeurs de bits, nécessitant une mise en œuvre minutieuse et parfois un réentraînement pour atténuer les baisses de précision. Le défi consiste à trouver le juste milieu qui maximise les gains d’efficacité tout en maintenant la dégradation des performances dans des limites acceptables.

Une autre stratégie courante est l’élagage (pruning). Cette technique repose sur le principe que de nombreux paramètres au sein d’un grand réseau neuronal peuvent être redondants ou contribuer de manière minimale à la sortie finale. Les algorithmes d’élagage identifient et suppriment ces connexions ou neurones moins importants, résultant en un modèle plus petit et plus épars (sparse). Comme la quantification, l’élagage réduit les besoins en mémoire et la charge de calcul. Cependant, identifier précisément quels paramètres sont ‘sûrs’ à supprimer est complexe. Un élagage agressif peut par inadvertance supprimer des composants cruciaux, entraînant une perte de précision substantielle. Un ajustement fin (fine-tuning) du modèle après l’élagage est souvent nécessaire pour récupérer les performances, ajoutant de la complexité au flux de travail. Un calibrage minutieux est essentiel pour garantir que le modèle élagué reste efficace.

Une approche architecturalement plus distincte est le modèle Mixture-of-Experts (MoE). Au lieu de traiter chaque entrée à travers l’ensemble du réseau, les modèles MoE se composent de plusieurs sous-réseaux ‘experts’ (généralement des FFN). Pour chaque token d’entrée, un mécanisme de routage (gating) sélectionne dynamiquement un petit sous-ensemble de ces experts pour effectuer le calcul. Ce calcul conditionnel signifie que seule une fraction des paramètres totaux du modèle est activée pour une entrée donnée, conduisant à des économies de calcul significatives, en particulier lors de l’entraînement et de l’inférence sur de très grands modèles. Les modèles MoE peuvent atteindre des trillions de paramètres tout en maintenant des coûts de calcul raisonnables. Cependant, leur efficacité dépend fortement de la charge de travail. Ils excellent dans le traitement de très grandes tailles de lots (batch sizes) où le schéma d’activation sélective conduit à une bonne utilisation du matériel. À des tailles de lots plus petites ou intermédiaires, les modèles MoE peuvent souffrir d’une sous-utilisation des ressources de calcul, car le matériel parallèle peut ne pas être constamment occupé par les experts activés de manière éparse. De plus, la mise en œuvre et l’équilibrage de charge des modèles MoE peuvent être plus complexes que le déploiement d’architectures ‘denses’ standard.

Bien que la quantification, l’élagage et les modèles MoE représentent des avancées précieuses dans l’optimisation des LLM, leurs limitations inhérentes soulignent le besoin de stratégies alternatives ou complémentaires. La quête se poursuit pour des méthodes capables d’apporter des améliorations d’efficacité générales dans divers scénarios, idéalement avec moins de compromis sur la précision ou la complexité de mise en œuvre, en particulier pour les architectures de modèles denses qui restent populaires en raison de leur relative simplicité d’entraînement et de déploiement.

FFN Fusion : Repenser le parallélisme dans les Transformeurs

Au milieu de ce paysage de techniques d’optimisation, des chercheurs de NVIDIA ont introduit une nouvelle approche convaincante appelée FFN Fusion. Cette technique s’attaque directement au goulot d’étranglement séquentiel inhérent à l’architecture transformer, non pas en modifiant les paramètres ou en activant sélectivement des parties, mais en repensant fondamentalement la manière dont les séquences de calculs peuvent être parallélisées. L’innovation découle d’une observation cruciale sur le comportement des couches FFN au sein des modèles transformer profonds.

À l’aide d’un outil de diagnostic nommé Puzzle, les chercheurs ont analysé le fonctionnement interne de grands modèles. Lorsqu’ils ont expérimentalement supprimé les couches d’attention, ils ont remarqué que les modèles conservaient souvent des séquences étonnamment longues de couches FFN consécutives. Plus important encore, l’analyse a révélé que les calculs effectués par ces FFN adjacents présentaient fréquemment une interdépendance minimale. Essentiellement, la sortie d’un FFN dans la séquence ne modifiait souvent pas radicalement le chemin directionnel ou les informations essentielles nécessaires au FFN immédiatement suivant. Cela suggérait que ces FFN, traditionnellement exécutés l’un après l’autre, pourraient potentiellement être exécutés simultanément, en parallèle, sans perturber de manière significative la fonction globale du modèle.

Cette intuition a constitué la base de FFN Fusion. L’idée maîtresse est élégamment simple mais puissante : identifier les séquences de couches FFN consécutives à faible dépendance computationnelle et les fusionner en une seule couche FFN plus large qui effectue le calcul équivalent en parallèle. Au lieu d’une chaîne comme Entrée -> FFN1 -> FFN2 -> FFN3 -> Sortie, la structure fusionnée devient Entrée -> FFN_Fusionné (Équivalent à FFN1+FFN2+FFN3 en parallèle) -> Sortie. Cette transformation architecturale raccourcit efficacement la profondeur séquentielle du réseau, remplaçant plusieurs étapes par une seule étape de calcul plus large. En ciblant ces séquences FFN à faible dépendance, FFN Fusion vise à réduire la latence et le coût de calcul tout en préservant la puissance de représentation et la précision du modèle. Le développement d’Ultra-253B-Base à partir de Llama-3.1-405B-Instruct a servi de démonstration principale du potentiel de cette technique.

L'alchimie architecturale : Comment fonctionne FFN Fusion

La magie derrière FFN Fusion réside dans sa manipulation astucieuse de la structure mathématique sous-jacente des réseaux feed-forward. Il ne s’agit pas simplement d’exécuter les couches existantes côte à côte ; cela implique de créer une nouvelle couche unifiée qui réplique le comportement collectif de la séquence originale mais le fait simultanément.

Considérez une séquence de k couches FFN consécutives. Dans un transformeur standard, l’entrée x passe par FFN1, sa sortie devient l’entrée pour FFN2, et ainsi de suite, jusqu’à FFNk. Chaque étape dépend explicitement de l’achèvement de la précédente. FFN Fusion brise cette chaîne de dépendance. Mathématiquement, un FFN implique généralement deux transformations linéaires avec une fonction d’activation non linéaire (comme GeLU ou SwiGLU) entre les deux : FFN(x) = W_out * Activation(W_in * x). FFN Fusion tire parti du fait que les transformations linéaires peuvent souvent être combinées.

Le processus de fusion fonctionne en concaténant les poids des couches FFN individuelles. Plus précisément, les matrices de poids d’entrée (W_in) des FFN consécutifs sont combinées (par exemple, en bloc diagonal) en une seule matrice de poids d’entrée plus grande pour la couche fusionnée. De même, les matrices de poids de sortie (W_out) sont concaténées pour former une seule matrice de poids de sortie plus large. La fonction d’activation est appliquée élément par élément au sein de cette structure plus grande. Cette construction garantit que le FFN fusionné opère sur l’entrée originale x simultanément à travers des chemins parallèles correspondant aux FFN originaux. Les sorties de ces chemins parallèles sont ensuite implicitement agrégées par la structure des poids de sortie concaténés.

Le fondement théorique confirme que cette structure fusionnée peut maintenir la même capacité de représentation que la séquence originale de FFN, à condition que les dépendances entre les couches originales soient effectivement faibles. La clé est d’identifier quelles séquences sont adaptées à la fusion. Pour ce faire systématiquement, les chercheurs de NVIDIA ont utilisé une technique d’analyse de dépendance. Ils ont mesuré la distance cosinus entre les états cachés de sortie des couches FFN consécutives pour un ensemble représentatif de tokens d’entrée. Une faible distance cosinus indique que le vecteur de sortie d’un FFN pointe dans une direction très similaire à celle du vecteur de sortie du FFN suivant dans la séquence. Cette similarité suggère une faible dépendance fonctionnelle – le second FFN ne change pas radicalement la représentation de l’information établie par le premier. Les séquences de FFN présentant des distances cosinus constamment faibles entre les couches ont été identifiées comme des candidats de choix pour la fusion, car leur fusion était moins susceptible de perturber les représentations apprises par le modèle et ses performances globales. Cette approche basée sur les données permet une application ciblée de FFN Fusion aux parties du modèle où elle sera la plus efficace et la moins perturbatrice.

Du Béhémoth au Sprinter : La transformation Ultra-253B-Base

La puissance pratique de FFN Fusion a été démontrée de manière éclatante par son application à l’un des plus grands modèles connus publiquement à l’époque, Llama-3.1-405B-Instruct. Ce modèle, doté de 405 milliards de paramètres, représentait une entreprise computationnelle significative pour l’inférence. Les chercheurs ont entrepris un processus de raffinement architectural, combinant FFN Fusion avec un élagage stratégique, pour créer un nouveau modèle plus efficace baptisé Ultra-253B-Base.

Le processus de transformation a impliqué plusieurs étapes :

1. Analyse : À l’aide de leurs outils d’analyse de dépendance (mesurant les distances cosinus), les chercheurs ont identifié des séquences de couches FFN consécutives au sein de l’architecture Llama-405B qui présentaient une faible dépendance inter-couches.
2. Fusion : Ces séquences FFN identifiées ont ensuite été fusionnées en couches FFN uniques et plus larges comme décrit précédemment (concaténation des poids). Cela a directement réduit le nombre d’étapes séquentielles dans le réseau.
3. Élagage : Simultanément ou ultérieurement, les paramètres jugés moins critiques (potentiellement identifiés par des techniques d’élagage standard ou informés par le processus de fusion) ont été supprimés du modèle.

Cette approche combinée a abouti à Ultra-253B-Base, un modèle avec 253 milliards de paramètres. Cela représente une réduction substantielle – plus de 37% de paramètres en moins que le modèle original de 405B. Les changements architecturaux réalisés grâce à la fusion ont été essentiels pour permettre une réduction de taille aussi significative tout en visant à conserver les performances. L’objectif n’était pas seulement un modèle plus petit, mais un modèle fondamentalement plus rapide et plus économe en calcul, grâce au parallélisme accru débloqué par FFN Fusion. Cette étude de cas a servi de preuve de concept cruciale, montrant que les modèles à grande échelle pouvaient être substantiellement restructurés pour l’efficacité.

Mesurer les gains : Performance, vitesse et économies de ressources

Le véritable test de toute technique d’optimisation réside dans son impact mesurable. Pour Ultra-253B-Base, les résultats dérivés de l’application de FFN Fusion et de l’élagage à la base Llama-405B ont été convaincants, démontrant des améliorations significatives sur plusieurs dimensions sans compromis substantiels sur les capacités.

Vitesse d’inférence et coût : Les gains les plus frappants ont été observés dans l’efficacité de l’inférence. Comparé au modèle original de 405B paramètres, Ultra-253B-Base a atteint :

• Une amélioration de 1.71x de la latence d’inférence. Cela signifie que le modèle pouvait générer des réponses beaucoup plus rapidement, ce qui est crucial pour les applications en temps réel.
• Une réduction de 35x du coût computationnel par token mesuré à une taille de lot de 32. Cette diminution spectaculaire des opérations de calcul (FLOPs) par token se traduit directement par une consommation d’énergie plus faible et des exigences matérielles réduites pour servir le modèle.

Benchmarks de performance du modèle : Fait critique, ces améliorations d’efficacité ne se sont pas faites au détriment de l’intelligence ou des capacités du modèle. Ultra-253B-Base a été rigoureusement évalué sur une suite de benchmarks LLM standard, obtenant des scores très compétitifs, et dans certains cas supérieurs, à ceux du modèle original beaucoup plus grand :

• MMLU (Massive Multitask Language Understanding) : 85.17%
• MMLU-Pro (Une version plus difficile) : 72.25%
• Arena Hard (Évaluation de préférence humaine sur des invites difficiles) : 84.92%
• HumanEval (Capacité de génération de code) : 86.58%
• MT-Bench (Qualité de conversation multi-tours) : 9.19

Ces scores indiquent que le modèle fusionné et élagué a conservé un très haut niveau de compréhension, de raisonnement, de capacité de codage et de qualité conversationnelle, comparable à son progéniteur de 405B paramètres malgré ses seulement 253 milliards de paramètres.

Efficacité mémoire : Au-delà de la vitesse et du coût de calcul, FFN Fusion a également contribué aux économies de mémoire. Les changements architecturaux, potentiellement combinés à d’autres optimisations permises par la fusion, ont conduit à une réduction de 2x de la taille du cache clé-valeur (KV cache) requis pendant l’inférence. Le KV cache stocke les activations intermédiaires (clés et valeurs d’attention) et peut consommer une mémoire GPU substantielle, en particulier pour les longues séquences d’entrée. Réduire de moitié cette exigence permet d’exécuter le modèle sur du matériel moins gourmand en mémoire ou de traiter des contextes plus longs dans les mêmes contraintes de mémoire.

Ces résultats quantifiables soulignent l’efficacité de FFN Fusion. Il a permis la création d’un modèle qui était non seulement plus petit mais fondamentalement plus efficace en termes de vitesse, d’opérations de calcul et d’utilisation de la mémoire, tout en maintenant des performances de premier plan sur des benchmarks exigeants.

Préserver les connaissances : Le rôle crucial de l'entraînement et du fine-tuning

Modifier architecturalement un modèle de langage massif et pré-entraîné comme Llama-405B par des techniques telles que FFN Fusion et l’élagage perturbe inévitablement l’équilibre délicat de ses paramètres appris. Bien que l’équivalence mathématique vise à préserver la fonction localement, le comportement global du réseau peut changer. Pour s’assurer que le modèle Ultra-253B-Base résultant devienne non seulement plus efficace mais conserve également son haut niveau de performance, un processus d’entraînement post-modification soigneusement orchestré était essentiel.

Ce processus comprenait deux phases principales :

1. Distillation des connaissances (Knowledge Distillation) : La première étape consistait à transférer les connaissances du modèle original plus grand (ou d’un modèle enseignant approprié) vers l’architecture modifiée. Cela a été réalisé par distillation, où le modèle Ultra-253B-Base a été entraîné à imiter les sorties ou les représentations internes du modèle enseignant. Cette phase a utilisé un jeu de données substantiel, spécifiquement 54 milliards de tokens, traités avec une fenêtre de contexte de 8k. La distillation aide le modèle fusionné et élagué à recapturer les nuances et les capacités qui auraient pu être légèrement perturbées lors des changements architecturaux.

2. Fine-tuning par étapes : Après la distillation, le modèle a subi une série d’étapes de fine-tuning spécifiquement conçues pour l’adapter à la gestion de longueurs de contexte progressivement plus longues. Ceci est crucial pour les LLM modernes, qui sont souvent censés traiter et générer du texte basé sur des entrées étendues. Le fine-tuning s’est déroulé par étapes :

   • Fine-tuning avec une fenêtre de contexte de 16k.
   • Fine-tuning supplémentaire avec une fenêtre de contexte de 32k.
   • Étape finale de fine-tuning avec une fenêtre de contexte de 128k.

Cette approche par étapes permet au modèle d’adapter progressivement ses paramètres, y compris les couches FFN fusionnées nouvellement formées et les mécanismes optimisés du KV cache, pour gérer efficacement les dépendances et le flux d’informations sur de très longues séquences. Chaque étape s’appuie sur la précédente, garantissant la stabilité et des performances robustes pour différentes tailles de contexte.

Ce régime d’entraînement méticuleux, combinant une distillation à grande échelle avec un fine-tuning par étapes pour contexte long, a été déterminant pour combler le fossé entre l’efficacité architecturale et les performances haute fidélité. Il a garanti que les avantages en termes de vitesse, de coût et de mémoire apportés par FFN Fusion ne compromettent pas la précision et les capacités du modèle sur des benchmarks exigeants.

Horizons élargis : Généralisabilité et orientations futures

La transformation réussie de Llama-405B en Ultra-253B-Base fournit une preuve solide du potentiel de FFN Fusion, mais sa vraie valeur réside dans son applicabilité plus large et les perspectives qu’elle offre pour la conception future des LLM. La recherche a démontré qu’il ne s’agissait pas simplement d’une astuce ponctuelle applicable uniquement aux modèles énormes.

Validation à différentes échelles : Les chercheurs de NVIDIA ont explicitement testé la méthodologie FFN Fusion sur des modèles de tailles variées. Ils ont appliqué avec succès la technique à des modèles de 70 milliards de paramètres, obtenant des gains d’efficacité similaires par rapport à leurs homologues d’origine. Ils ont également rapporté une validation à l’échelle 49B, renforçant davantage l’idée que l’indépendance des FFN et le potentiel de fusion ne sont pas des caractéristiques exclusives des plus grands modèles mais pourraient être une propriété plus générale de l’architecture transformer, potentiellement plus prononcée à plus grande échelle où des séquences FFN plus profondes apparaissent naturellement. Cela suggère que FFN Fusion pourrait devenir un outil standard dans l’arsenal d’optimisation des LLM, applicable à une gamme de tailles de modèles.

Fusion FFN vs Fusion de blocs complets : La recherche a également éclairé le rôle spécifique des couches FFN par rapport aux couches d’attention au sein du bloc transformer. Alors que les couches FFN consécutives montraient souvent une faible dépendance, les rendant idéales pour la fusion, les tentatives de paralléliser des blocs transformer entiers (incluant à la fois les couches d’attention et FFN) se sont avérées plus difficiles. L’analyse a indiqué des interdépendances plus fortes impliquant les mécanismes d’attention. La fusion simultanée de blocs entiers a entraîné une dégradation plus significative des performances, suggérant que les couches d’attention jouent un rôle plus critique et séquentiellement dépendant dans l’intégration des informations à travers les tokens. Cette découverte aide à délimiter les frontières d’une parallélisation efficace – les séquences FFN sont un terrain fertile, tandis que les mécanismes d’attention pourraient nécessiter des stratégies d’optimisation différentes.

Implications pour l’architecture LLM : FFN Fusion offre plus qu’une simple technique d’optimisation post-hoc ; elle fournit des informations précieuses pour la conception des futurs LLM. La découverte que des séquences de FFN peuvent souvent être traitées comme des unités parallélisables remet en question l’hypothèse strictement séquentielle qui sous-tend souvent la conception des transformeurs. Cela pourrait inspirer de nouvelles architectures intrinsèquement plus adaptées au parallélisme dès le départ. Les futurs modèles pourraient être conçus avec des structures FFN explicitement destinées à la fusion ou à l’exécution parallèle, conduisant potentiellement à une co-conception matériel-logiciel où les architectures GPU sont davantage optimisées pour exploiter ce type de parallélisme. La méthode systématique utilisant la distance cosinus pour quantifier la dépendance inter-couches fournit également un outil analytique précieux pour comprendre et redessiner les structures des réseaux neuronaux. En démontrant que des gains d’efficacité significatifs sont possibles grâce à une refonte architecturale réfléchie axée sur la parallélisation des composants existants, FFN Fusion ouvre la voie au développement de LLM à la fois puissants et plus durables sur le plan computationnel. Elle met en évidence une voie vers l’atténuation des demandes croissantes en ressources de l’IA de pointe.