大语言模型剪枝的优化路径
近年来,大语言模型(LLMs)彻底改变了自然语言处理领域,并对计算机视觉、语音识别和机器翻译做出重大贡献。LLMs有效性的关键因素之一是其训练所使用的超大规模数据集。但代价是模型规模过大,导致运行速度变慢和计算资源消耗增加。AI研究人员深知这些挑战,许多团队正在寻求在保持性能的同时使大模型更紧凑的方法。
为此,我们提出了一种新颖的“轻剪常尝”理念,专注于一种新的剪枝方法。剪枝是一种压缩过程,可移除LLM神经网络层中不重要的连接。在计算语言学协会(ACL)会议上发表的论文中,我们描述了Wanda++框架,该框架可在单个GPU上10分钟内压缩具有70亿参数的模型。
以困惑度(概率分布预测给定样本的能力)衡量,该方法比领先的前身方法Wanda提升了32%的性能。
剪枝技术演进
剪枝面临多重挑战。首先,训练大型LLMs成本高昂,训练完成后运行时成本也很高。虽然剪枝可以降低运行时成本,但如果在构建过程后期进行会损害性能;如果在构建过程过早进行,则会加剧训练成本问题。
模型训练时会根据训练数据构建语义连接图谱。这些称为参数的连接会随着更多训练数据的引入而获得或失去重要性(权重)。在训练阶段进行的“剪枝感知训练”需要执行全模型权重扫描,计算成本高昂,且需要大量全规模运行的试验负担。
另一种方法是在LLM训练完成后进行剪枝。这种方法成本较低,仅需几分钟到几小时(而训练可能需要数周),且不需要大量GPU。工程师会逐层扫描模型中的不重要权重,但这种方法不够“温和”:它会冲击模型结构,导致准确性下降。
平衡之道
我们的理念提供了第三条路径。模型完全训练后,我们逐块扫描模型,在解码块级别(而非全模型或层级)分析权重。解码块是构成LLM的大部分重复构建块。
在每个解码块内,我们输入少量数据并收集输出来校准权重,修剪不重要的权重并更新保留的权重。由于解码块规模较小,该方法仅需单个GPU即可在几分钟内完成扫描。
从技术角度看,关键在于聚焦解码块(由注意力层、多头注意力层和多层感知机等神经网络层组成)。即使具有70亿参数的LLM可能只有32个解码块。每个块足够小(约2亿参数),可轻松被单个GPU扫描。在块级别进行剪枝可节省GPU内存资源。
虽然所有剪枝过程最初都会降低性能,但我们的方法能使其恢复。每次扫描块时,我们都会平衡剪枝与性能直至优化完成,然后处理下一个块。这既保持了块级性能,也保障了整体模型质量。
技术影响与扩展应用
在解码块级别进行剪枝是“温和的”,因为剪枝影响局部化,对模型整体行为影响较小。重复每个块的剪枝过程就像厨师“常尝”调味以确保风味平衡。
该理念还有助于解决大公司LLM开发的痛点。在LLM时代之前,每个团队构建自己的模型,而现在单个LLM通过协调这些模型提供服务。通过我们的方法,从事运行时性能优化的团队可以重新获得更多GPU资源。
“轻剪常尝”理念还可应用于其他架构优化。例如,在解码块级别校准模型可将密集多层感知机转换为计算强度较低的混合专家网络(MoE)。本质上,按解码块校准能够通过用更高效的替代方案(如Kolmogorov-Arnold网络)替换通用组件,实现模型的外科手术式重新设计。