您所在的位置:首页 - 教育 - 正文教育

英伟达玩转剪枝蒸馏:把参数减半,性能同尺寸更强

admin admin 08-19 【教育】 383人已围观

摘要机器之心报道编辑:杜伟、陈陈、泽南小模型崛起了。上个月,Meta发布了Llama3.1系列模型,其中包括Meta迄今为止最大的405B模型,以及两个较小的模型,参数量分别为700亿和80亿。Llama3.1被认为是引领了开源新时代。然而,新一代的模型虽然性能强大,但部署时仍需要大量计算资源。因此,业界出现了另一种趋势,即开发小型语言模型(SLM),这种模型在许多语言任务中表现足够出色,部署起来也非常便宜。最近,英伟达研究表明,结构化权重剪枝与知识蒸馏相结合,可以从初始较大的模型中逐步获得较小的语

机器之心报道

编辑:杜伟、陈陈、泽南

小模型崛起了。

上个月,Meta发布了Llama3.1系列模型,其中包括Meta迄今为止最大的405B模型,以及两个较小的模型,参数量分别为700亿和80亿。

Llama3.1被认为是引领了开源新时代。然而,新一代的模型虽然性能强大,但部署时仍需要大量计算资源。

因此,业界出现了另一种趋势,即开发小型语言模型(SLM),这种模型在许多语言任务中表现足够出色,部署起来也非常便宜。

最近,英伟达研究表明,结构化权重剪枝与知识蒸馏相结合,可以从初始较大的模型中逐步获得较小的语言模型。

图灵奖得主、Meta首席AI科学家YannLeCun也点赞转帖了该研究。

经过剪枝和蒸馏,英伟达研究团队将Llama3.18B提炼为Llama-3.1-Minitron4B开源了出来。这是英伟达在Llama3.1开源系列中的第一个作品。

Llama-3.1-Minitron4B的表现优于类似大小的最先进的开源模型,包括Minitron4B、Phi-22.7B、Gemma22.6B和Qwen2-1.5B。

这项研究的相关论文早在上个月已经放出了。

剪枝和蒸馏

剪枝使模型变得更小、更精简,可以通过删除层(深度剪枝)或删除神经元和注意力头以及嵌入通道(宽度剪枝)来实现。剪枝通常伴随着一定程度的再训练,英伟达玩转剪枝蒸馏:把参数减半,性能同尺寸更强以恢复准确率。

模型蒸馏是一种将知识从大型复杂模型(通常称为教师模型)迁移到较小、较简单的学生模型的技术。目标是创建一个更高效的模型,该模型保留了原始较大模型的大部分预测能力,同时运行速度更快且资源消耗更少。

蒸馏方式主要包括两种:SDG微调与经典知识蒸馏,这两种蒸馏方式互补。本文主要关注经典知识蒸馏方法。

英伟达采用将剪枝与经典知识蒸馏相结合的方式来构造大模型,下图展示了单个模型的剪枝和蒸馏过程(上)以及模型剪枝和蒸馏的链条(下)。具体过程如下:

1.英伟达从15B模型开始,评估每个组件(层、神经元、头和嵌入通道)的重要性,然后对模型进行排序和剪枝,使其达到目标大小:8B模型。

2.接着使用模型蒸馏进行了轻度再训练,原始模型作为老师,剪枝后的模型作为学生。

3.训练结束后,以小模型(8B)为起点,剪枝和蒸馏为更小的4B模型。

从15B模型进行剪枝与蒸馏的过程。

需要注意的点是,在对模型剪枝之前,需要先了解模型的哪部分是重要的。英伟达提出了一种基于激活的纯重要性评估策略,该策略可以同时计算所有相关维度(深度、神经元、头和嵌入通道)的信息,使用一个包含1024个样本的小型校准数据集,并且只需要前向传播。这种方法相比依赖梯度信息并需要反向传播的策略更加简单且具有成本效益。

在剪枝过程中,你可以针对给定轴或轴组合在剪枝和重要性估计之间进行迭代交替。实证研究显示,使用单次重要性估计就足够了,迭代估计不会带来额外的好处。

利用经典知识蒸馏进行重新训练

下图2展示了蒸馏过程,其中N层学生模型(剪枝后的模型)是从M层教师模型中(原始未剪枝模型)蒸馏而来。学生模型通过最小化嵌入输出损失、logit损失以及映射到学生块S和教师块T的Transformer编码器特定损失组合来学习。

图2:蒸馏训练损失。

剪枝和蒸馏最佳实践

英伟达基于紧凑语言模型中剪枝和知识蒸馏的广泛消融研究,将自己的学习成果总结为以下几种结构化压缩最佳实践。

一是调整大小。

二是剪枝。

三是重新训练。

Llama-3.1-Minitron:将最佳实践付诸应用

Meta最近推出了功能强大的Llama3.1开源模型系列,在许多基准测试中可与闭源模型相媲美。Llama3.1的参数范围从巨大的405B到70B、8B。

凭借Nemotron蒸馏的经验,英伟达着手将Llama3.18B模型蒸馏为更小、更高效的4B模型,采取以下措施:

教师微调

为了纠正模型训练所基于的原始数据集的分布偏差,英伟达首先在他们的数据集上(94Btoken)对未剪枝的8B模型进行了微调。实验表明,如果不纠正分布偏差,教师模型在蒸馏时会为数据集提供次优指导。

Depth-only剪枝

为了从8B降到4B,英伟达剪枝了16层(50%)。他们首先通过从模型中删除每个层或连续子层组来评估它们的重要性,并观察下游任务中LM损失的增加或准确率的降低。

下图5显示了删除1、2、8或16层后验证集上的LM损失值。例如,第16层的红色图表示如果删除前16层,则出现LM损失。第17层表示如果保留第一层并删除第2至第17层,也出现LM损失。英伟达观察到:开始和结束的层是最重要的。

图5:depth-only剪枝中层的重要性。

然而,英伟达观察到,这种LM损失不一定与下游性能直接相关。

下图6显示了每个剪枝模型的Winogrande准确率,它表明最好删除第16到第31层,其中第31层是倒数第二层,剪枝模型的5-shot准确率明显高于随机准确率(0.5)。英伟达采纳了这一见解,删除了第16到第31层。

图6:当删除16层时,在Winogrande任务上的准确率。

Width-only剪枝

英伟达沿宽度轴剪枝了嵌入(隐藏)和MLP中间维,以压缩Llama3.18B。具体来说,他们使用前面描述的基于激活的策略来计算每个注意头、嵌入通道和MLP隐藏维度的重要性分数。

在重要性估计之后,英伟达选择

值得一提的是,在单样本剪枝之后,宽度剪枝的LM损失高于深度剪枝。然而,经过短暂的重新训练后,趋势发生了逆转。

准确率基准

英伟达使用以下参数对模型进行蒸馏

下表1显示了Llama-3.1-Minitron4B模型变体(宽度剪枝和深度剪枝)与原始Llama3.18B模型、其他类似大小的模型在跨多个领域的基准测试中的性能比较。总体而言,英伟达再次证实了宽度剪枝策略相较于遵循最佳实践的深度剪枝的有效性。

表1:Minitron4Bbase模型相较于类似规模base模型的准确率比较。

为了验证蒸馏后的模型是否可以成为强大的指令模型,英伟达使用NeMo-Aligner对Llama-3.1-Minitron4B模型进行了微调。

他们使用了Nemotron-4340B的训练数据,在IFEval、MT-Bench、ChatRAG-Bench和BerkeleyFunctionCallingLeaderboard(BFCL)上进行了评估,以测试指令遵循、角色扮演、RAG和函数调用功能。最后确认Llama-3.1-Minitron4B模型可以成为可靠的指令模型,其表现优于其他基线SLM。

表2:对齐Minitron4Bbase模型与类似规模的对齐模型的准确率比较。

性能基准

英伟达利用NVIDIATensorRT-LLM(一种用于优化LLM推理的开源工具包)优化了Llama3.18B和Llama-3.1-Minitron4B模型。

下两张图显示了不同模型在不同用例下以FP8和FP16精度每秒的吞吐量请求,表示为8B模型的batchsize为32的输入序列长度/输出序列长度(ISL/OSL)组合以及4B模型的batchsize为64的输入序列长度/输出序列长度(ISL/OSL)组合,这要归功于在一块英伟达H10080GBGPU上,较小的权重允许较大的batchsize。

Llama-3.1-Minitron-4B-Depth-Base变体是最快的,平均吞吐量约为Llama3.18B的2.7倍,而Llama-3.1-Minitron-4B-Width-Base变体的平均吞吐量约为Llama3.18B的1.8倍。与BF16相比,在FP8中部署还可使这三种型号的性能提高约1.3倍。

图8:组合:Llama3.18B为BS=32,Llama-3.1-Minitron4B型号为BS=64。1xH10080GBGPU。

结论

剪枝和经典知识提炼是一种非常经济高效的方法,可以逐步获得更小尺寸的LLM,与在所有领域从头开始训练相比,可实现更高的准确性。与合成数据式微调或从头开始预训练相比,这是一种更有效且数据效率更高的方法。

Llama-3.1-Minitron4B是英伟达首次尝试使用最先进的开源Llama3.1系列完成的探索。要在NVIDIANeMo中使用Llama-3.1的SDG微调,可参阅GitHub上的/sdg-law-title-generation部分。

有关更多信息,请参阅以下资源:

参考链接:

https://developer.nvidia.com/blog/how-to-prune-and-distill-llama-3-1-8b-to-an-nvidia-llama-3-1-minitron-4b-model/

Tags: 英伟达玩转剪枝蒸馏把参数减半性能同尺寸更强

icp沪ICP备2023034384号-20 icp粤公网安备 44030902003287号
取消
微信二维码
支付宝二维码

目录[+]