FFN Fusion: Repensando o Cálculo Sequencial em Modelos de Linguagem de Grande 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

Resumo

Apresentamos o FFN Fusion, uma técnica de otimização arquitetônica que reduz a computação sequencial em modelos de linguagem de grande escala ao identificar e explorar oportunidades naturais de paralelização. Nossa principal percepção é que sequências de camadas de Feed-Forward Network (FFN), particularmente aquelas que permanecem após a remoção de camadas de atenção específicas, podem frequentemente ser paralelizadas com impacto mínimo na precisão. Desenvolvemos uma metodologia fundamentada para identificar e fundir tais sequências, transformando-as em operações paralelas que reduzem significativamente a latência de inferência enquanto preservam o comportamento do modelo. Aplicando essas técnicas ao Llama-3.1-405B-Instruct, criamos o Llama-Nemotron-Ultra-253B-Base (Ultra-253B-Base), um modelo eficiente e em breve disponível publicamente que alcança um aumento de velocidade de 1,71X na latência de inferência e um custo 35X menor por token, mantendo um desempenho forte em benchmarks. Por meio de extensos experimentos em modelos com parâmetros variando de 49B a 253B, demonstramos que o FFN Fusion se torna cada vez mais eficaz em escalas maiores e pode complementar técnicas de otimização existentes, como quantização e poda. Mais intrigantemente, descobrimos que até blocos completos de transformadores contendo tanto camadas de atenção quanto FFN podem às vezes ser paralelizados, sugerindo novas direções para o design de arquiteturas neurais.

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.