英伟达神操作!剪枝、蒸馏让 Llama 3.1 8B参数减半,同尺寸最强!

ff407846772b124cb0c9e1dd49269bdd.png机器之心报道

小模型崛起了。

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

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

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

1ffa3e22db1d66eb33edfcde75284574.png

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

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

Llama-3.1-Minitron 4B 的表现优于类似大小的最先进的开源模型,包括 Minitron 4B、Phi-2 2.7B、Gemma2 2.6B 和 Qwen2-1.5B。


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

4461b5ddb6724af2296b8bdad4c9a3c5.png

  • 论文链接:https://www.arxiv.org/pdf/2407.14679

  • 论文标题:Compact Language Models via Pruning and Knowledge Distillation

剪枝和蒸馏

剪枝使模型变得更小、更精简,可以通过删除层(深度剪枝)或删除神经元和注意力头以及嵌入通道(宽度剪枝)来实现。剪枝通常伴随着一定程度的再训练,以恢复准确率。

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

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

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

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

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

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

513652b5215ab57a2ab6a953d0b788f9.png

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

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

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

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

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

5ee3f86714dc6ecb35a37be0bdb1e377.png

图 2:蒸馏训练损失

剪枝和蒸馏最佳实践

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

一是调整大小。

  • 要训练一组 LLM,首先训练最大的一个,然后迭代地剪枝和蒸馏以获得较小的 LLM。

  • 如果使用多阶段训练策略来训练最大的模型,最好剪枝并对训练最后阶段获得的模型进行重新训练。

  • 对最接近目标大小的可用源模型进行剪枝。

二是剪枝。

  • 优先考虑宽度剪枝而不是深度剪枝,这对于 15B 参数规模以下的模型效果很好。

  • 使用单样本(single-shot)重要性估计,因为迭代重要性估计没有任何好处。

三是重新训练。

  • 仅使用蒸馏损失进行重新训练,而不是常规训练。

  • 当深度明显减少时,使用 logit、中间状态和嵌入蒸馏。

  • 当深度没有明显减少时,使用 logit-only 蒸馏。

结论

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

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


往期推荐

最近疯传的130个网站,收藏备用!

黑匣子被打开了?能玩的Transformer可视化解释工具!

GitHub 标星 18.3w,Python必备宝典!

新PyTorch API:几行代码轻松搞定多种注意力变体

从 Flask 切到 FastAPI 后,起飞了!


商务合作 | 交流学习 | 送书活动

添加vx:yuliang-bj(备注姓名-单位)

整理不易,点赞

### 关于 Llama-Factory 和 NVIDIA 的项目资源 #### Llama-Factory 资源概述 Llama-Factory 是一个致力于开发和优化大型语言模型(LLM)的社区项目,特别关注 Meta 发布的 LLaMA 模型系列。该项目不仅提供了预训练模型权重,还分享了一系列用于微调、评估以及部署这些模型的最佳实践指南[^3]。 对于希望基于 LLaMA 构建应用的研究人员来说,GitHub 上托管了许多有价值的资源库,其中包含了详细的文档说明和技术支持渠道。例如,在 `llama-factory` 组织页面下可以找到多个子仓库,涵盖了从基础架构搭建到高级特性实现的不方面。 #### NVIDIA 相关资源介绍 NVIDIA 提供了一套完整的 GPU 加速计算解决方案来助力 AI 开发者更高效地工作: - **硬件设备**:如前所述的基础环境中提到的 A800 80GB 显卡就是专为高性能计算而设计的产品之一;这类显卡具备强大的浮点运算能力和较大的显存容量,非常适合处理大规模深度学习任务。 - **软件栈**:包括 CUDA 工具包、cuDNN 库以及 NCCL 等组件构成了高效的并行编程框架,使得开发者能够充分利用 NVIDIA GPU 的潜力加速经网络训练过程[^2]。 此外,NVIDIA 还推出了专门针对自然语言处理领域的 SDK——NeMo Megatron-LM,它内置了多种先进的 Transformer 结构变体,并且已经预先配置好了适用于不类型 NLP 任务的工作流模板,极大地方便了研究人员快速上手实验新想法。 ```bash # 安装 NeMo Megatron-LM 所需依赖项 pip install nemo_toolkit[megatron] ``` 通过结合使用上述来自 llama-factory 社区的支持材料与 NVIDIA 提供的强大软硬件设施,研究团队可以在较短时间内建立起一套成熟稳定的 LLM 微调及推理平台。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值