人工智能芯片巨头英伟达的研究人员近日发布了一项名为“FFN融合”(FFN Fusion)的创新架构优化技术。该技术旨在通过解决Transformer架构中固有的串行计算瓶颈,显著提升大型语言模型(LLMs)的推理效率,为更广泛地部署高性能AI应用铺平道路.

近年来,大型语言模型在自然语言处理、科学研究和对话代理等领域展现出强大的能力。然而,随着模型规模和复杂性的不断增加,其推理过程所需的计算资源也大幅增长,导致了效率瓶颈。Transformer架构是LLM的基础,其交替的注意力机制和前馈网络(FFNs)层需要按顺序处理输入。这种固有的串行结构在模型规模扩大时,会显著增加计算和GPU之间的通信成本,降低效率并提高部署成本。尤其是在需要快速生成多个token的场景(如实时AI助手)中,这个问题更加突出。

英伟达AI研究人员推出FFN融合技术:加速大型语言模型推理

为了应对这一挑战,英伟达的研究人员提出了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的优化。

英伟达AI研究人员推出FFN融合技术:加速大型语言模型推理

研究人员使用余弦距离分析FFN层之间的输出,以识别低相互依赖性的区域,这些区域是进行融合的最佳选择。FFN融合技术已在不同规模的模型(包括490亿、700亿和2530亿参数)上得到验证,表明其具有良好的通用性

这项研究表明,通过深入分析和巧妙的架构设计,可以显著提升LLM的效率。FFN融合为设计更并行化、更适应硬件的LLM奠定了基础。虽然完全Transformer模块的并行化由于更强的层间依赖性而面临更多挑战,但FFN融合的成功无疑为未来LLM的效率优化指明了一个重要的方向。

论文:https://arxiv.org/abs/2503.18908