FFN融合技术亮相:AI芯片领军企业NVIDIA的研究人员,近期公布了一项名为“FFN融合”的突破性架构优化技术。此项技术旨在应对Transformer架构中固有的串行计算瓶颈,大幅提升大语言模型(LLMs)的推理效率,为高性能AI应用的广泛部署创造条件。
大语言模型的困境:近年来,大语言模型在自然语言处理、科研及对话代理等领域展现出强大能力。然而,随着模型规模与复杂度提升,推理所需计算资源大幅增加,出现效率瓶颈。Transformer架构作为LLMs的基础,其交替的注意力机制和前馈网络(FFN)层顺序处理输入,这种串行结构在扩大规模时,显著增加GPU间的计算与通信成本,降低效率并提高部署成本,在实时AI助手等需快速生成多个令牌的场景中尤为突出。
FFN融合的核心思路:为应对这一挑战,NVIDIA研究人员提出FFN融合技术。核心是将连续、依赖松散的FFN层合并为一个更宽的单一FFN。研究人员发现,去除注意力层后,LLMs中常存在长序列连续FFN,分析表明这些FFN层间依赖极小,可并行执行。
FFN融合的数学原理:FFN融合的数学基础在于将多个串行连接的FFN权重拼接,创建一个可并行计算的等效单一模块。例如,三个顺序堆叠的FFN,通过FFN融合消除依赖,使它们能同时处理相同输入并聚合输出。理论分析显示,融合后的FFN保留与原始FFN相同的表征能力。
Ultra-253B-Base的卓越表现:NVIDIA研究人员将FFN融合应用于Meta的LLaMA-3.1-405B-Instruct模型,通过剪枝与重构创建新模型Ultra-253B-Base。实验结果显示,推理速度和资源效率显著提升。具体而言,在批量大小为32时,推理延迟降低1.71倍,每个令牌的计算成本降低35倍。
性能与效率兼顾 :更令人惊喜的是,效率提升并未牺牲模型能力。Ultra-253B-Base在多个权威基准数据集上取得优异成绩,如MMLU85.17%、MMLU-Pro 72.25%、HumanEval 86.58%、Arena Hard 84.92%、MT-Bench9.19。其参数仅2530亿,却常与原4050亿参数模型表现相当甚至更优,且通过kv-cache优化,内存使用减半。
技术验证与通用性:研究人员利用FFN层输出间的余弦距离分析识别低依赖区域,作为融合的最佳候选。FFN融合已在不同规模模型(包括490亿、700亿和2530亿参数)上得到验证,展现出良好的通用性。
技术的意义与展望:此项研究表明,通过深入分析和巧妙架构设计,可显著提升LLMs效率。FFN融合为设计更易并行、对硬件更友好的LLMs奠定基础。尽管整个Transformer模块的并行化因层间依赖更强面临更多挑战,但FFN融合的成功无疑为未来LLM效率优化指明关键方向。