原文:
towardsdatascience.com/how-to-prune-llama-3-2-and-similar-large-language-models-cf18e9a2afb6
免责声明:本文最初是用西班牙语撰写的,并使用 AI 工具进行翻译以确保准确性和一致性。您可以在这里找到原始的西班牙语版本。
随着大型语言模型为了实现更大的能力而不断增长,对更高效、更小版本的需求比以往任何时候都更加迫切。然而,在不失去核心功能的前提下减小模型大小是一项微妙的平衡行为。
量化剪枝等技术通常用于减小模型大小,而知识蒸馏或迁移学习等方法有助于保留或恢复在缩减过程中丢失的能力。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/f8a7ac7adf2a6f0e5bbd097f779fd57f.png
由作者使用 GPT 4 生成的图像。
在这些方法中,剪枝是减少模型大小最有效的策略之一。与简化数值表示的量化不同,剪枝涉及移除模型的具体部分,如神经元或整个层。但这种有效性是有代价的:剪枝正确应用起来具有挑战性。您不仅需要确定要剪枝的模型部分,还必须仔细选择要移除的元素,以最大限度地减少对模型能力的影响。
本文重点介绍了结构宽度剪枝,其中选定的神经元被移除,并展示了如何在具有门控线性单元(GLU)结构的 MLP 层上有效地应用它。通过遵循概述的步骤,您将看到剪枝如何显著减少模型大小,同时保留其生成连贯输出和在关键基准测试中表现良好的能力。
什么是剪枝以及它如何影响模型?
如我之前所解释的,剪枝涉及移除被认为对最终输出贡献最小的模型部分。通过仔细选择这些不太关键的组件,剪枝旨在创建一个具有更少参数和降低计算需求的更高效模型,同时不牺牲其核心能力。
在剪枝的主要挑战在于决定要移除模型中的哪些部分。模型的不同部分对性能的影响并不相同;每个部分都服务于不同的目的。
为了说明这一点,让我们来分析本文中使用的模型结构:LLaMA 3.2–1B。
LlamaForCausalLM((model):LlamaModel((embed_tokens):Embedding(128256,2048)(layers):ModuleList((0-15):16x LlamaDecoderLayer((self_attn):LlamaSdpaAttention((q_proj):Linear(in_features=2048,out_features=2048,bias=False)(k_proj):Linear(in_features=2048,out_features=512,bias=False)(v_proj):Linear(in_features=2048,out_features=512,bias=False)(o_proj):Linear(in_features=2048,out_features=2048,bias=False)(rotary_emb):LlamaRotaryEmbedding())(mlp):LlamaMLP((gate_proj):Linear(in_features=2048,out_features=8192,bias=False)(up_proj):Linear(in_features=2048,out_features=8192,bias=False)(down_proj):Linear(in_features=8192,out_features=2048,bias=False)(act_fn):SiLU())(input_layernorm):LlamaRMSNorm((2048,),eps=1e-05)(post_attention_layernorm):LlamaRMSNorm((2048,),eps=1e-05)))(norm):LlamaRMSNorm((2048,),eps=1e-05)(rotary_emb):LlamaRotaryEmbedding())(lm_head):Linear(in_features=2048,out_features=128256,bias=False))在检查结构时,我们可以识别出三个主要块,它们可以是剪枝的目标:嵌入、自注意力机制和 MLP 层。为了决定这些中哪一个应该是剪枝过程的焦点,理解其潜在的好处和可能对模型的影响是至关重要的。
第一步是评估这些部分中的每一个在模型中占用的空间,这给我们一个关于潜在尺寸减少的想法。
参数分布分析。
嵌入和输出层 (embed_tokens, lm_head):
每层 128256 × 2048 ≈ 262M 个参数
两层总计 524M 个参数
自注意力机制 (self_attn):
16 层,每层包含四个投影子层
每层:2048 × (2048 + 512 + 512 + 2048) ≈ 10.5M 个参数
总计:10.5 × 16 ≈ 168M 个参数
MLP 层 (mlp):
16 层具有 GLU 结构 (_gateproj, _upproj, 和 _downproj)
每层:2048 × 8192 + 2048 × 8192 + 8192 × 2048 ≈ 50M 个参数
总计:50 × 16 ≈ 805M 个参数
如我们所见,MLP 层代表了模型大小的超过 50%,使它们成为剪枝的明显候选者。然而,在做出这个决定之前,了解每个部分对模型行为的贡献是至关重要的。
影响分析。
嵌入层负责将输入转换为模型可以有效地处理的密集向量表示。剪枝嵌入层可能会导致模型理解某些单词的能力下降,或者至少会减少创建正确捕捉输入语义意义的向量的能力。如果你想要创建一个高度特定的模型,该模型只使用其输入词汇表的一个非常特定的部分,例如,用于金融或医学分析的模型,剪枝这一层可能是一个选择。
注意力机制允许模型在处理每个标记时关注输入序列中最相关的部分。它计算输入序列中每对标记之间的加权重要性分数,使模型能够捕捉上下文并关注相关信息。剪枝这一部分可能会降低模型执行需要广泛理解输入上下文的任务的能力,例如文本摘要或翻译。它还会影响生成文本的连贯性。
MLP 层伴随着注意力机制,通过一系列数据扩展和收缩来增强模型理解复杂模式的能力。剪枝这一部分可能会限制模型对未见过数据或训练期间未覆盖的任务的反应。换句话说,它降低了模型的一般化能力和对不熟悉输入提供连贯响应的能力。
一旦你决定了要针对模型的哪个部分,下一步就是确定是执行宽度剪枝,移除单个神经元,还是深度剪枝,移除整个层。
正如你所见,剪枝模型是一个相当复杂的过程,涉及到许多决策。你不仅要评估最终模型的性能,还要评估其训练能力。这些模型的设计目的是为了进行微调,通常是为了特定的任务,因此它们在执行其创建的任务时可以比基础模型更有效率和高效。
门控线性单元的特性
门控线性单元(GLU)架构在现代神经网络中很常见,包括 LLaMA、Gemma、Mistral、Qwen 和类似的大型语言模型。GLU 引入了一个逐元素的门控机制,允许模型选择性地过滤和控制信息的流动。这个架构由成对的层组成,通常是:gate_proj、up_proj 和 down_proj(如上模型结构所示),它们共同工作以扩展和收缩数据。
这个机制使得模型能够在保持效率的同时处理更复杂的模式。然而,这也意味着 GLU 结构内部的层紧密耦合,剪枝这些层需要仔细考虑。
对某一层进行的任何操作(例如,移除神经元)必须在相应的成对层中进行镜像。例如,如果从 _gateproj中移除了一个神经元,那么相同的神经元也必须从 up_proj 中移除,并且 _downproj层的大小必须相应调整。最重要的是,在计算神经元的相对重要性以决定保留哪些神经元时,你需要评估成对神经元。
打破这些层的平衡可能导致性能下降甚至模型完全失效,即使只移除了少量神经元。
剪枝 Llama 3.2 模型。
示例将使用 Llama 模型进行演示,但代码也已在 Gemma 和 QWen 模型上成功测试。
你可以在我的 GitHub 仓库上的笔记本中访问完整的代码。
GitHub – peremartra/Large-Language-Model-Notebooks-Course: 实践课程关于大型语言模型…
我在内存中的原始模型上采取的第一步是执行一个小的提示并保存结果。这使得我能够轻松、直观、快速地检查通过剪枝过程生成的模型是否连贯,或者相反,是否失去了生成可理解文本的能力。
让我向你保证,在第一次尝试中,由于没有尊重模型的 GLU 结构,生成的文本毫无疑问地表明剪枝过程存在根本性的缺陷。
原始提示是:“巴黎是…的首都。”让我们看看原始模型的响应,并将其与我的第一次、失败的剪枝尝试返回的响应进行比较。
基础模型:
“巴黎是法国的首都,也是世界上最受游客欢迎的城市之一。这座城市以艺术、文化、时尚和美食而闻名。这座城市拥有丰富的历史,并拥有许多著名的地标,包括埃菲尔铁塔等。”
仅剪枝 20% 的错误模型:
“巴黎是法国的首都。这是主要区域。这是法国的城市。这是法国的主要区域。这是法国的城市。这是法国的主要区域。这是法国的城市。这是法国的主要区域。这是法国的城市。这是法国的主要区域。这是法国的城市。这是法国的主要区域。这是法国的城市。这是法国的主要区域。这是法国的城市。”
很明显,第一次尝试中有些东西没有工作。这可能看起来微不足道,但这样的经验检查可以为你节省几个小时。
实现细节
让我们先看看负责计算神经元重要性的函数,它最终将决定哪些神经元保留在模型中,哪些神经元被移除。
defcompute_neuron_pair_importance(gate_weight,up_weight):""" compute neuron pair importance scores (Maximum Absolute Weight) Args: - gate_weight: Weight matrix from the gate_proj layer. - up_weight: Weight matrix from the up_weight layer. Returns: - importance_scores: Importance scores for each neuron pair. """gate_max_abs=torch.max(gate_weight,dim=1).values+torch.abs(torch.min(gate_weight,dim=1).values)up_max_abs=torch.max(up_weight,dim=1).values+torch.abs(torch.min(up_weight,dim=1).values)importance_scores=gate_max_abs+up_max_absreturnimportance_scores该函数接收一个 _gateproj层和一个 _upproj层的权重,正如我之前解释的,它们是成对工作的。因此,必须联合计算神经元的相对重要性。
计算非常简单:它计算每个神经元的权重的绝对值。考虑正负值,因为在理论上,具有最极端值的神经元通过显著改变通过它们的值对模型输出的影响而具有更大的影响。
在这里,我必须感谢 MariusZ Kurman 对将最小值纳入计算的贡献。虽然没有它们方法也能正确工作,但它们的加入提高了结果。
对于每个层,重要性是单独计算的,但函数返回的是综合值。
defprune_neuron_pairs(mlp,prune_percent):""" Reduces the dimensions of the **gate_proj**,**up_proj**, **down_proj** layers removing the least important neurons. Args: - mlp: Layers to prune. - prune_percent: Percentage of neurons to prune. Returns: - new_gate_proj, new_up_proj, new_down_proj: New pruned layers. - k: New intermediate size. """# Extract weights from MLP layersgate_weight=mlp.gate_proj.weight.data.float()up_weight=mlp.up_proj.weight.data.float()# Compute importance scoresimportance_scores=compute_neuron_pair_importance(gate_weight,up_weight)original_intermediate_size=gate_weight.size(0)# Calculate neurons to keepnum_neuron_pairs_to_prune=min(int(prune_percent*original_intermediate_size),original_intermediate_size-1)k=original_intermediate_size-num_neuron_pairs_to_prune# Validation checkifk<=0:raiseValueError(f"Invalid number of neuron pairs to keep:{k}")# Select neurons to keep_,indices_to_keep=torch.topk(importance_scores,k,largest=True,sorted=True)indices_to_keep=indices_to_keep.sort().values# Create and populate new layersnew_gate_proj=nn.Linear(mlp.gate_proj.in_features,k,bias=False).to(device)new_up_proj=nn.Linear(mlp.up_proj.in_features,k,bias=False).to(device)new_down_proj=nn.Linear(k,mlp.down_proj.out_features,bias=False).to(device)# Copy selected weightsnew_gate_proj.weight.data=mlp.gate_proj.weight.data[indices_to_keep,:]new_up_proj.weight.data=mlp.up_proj.weight.data[indices_to_keep,:]new_down_proj.weight.data=mlp.down_proj.weight.data[:,indices_to_keep]returnnew_gate_proj,new_up_proj,new_down_proj,k这个函数在保留最重要的神经元的同时创建新的、更小的层。这个过程包括:
- 提取当前权重:
# Extract weights from MLP layersgate_weight=mlp.gate_proj.weight.data.float()up_weight=mlp.up_proj.weight.data.float()- 计算神经元对的重要性分数:
# Compute importance scoresimportance_scores=compute_neuron_pair_importance(gate_weight,up_weight)original_intermediate_size=gate_weight.size(0)获得一个张量,其中包含为每个神经元计算的重要性分数。这些分数反映了每个神经元对最终输出的贡献,指示哪些应该被保留。
- 确定要保留多少神经元:
# Calculate neurons to keepnum_neuron_pairs_to_prune=min(int(prune_percent*original_intermediate_size),original_intermediate_size-1)k=original_intermediate_size-num_neuron_pairs_to_prune要保留的神经元总数是通过提供的剪枝百分比参数和原始层的大小来计算的。
- 选择最重要的神经元:
# Select neurons to keep_,indices_to_keep=torch.topk(importance_scores,k,largest=True,sorted=True)indices_to_keep=indices_to_keep.sort().values使用 Torch 获取具有最高重要性分数的神经元,同时将它们从最重要到最不重要进行排序。由于 torch 返回的数据是降序的,因此使用排序方法将它们重新排列为升序,这是我们需要的。
- 创建新的、更小的层:
# Create and populate new layersnew_gate_proj=nn.Linear(mlp.gate_proj.in_features,k,bias=False).to(device)new_up_proj=nn.Linear(mlp.up_proj.in_features,k,bias=False).to(device)new_down_proj=nn.Linear(k,mlp.down_proj.out_features,bias=False).to(device)创建了三个新层,其尺寸根据所选索引进行调整。在 _new_gateproj和 _new_upproj中,保留了输入维度,而输出维度被减小。相反,在 _new_downproj中,调整了输入维度,而输出维度保持不变。
- 将选定的权重复制到新层:
#copy weights to the new layers.new_gate_proj.weight.data=mlp.gate_proj.weight.data[indices_to_keep,:]new_up_proj.weight.data=mlp.up_proj.weight.data[indices_to_keep,:]new_down_proj.weight.data=mlp.down_proj.weight.data[:,indices_to_keep]相关权重从原始层转移到新层,确保只保留对应于所选神经元的权重。
现在,让我们看看负责遍历所有层并构建修改后模型的函数。
defupdate_model(model,prune_percent):""" Modifies each MLP layer in the model to retain only the most important neurons. Args: - model: Model to prune. - prune_percent: Percentage of neurons to prune. Returns: - model: New pruned model. """new_intermediate_size=Noneforidx,layerinenumerate(model.model.layers):mlp=layer.mlp new_gate_proj,new_up_proj,new_down_proj,new_size=prune_neuron_pairs(mlp,prune_percent)mlp.gate_proj=new_gate_proj mlp.up_proj=new_up_proj mlp.down_proj=new_down_projifnew_intermediate_sizeisNone:new_intermediate_size=new_size model.config.intermediate_size=new_intermediate_sizereturnmodel此函数遍历模型的每一层,应用剪枝过程并更新模型的配置以反映新的架构。
如果配置文件没有更新,那么在保存后,无论是在 Hugging Face 还是本地,模型都无法使用。许多库,如 Hugging Face 的 Transformers,依赖于model.config来解释模型的架构。如果配置与实际结构不匹配,那么通过这些库进行的微调或推理操作可能会失败。
结果分析。
通过这段代码,我创建了几个模型,这些模型可在 Hugging Face Hub 上找到。
这些包括:
从 Llama-3.2–1b 衍生出的三个模型,分别剪除了 MLP 层中 20%、40% 和 60% 的神经元。
一个基于 Gemma-2–2B 的模型,通过 40% 剪枝。
您可以下载这些模型,除了使用它们之外,还可以研究它们的架构以及与它们所基于的原模型相比发生了哪些变化。
让我们分析在 Llama3.2–1b 模型上应用 20% 剪枝后架构的变化。
LlamaForCausalLM((model):LlamaModel((embed_tokens):Embedding(128256,2048)(layers):ModuleList((0-15):16x LlamaDecoderLayer((self_attn):LlamaSdpaAttention((q_proj):Linear(in_features=2048,out_features=2048,bias=False)(k_proj):Linear(in_features=2048,out_features=512,bias=False)(v_proj):Linear(in_features=2048,out_features=512,bias=False)(o_proj):Linear(in_features=2048,out_features=2048,bias=False)(rotary_emb):LlamaRotaryEmbedding())(mlp):LlamaMLP((gate_proj):Linear(in_features=2048,out_features=6554,bias=False)(up_proj):Linear(in_features=2048,out_features=6554,bias=False)(down_proj):Linear(in_features=6554,out_features=2048,bias=False)(act_fn):SiLU())(input_layernorm):LlamaRMSNorm((2048,),eps=1e-05)(post_attention_layernorm):LlamaRMSNorm((2048,),eps=1e-05)))(norm):LlamaRMSNorm((2048,),eps=1e-05)(rotary_emb):LlamaRotaryEmbedding())(lm_head):Linear(in_features=2048,out_features=128256,bias=False))模型的结构保持不变,除了 MLP 块中中间层的大小。正如您所看到的,_gateproj和 _upproj层的特征数已从 8192 减少到 6554,而 _downproj层也经历了相同的变化,但是在其输入特征上。
这个变化与代码所做的一致:修改这些层的同时保留对模型性能最关键的神经元。如果我们从 8192 中移除 20%,我们得到 6553.6,这证实了剪枝的神经元百分比是正确的。
经验性提示测试。
现在,让我们看看剪枝模型在测试提示下的表现:
巴黎是法国的首都。它也是世界上最美丽的城市之一。在巴黎有太多可看可做的事情,一天之内不可能全部看完。然而,有一些事情你
响应与原始模型的响应不完全相同,但保持了连贯性。这表明模型保留了其大部分能力,更重要的是,它有可能通过知识蒸馏或微调恢复任何损失。
EleutherAI / lm-evaluation。
除了这个经验性检查之外,我还使用一些最常用的基准对模型进行了评估。让我们分析不同程度的剪枝如何影响模型的表现。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/9ab8a5873c83b58fe4e3daba6483c4ec.png
图片由作者提供。不同指标上的模型性能
如我们所见,剪枝的影响在一定程度上是不对称的。BoolQ 测试评估的任务没有经历显著的退化,对于一个失去了 MLP 层 40% 神经元的模型,下降幅度仅为约 2%。
相比之下,对 Lambada 测试的影响是显著的,准确率下降了超过 50%。
这表明模型保留了大部分的理解能力,但在需要更多开放式生成的测试中表现不佳。
BoolQ 简单地向模型展示一段文本和一个需要用 Yes/No 回答的问题。这是一个专注于测量模型理解输入文本中关系的能力的测试。
另一方面,Lambada 要求模型猜测段落中的最后一个单词,这是一个复杂的任务,最后一个单词测试了模型在复杂语言建模方面的能力。
Hugging Face Open LLM 领先板。
在 Hugging Face Open LLM 领先板上剪枝到 20% 的模型结果可能更加令人惊讶,因为它优于其基础模型以及广泛使用的 TinyLlama-1.1B-v1.1。
在这张图中,我们可以看到两个模型的测试结果。
https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/b767432db04b618d6e2c7c1e9dddc5af.png
图片由作者使用 GPT 生成。
从这张图中,我们可以得出以下结论:剪枝模型在平均性能上优于基础模型(4.86 vs. 4.03)。这表明剪枝过程在关键领域有效地保留了或增强了性能,同时减少了冗余。
通过研究这些结果,我们可以识别出剪枝模型的优缺点。
优势:
IFEval:显著的改进(19.94 vs. 14.78)表明剪枝要么减少了过拟合,要么提高了模型高效提取信息的能力。
MUSR:更好的性能(4.39 vs. 2.56)表明剪枝模型在处理需要长上下文推理或叙事理解的任务上表现更好,这可能是由于聚焦的权重。
劣势:
BBH:在不确定性下的推理能力下降(3.19 vs. 4.37)可能表明剪枝减少了模型处理模糊或多解释场景的能力。
MMLU-PRO:专业领域特定任务下降(1.36 vs. 2.26)可能是由于移除了保留特定领域详细知识所需的关键权重。
能源效率:剪枝模型在能源效率上略高(0.4 kg vs. 0.42 kg CO₂),这与减少计算开销同时保持竞争性能的目标相一致。
对模型在不同排名上的性能进行更全面的研究是必要的,但这些结果表明,我们有一个有希望的模型,通过适当的知识蒸馏或微调可以显著改进。最重要的是,这些结果与对 MLP 层进行的剪枝过程相一致。
结论。
模型的剪枝过程是成功的。这种处理 GLU 层的方法允许我们在保留模型大部分能力的同时进行剪枝,从而显著减小其大小和资源消耗。
需要注意的是,测试结果是在剪枝模型在进行任何能力恢复过程(如知识蒸馏或微调)之前获得的,这些过程通常是针对经过剪枝的模型进行的。
未来工作。
有许多剪枝技术值得探索。可能最直接的是深度剪枝,这涉及到移除对模型性能贡献最小的层。
另一个重要的研究领域是将这些剪枝模型提交给知识蒸馏过程,并评估它们是否保留了学习新任务的能力。这可能会使它们的性能更接近基线模型,特别是在剪枝模型显示出最大损失的基准测试中。
开发更轻量、更高效的模型仍然是一个有吸引力的领域,尤其是对于寻求在不需大量基础设施的情况下部署 LLM 能力的公司。这项工作为将这些强大的模型变得更加可访问和可部署提供了基础。
**本文是关于大型语言模型的完整课程的一部分,可在 GitHub 上找到。为了了解新文章的更新,请考虑关注该存储库或给它加星。**这样,您将在新内容添加时收到通知。
我是 Apress 出版社出版的书籍《大型语言模型项目:应用和实现大型语言模型的策略》的作者。
我经常撰写关于生成式 AI、深度学习和 TensorFlow 的文章。**请考虑在 Medium 上关注我,以获取新文章的更新。**当然,您也可以在LinkedIn上与我建立联系。