人工智能芯片巨头英伟达的研究人员近日发布了一项名为“FFN融合”（FFN Fusion）的创新架构优化技术。该技术旨在通过解决Transformer架构中固有的串行计算瓶颈，显著提升大型语言模型（LLMs）的推理效率，为更广泛地部署高性能AI应用铺平道路.
近年来，大型语言模型在自然语言处理、科学研究和对话代理等领域展现出强大的能力。然而，随着模型规模和复杂性的不断增加，其推理过程所需的计算资源也大幅增长，导致了效率瓶颈。Transformer架构是LLM的基础，其交替的注意力机制和前馈网络（FFNs）层需要按顺序处理输入。这种固有的串行结构在模型规模扩大时，会显著增加计算和GPU之间的通信成本，降低效率并提高部署成本。尤其是在需要快速生成多个token的场景（如实时AI助手）中，这个问题更加突出。
为了应对这一挑战，英伟达的研究人员提出了FFN融合技术。该方法的核心思想是将模型中连续的、相互依赖性较低的FFN层合并为一个更宽的FFN。研究人员观察到，在移除注意力层后，LLM中通常会存在较长的连续FFN序列。通过分析这些序列，他们发现这些FFN层之间的依赖性很小，因此可以并行执行。
FFN融合的数学基础在于将多个串联FFN的权重进行拼接，从而创建一个等效的、可以并行计算的单一模块。例如，如果三个FFN顺序堆叠，每个FFN的输出作为下一个的输入，FFN融合则消除了这种依赖关系，使得这三个FFN能够同时处理相同的输入，并将其输出进行聚合。理论分析表明，融合后的FFN保持了与原始FFN相同的表示能力。
Ultra-253B-Base:性能与效率的双重提升
英伟达的研究人员将FFN融合技术应用于Meta的Llama-3.1-405B-Instruct模型，通过剪枝和重构创建了一个名为Ultra-253B-Base的新模型。实验结果表明，Ultra-253B-Base在推理速度和资源效率方面取得了显著的提升。具体来说，该模型在批量大小为32时，推理延迟降低了1.71倍，每个token的计算成本降低了35倍。
更令人印象深刻的是，效率的提升并没有以牺牲模型能力为代价。Ultra-253B-Base在多个权威评测基准上取得了优异的成绩，例如:MMLU85.17%，MMLU-Pro72.25%，HumanEval86.58%，Arena Hard84.92%，MT-Bench9.19。这些结果通常与原始的4050亿参数模型相当甚至更高，而Ultra-253B-Base仅包含2530亿参数。此外，该模型的内存使用量也减少了一半，这得益于kv-cache的优化。
研究人员使用余弦距离分析FFN层之间的输出，以识别低相互依赖性的区域，这些区域是进行融合的最佳选择。FFN融合技术已在不同规模的模型（包括490亿、700亿和2530亿参数）上得到验证，表明其具有良好的通用性。
这项研究表明，通过深入分析和巧妙的架构设计，可以显著提升LLM的效率。FFN融合为设计更并行化、更适应硬件的LLM奠定了基础。虽然完全Transformer模块的并行化由于更强的层间依赖性而面临更多挑战，但FFN融合的成功无疑为未来LLM的效率优化指明了一个重要的方向。
论文:https://arxiv.org/abs/2503.18908
