FFN Fusión: Replanteando el Cálculo Secuencial en Modelos de Lenguaje a Gran Escala

FFN Fusion: Rethinking Sequential Computation in Large Language Models

March 24, 2025

Autores: 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

Resumen

Presentamos FFN Fusion, una técnica de optimización arquitectónica que reduce el cómputo secuencial en modelos de lenguaje grandes al identificar y aprovechar oportunidades naturales de paralelización. Nuestra idea clave es que las secuencias de capas de Redes Feed-Forward (FFN), particularmente aquellas que permanecen tras la eliminación de capas de atención específicas, a menudo pueden paralelizarse con un impacto mínimo en la precisión. Desarrollamos una metodología fundamentada para identificar y fusionar dichas secuencias, transformándolas en operaciones paralelas que reducen significativamente la latencia de inferencia mientras preservan el comportamiento del modelo. Al aplicar estas técnicas a Llama-3.1-405B-Instruct, creamos Llama-Nemotron-Ultra-253B-Base (Ultra-253B-Base), un modelo eficiente y próximamente disponible públicamente que logra una aceleración de 1.71X en la latencia de inferencia y un costo 35X menor por token, manteniendo un rendimiento sólido en diversos benchmarks. A través de experimentos exhaustivos en modelos que van desde 49B hasta 253B parámetros, demostramos que FFN Fusion se vuelve cada vez más efectiva a escalas mayores y puede complementar técnicas de optimización existentes como la cuantización y la poda. Más intrigantemente, encontramos que incluso bloques completos de transformadores que contienen tanto capas de atención como FFN pueden paralelizarse en ocasiones, sugiriendo nuevas direcciones para el diseño de arquitecturas 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.