FFN Fusion: Rethinking Sequential Computation in Large Language Models | Summary

This article explains the key points of FFN Fusion: Rethinking Sequential Computation in Large Language Models.

• 2025-03-24 (arXiv)
• Bercovich, Akhiad, Dabbah, Mohammad, Puny, Omri, Galil, Ido, Geifman, Amnon, Geifman, Yonatan, Golan, Izhak, Karpas, Ehud, Levy, Itay, Moshe, Zach, et al.
• NVIDIA
• Paper

Read this article in Korean

Summary

• Executing large language model layers sequentially creates inference delays through repeated communication and synchronization across graphics processing units (GPUs). The paper’s complete optimization pipeline produces Ultra-253B-Base, which generates 70.92 tokens/second versus 41.44 for Llama-3.1-405B-Instruct under identical eight-way tensor parallelism, which distributes layer computations across GPUs, on one 8×H100 node.
• FFN Fusion replaces selected sequences of feed-forward networks (FFNs), modules that transform each token independently, with one wider FFN whose components share an input. Attention pruning, which removes modules that model relationships between tokens, exposes suitable sequences; dependency measurements identify regions where replacing sequential execution causes limited disruption.
• The strongest quality results follow pruning, fusion, knowledge distillation (KD), training a derivative model to preserve its parent’s behavior, and alignment through training on instructions and human preferences. The fusion step increases generation rate from 63.38 to 70.92 tokens/second on the already-pruned 253B model, while smaller-model ablations show that retaining FFN computations preserves quality better than removing them.

1 Introduction

The introduction motivates reducing execution depth alongside quantization, which lowers numerical precision, and pruning, which removes model components. Consecutive layers impose synchronization costs across GPUs, and small distributed operations can underuse available compute. FFN Fusion concentrates retained computations into fewer, larger operations; its central empirical claim concerns weak dependencies in FFN sequences exposed by attention removal.

The two-stage workflow shows that pruning creates consecutive FFN runs before fusion reduces their execution depth. Parameter reduction comes from pruning; fusion retains the selected FFN weights while changing how their computations are executed.

2 Preliminaries

A standard transformer block applies attention and then an FFN, adding each module’s output to its input through a residual connection. The experiments use SwiGLU, an FFN that combines two input projections through an elementwise gate before an output projection, with token-level normalization that rescales each token vector by its norm and learned scale factors. Puzzle is the prior architecture-search framework used to prune or reconfigure pretrained blocks and recover quality through distillation; it creates both variable FFN widths and attention-free sequences.

3 FFN Fusion

For consecutive attention-free blocks, the parallel replacement is \(\hat f^{[i,i+c]}(X)=X+\sum_{j=0}^{c}\operatorname{FFN}^{i+j}(\eta_2(X))\), using the normalization from the sequence’s final layer. Every FFN receives the same normalized input, making its computation independent of the others. Theorem 3.1 proves that this sum can be implemented exactly as one wider SwiGLU FFN by concatenating its projection matrices; replacing the original sequential computation with this parallel sum remains an approximation.

Candidate regions are evaluated by removing one block and measuring the cosine distance, a measure of directional change, between another block’s residual contributions with and without that removal. Small distances indicate weak pairwise dependencies on Distillation Mix, the paper’s mixture of text and code data. Ultra-PreFusion, the pruned 253B model before fusion, exhibits a long attention-free region with low dependency values, but GPU memory limits require dividing it into several fusion groups.

The boxed attention-free region contains consistently weak pairwise dependencies, supporting its selection for fusion. Strong early rows show that this property does not hold uniformly across the model.

4 Producing Large-Scale Models with FFN Fusion

Puzzle first searches for a Llama-405B derivative targeting a 1.5× latency speedup and memory compatibility with an 8×H100 node containing 640 GB and a single B100 GPU containing 192 GB. The resulting 253B model has 50 consecutive attention-free blocks; 49 FFNs are fused into four groups, [66,73], [74,85], [86,100], and [101,114], leaving the final FFN separate. Before further training, fusion changes MMLU, a multitask language-understanding benchmark, from 84.23 to 82.76 and MT-Bench, a conversational evaluation, from 8.83 to 8.35.

Recovery KD uses 54B tokens at 8k context, 5B tokens each at 16k and 32k, and 0.8B tokens at 128k, raising MMLU and MT-Bench to 85.17 and 9.10. Instruction tuning and reinforcement learning from human feedback (RLHF), optimization using human preference signals, subsequently improve several reported benchmarks. A separate continual pretraining (CPT) branch, which continues training on additional data, adds 73B tokens at 8k context and 15B at 258k without alignment, improving the evaluations shown in Figure 3.

Additional continual pretraining improves all four displayed evaluations, with particularly large changes on MATH500, the mathematical evaluation, and RULER 128K, the long-context evaluation. These improvements belong to the additional-training branch and do not isolate fusion's effect.

On one 8×H100 node, generation rates are 41.44 tokens/second for Llama-405B, 63.38 for Ultra-PreFusion, and 70.92 for Ultra-253B-Base; Figure 4 specifies FP8, eight-bit floating-point arithmetic. After KD and alignment, Ultra-253B-Base exceeds its parent on all five metrics in Table 3: MMLU Instruct, an instruction-based MMLU evaluation; MMLU-Pro, a more challenging multitask language-understanding benchmark; Arena Hard, a difficult prompt benchmark; HumanEval, a code-generation benchmark; and MT-Bench. The paper also reports 35× lower per-token cost at batch size 32 and half the key-value cache memory, the storage for attention’s cached states, but provides no corresponding cost breakdown.

Alignment raises Arena Hard from 71.81 to 84.92 and improves MMLU-Pro, HumanEval, and MT-Bench, while MMLU changes from 85.17 to 85.13. This separates alignment gains from the preceding fusion and distillation stages.

The intermediate Ultra-PreFusion baseline isolates fusion's incremental speed contribution: generation rises from 63.38 to 70.92 tokens/second. The 1.71× comparison uses the 41.44 tokens/second parent baseline and therefore includes the preceding pruning stage.

The aligned 253B model improves the plotted quality-throughput tradeoff relative to the two Llama baselines under FP8 and tensor parallelism of eight. Its accuracy axis averages five evaluations after multiplying MT-Bench by ten, so the plotted score is a composite rather than one benchmark result.

The final derivative exceeds Llama-405B on each listed metric after fusion, KD, and alignment: MMLU Instruct 87.54 versus 86.53, MMLU-Pro 72.25 versus 71.92, Arena Hard 84.92 versus 72.56, HumanEval 86.58 versus 85.97, and MT-Bench 9.19 versus 9.06. The comparison establishes the complete pipeline's outcome.

5 Additional Empirical Studies

The additional studies test fusion on a 49B derivative, compare fusion with FFN removal, examine sensitivity at sequence boundaries, and investigate why parallelization preserves behavior. These experiments distinguish weak dependence on execution order from dispensability of the retained FFN computations.

5.1 FFN Fusion in 70B Scale

Puzzle-49B derives from Llama-3.1-70B-Instruct and contains attention-free runs at layers 42–51 and 53–70. Excluding each run’s final FFN, four progressively stronger fusion settings replace 26 FFNs with 12, 6, 3, or 2 layers without additional training; their combined MMLU and MT-Bench scores are 83.92, 82.84, 81.69, and 81.24, versus 84.71 for the baseline. The score is \(\operatorname{Accuracy}=(\operatorname{MMLU}+10\times\operatorname{MT\text{-}Bench})/2\), and KD raises the three-layer variant to 85.28.

Increasing fusion strength gradually lowers the combined score without further training, while the three-layer variant with KD exceeds the baseline combined score. This demonstrates both the quality cost of aggressive fusion and the role of recovery training.

5.2 Removing FFNs vs. FFN Fusion

The removal baseline drops the least important FFNs according to Puzzle’s importance scores, restricting removal to the regions eligible for fusion. Step-3 fusion gives a reported 10% latency improvement with a 3.5% accuracy decline, whereas removing 15 or 20 FFNs gives 7.5% or 11% latency improvements with 5% or 8.7% accuracy declines. After 25B tokens of KD, the fused model scores 80.56 on MMLU and 9.00 on MT-Bench, exceeding the baseline combined score while remaining slightly below its 80.73 MMLU.

Fusion retains higher combined accuracy than FFN removal at nearby latency values. The KD-recovered fusion point rises above the baseline's combined accuracy while retaining the fused architecture's lower latency.

5.3 The Final FFN in Each Sequence is Sensitive to Fusion

Boundary-layer ablations support leaving the last FFN outside each fused group, especially when both runs are fused. Fusing [42,50] and [53,69] gives MT-Bench 8.25, while including both final FFNs through [42,51] and [53,70] reduces it to 7.30. The single-run evidence is less uniform: shifting [42,50] to [43,51] slightly increases MT-Bench from 8.39 to 8.42, whereas shifting [53,69] to [54,70] lowers it from 8.49 to 8.12.

Including final FFNs in both fused runs substantially reduces MT-Bench, supporting the practice of leaving these FFNs separate. The isolated first-run comparison is mixed, so the table supports a practical tendency rather than a universal rule for every metric.

5.4 Fusion Explainability

The proposed explanation is that normalization makes input direction important, while FFNs in fusion regions make relatively small changes to that direction. Per-layer measurements show both smaller input-output cosine distances and smaller residual-to-input norm ratios, measuring the size of a block’s update relative to its input, in these regions; reversing selected FFN sequences also produces scores close to fusion. This explanation assumes that small input changes cause small output changes and supplies empirical support rather than a formal error bound for replacing sequential execution.

For four selected sequences, reversing execution order gives combined scores close to those obtained by fusion; for [53,60], the scores are 83.25 and 83.24. This supports weak order sensitivity in these selected regions without establishing that the computations can be discarded.

Both model scales show lower directional changes and smaller residual contributions in the regions used for fusion. These measurements support the small-update explanation, while higher values outside those regions caution against applying it uniformly.

6 Block Parallelization

Full block parallelization runs complete attention-plus-FFN blocks independently, so it cannot use the same simple weight concatenation. A greedy search selects nonoverlapping groups of four consecutive full blocks with the smallest maximum pairwise dependency, breaking ties by the dependency sum; on Puzzle-49B, parallelizing [38,41] and [71,74] changes MMLU from 80.73 to 80.04 and MT-Bench from 8.87 to 8.67. Adding [32,35] lowers the scores to 77.86 and 8.20, and adding [26,29] lowers them to 56.12 and 6.62; implementation uses HuggingFace Transformers, a model implementation library, with optimized-framework speedups left unmeasured.

The first two selected full-block groups preserve scores relatively closely, but the third and fourth groups produce increasingly severe degradation. This limits the evidence for extending the approach from FFN-only sequences to complete transformer blocks.

The dependency heatmap and four-block statistics show that low-dependency FFN-only regions differ from the full-block candidates. The search excludes attention-free regions and ranks remaining groups using their maximum dependency and dependency sum.

7 Concluding Remarks

The conclusion emphasizes that attention-free regions can contain useful FFN computations whose execution order is only weakly constrained, with final FFNs often more sensitive to fusion. The evidence covers 49B and 253B derivatives and a complete 253B optimization pipeline. Extensions to mixture-of-experts (MoE), architectures that activate selected expert modules for each token, and broader block parallelization are proposed research directions.

Appendix

• A Proof of Theorem 3.1 establishes the two-FFN case by concatenating the input and gate projections and combining the output projections. Associativity of addition extends the construction to more FFNs. This proves exact equivalence between a shared-input sum and a wider FFN, without proving equivalence to the original sequential residual blocks.
• B Datasets identifies Distillation Mix as 224 billion tokens from FineWeb, Dolma, and Buzz-V1.2, covering source code, Wikipedia, books, news, and other domains. Ultra-253B-Base KD additionally uses synthetic data generated by Llama-3.1-405B-Instruct following the cited data-generation approach.
• C Ultra-253B-Base and Puzzle-49B Architecture Overview describes the 253B configuration across 126 blocks, with FFN width multipliers from 0.4875 to 4.875 and 50 consecutive layers that bypass attention through no_op, an operation that leaves its input unchanged. Puzzle-49B has a similar pattern of variable widths and attention removal at a smaller scale.
• D MoE Inference Performance discusses how sparse expert activation can reduce computation while introducing deployment overheads. Its argument motivates larger dense fused operations, but the appendix does not present a measured MoE comparison.
• D.1 Problems with Small Blocks attributes inefficiency to underused GPU compute and all-reduce communication, the synchronization operation that combines distributed results. The authors argue that small experts and routing modules, which select experts, are particularly exposed to these overheads, while fusion increases operation size and reduces synchronization points.
• D.2 Batch Size Effects argues that increasing batch size activates more experts before expert coverage saturates, limiting efficiency at intermediate workloads. It also argues that verifying a few proposed tokens during speculative decoding, generation using a draft model followed by verification, can encounter this regime. These broad claims are qualitative and lack batch-size measurements in this appendix.
• E Further Pairwise Block Dependency Analysis expands the Ultra-PreFusion heatmap interpretation. Strong early-block dependencies, local bands of sequential reliance, and stronger dependencies near the final layers coexist with the weakly coupled attention-free region. These patterns support selecting fusion regions individually rather than parallelizing uniformly across depth.

Brief Thoughts

The removal ablation provides the clearest evidence for preserving FFN computations while changing their execution order: removal causes greater quality loss at comparable latency. The 1.71× headline includes Puzzle pruning, while fusion’s incremental generation-rate gain is approximately 1.12×, from 63.38 to 70.92 tokens/second, and the best quality results include recovery training and alignment. The experiments support fusion in selected attention-free regions of the tested derivatives; optimized full-block speedups remain unmeasured, the reported 35× cost reduction lacks a breakdown, and the two tested model scales do not establish a general scaling law.