FFN Fusion: Переосмысление последовательных вычислений в больших языковых моделях

FFN Fusion: Rethinking Sequential Computation in Large Language Models

March 24, 2025

Авторы: 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

Аннотация

Мы представляем FFN Fusion — метод оптимизации архитектуры, который сокращает последовательные вычисления в больших языковых моделях за счет выявления и использования естественных возможностей для параллелизации. Наше ключевое наблюдение заключается в том, что последовательности слоев Feed-Forward Network (FFN), особенно оставшиеся после удаления определенных слоев внимания, часто могут быть параллелизованы с минимальным влиянием на точность. Мы разработали принципиальную методологию для выявления и объединения таких последовательностей, преобразуя их в параллельные операции, которые значительно сокращают задержку вывода при сохранении поведения модели. Применив эти методы к модели Llama-3.1-405B-Instruct, мы создали Llama-Nemotron-Ultra-253B-Base (Ultra-253B-Base) — эффективную модель, которая скоро станет общедоступной и демонстрирует ускорение задержки вывода в 1.71 раза и снижение стоимости обработки одного токена в 35 раз при сохранении высокой производительности на тестовых наборах. В ходе обширных экспериментов с моделями от 49B до 253B параметров мы показываем, что FFN Fusion становится все более эффективным на больших масштабах и может дополнять существующие методы оптимизации, такие как квантование и обрезка. Наиболее интригующе то, что мы обнаружили, что даже полные блоки трансформера, содержащие как слои внимания, так и FFN, иногда могут быть параллелизованы, что открывает новые направления в проектировании нейронных архитектур.

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.