英偉達推出FFN融合技術，大幅提升大語言模型推理效率

近期，人工智能芯片领域的领军企业英伟达的研究团队推出了一项名为“FFN融合”的创新技术。这项技术针对Transformer架构中的串行计算瓶颈进行了优化，旨在大幅提升大型语言模型（LLMs）的推理效率，为高性能AI应用的广泛应用提供支持。 近年来，随着LLMs在自然语言处理、科学研究及对话代理等领域展现的强大能力，其模型规模和复杂度也随之迅速增长。然而，这带来了更高的计算需求，导致推理效率下降。Transformer架构虽是LLM的基础，但其由注意力机制和前馈网络（FFNs）组成的交替结构需要按顺序处理输入，这在大规模模型中会显著增加计算和GPU之间的通信开销，尤其在需要快速生成多个token的应用场景中表现得更为明显。 为了解决上述问题，英伟达的研究人员提出了一种新的优化策略——FFN融合技术。该技术的核心是将模型中连续且相互依赖性较低的FFN层整合成一个更宽的FFN模块。研究发现，在去除注意力层后，LLM中通常会出现较长的连续FFN序列，而这些FFN层之间实际上依赖性较小，因此可以实现并行计算。通过将多个串联FFN的权重拼接，形成一个可以并行计算的单一模块，FFN融合不仅提高了计算效率，还确保了与原始FFN相同的表示能力。 为了验证FFN融合的实际效果，研究团队将其应用于Meta的Llama-3.1-405B-Instruct模型，通过剪枝和重构得到了一个新的模型——Ultra-253B-Base。实验结果显示，Ultra-253B-Base在推理速度和资源利用方面有了显著提升。例如，在批量大小为32的情况下，推理延迟减少了1.71倍，每个token的计算成本更是降低了35倍。尽管参数数量从4050亿减少到了2530亿，该模型在多个权威评测基准上的表现依然出色，包括MMLU得分85.17%，MMLU-Pro得分72.25%，HumanEval得分86.58%，Arena Hard得分84.92%，以及MT-Bench得分9.19。此外，Ultra-253B-Base的内存使用量也减少了约一半，这得益于kv-cache的优化。 进一步研究显示，FFN融合技术适用于不同规模的模型，从490亿参数到700亿参数再到2530亿参数，均能取得良好效果，体现了该技术的通用性。这一突破不仅提升了LLM的推理效率，也为未来设计更加并行化、更适合硬件特性的LLM提供了重要参考。 通过这项研究可以看出，通过对现有架构的深入分析与创新性改造，可以在不牺牲模型能力的前提下实现效率的巨大飞跃。虽然完全并行化Transformer模块仍面临诸多挑战，但FFN融合技术的成功无疑为未来LLM的优化开辟了一条新路径。相关论文已发布于arXiv平台，供全球科研人员查阅。（文中图片详见原文链接）