FFN Fusion : Repenser le calcul séquentiel dans les grands modèles de langage

FFN Fusion: Rethinking Sequential Computation in Large Language Models

March 24, 2025

Auteurs: Akhiad Bercovich, Mohammad Dabbah, Omri Puny, Ido Galil, Amnon Geifman, Yonatan Geifman, Izhak Golan, Ehud Karpas, Itay Levy, Zach Moshe, Najeeb Nabwani, Tomer Ronen, Itamar Schen, Elad Segal, Ido Shahaf, Oren Tropp, Ran Zilberstein, Ran El-Yaniv

cs.AI

Résumé

Nous présentons FFN Fusion, une technique d'optimisation architecturale qui réduit les calculs séquentiels dans les grands modèles de langage en identifiant et en exploitant les opportunités naturelles de parallélisation. Notre idée clé est que les séquences de couches de réseaux feed-forward (FFN), en particulier celles qui subsistent après la suppression de certaines couches d'attention, peuvent souvent être parallélisées avec un impact minimal sur la précision. Nous développons une méthodologie rigoureuse pour identifier et fusionner ces séquences, les transformant en opérations parallèles qui réduisent considérablement la latence d'inférence tout en préservant le comportement du modèle. En appliquant ces techniques à Llama-3.1-405B-Instruct, nous créons Llama-Nemotron-Ultra-253B-Base (Ultra-253B-Base), un modèle efficace et bientôt disponible publiquement, qui atteint une accélération de 1,71X en latence d'inférence et un coût par token 35 fois inférieur, tout en maintenant des performances solides sur les benchmarks. À travers des expériences approfondies sur des modèles de 49B à 253B paramètres, nous démontrons que FFN Fusion devient de plus en plus efficace à plus grande échelle et peut compléter des techniques d'optimisation existantes comme la quantification et l'élagage. Plus intriguant encore, nous constatons que même des blocs de transformateurs complets contenant à la fois des couches d'attention et des FFN peuvent parfois être parallélisés, suggérant de nouvelles directions pour la conception d'architectures neuronales.

English

We introduce FFN Fusion, an architectural optimization technique that reduces sequential computation in large language models by identifying and exploiting natural opportunities for parallelization. Our key insight is that sequences of Feed-Forward Network (FFN) layers, particularly those remaining after the removal of specific attention layers, can often be parallelized with minimal accuracy impact. We develop a principled methodology for identifying and fusing such sequences, transforming them into parallel operations that significantly reduce inference latency while preserving model behavior. Applying these techniques to Llama-3.1-405B-Instruct, we create Llama-Nemotron-Ultra-253B-Base (Ultra-253B-Base), an efficient and soon-to-be publicly available model that achieves a 1.71X speedup in inference latency and 35X lower per-token cost while maintaining strong performance across benchmarks. Through extensive experiments on models from 49B to 253B parameters, we demonstrate that FFN Fusion becomes increasingly effective at larger scales and can complement existing optimization techniques like quantization and pruning. Most intriguingly, we find that even full transformer blocks containing both attention and FFN layers can sometimes be parallelized, suggesting new directions for neural architecture design.