NVIDIA 推出高效 LLM 的剪枝和蒸馏技术

realtime news  Aug 16, 2024 19:30  UTC 11:30

1 Min Read

大型语言模型(LLM)在自然语言处理和理解中越来越重要,因为它们的有效性和多功能性。然而,它们的部署资源密集。根据NVIDIA 技术博客,NVIDIA 研究人员已经证明,通过将结构化权重剪枝与知识蒸馏相结合,可以高效地创建更小、更经济的语言模型。

剪枝和蒸馏

剪枝通过删除层(深度剪枝)或神经元、注意力头和嵌入通道(宽度剪枝)来减少模型的大小。这个过程通常随后会进行重新训练,以恢复准确性。模型蒸馏将较大、复杂模型(教师模型)的知识转移到较小、简单的模型(学生模型)中,旨在保留原始模型的大部分预测能力,同时更快且资源需求较少。

经典知识蒸馏与 SDG 微调

蒸馏可以分为两种主要方式:

  • SDG 微调:使用从较大教师模型生成的合成数据来微调较小的预训练学生模型,模仿最终的标记预测。
  • 经典知识蒸馏:学生模拟教师在训练数据集上的 logit 和其他中间状态,提供更丰富的反馈,提高训练准确性和效率。

这些方法是互补的,NVIDIA 的方法侧重于经典知识蒸馏。

剪枝和蒸馏程序

NVIDIA的程序包括:

  1. 从15B模型开始,估计每个组件的重要性,并将其剪枝到8B模型。
  2. 使用模型蒸馏进行轻度重新训练,原始模型作为教师模型,剪枝后的模型作为学生模型。
  3. 进一步剪枝和蒸馏小的8B模型为4B模型。

这种迭代方法确保一个阶段的输出模型作为下一个阶段的输入模型,优化了资源效率。

重要性分析

理解模型的重要部分对于剪枝至关重要。NVIDIA提出了一种基于激活的重要性估计策略,与基于梯度的策略相比,它具有成本效益且简单。

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

重新训练包括最小化嵌入输出损失、logit损失和变压器编码器特定的损失组合,确保较小的模型保留原始模型的大部分准确性。

剪枝和蒸馏的最佳实践

NVIDIA 的广泛研究确定了几个最佳实践:

  • 大小:首先训练最大的模型,然后迭代剪枝和蒸馏。
  • 剪枝:对于 ≤ 15B 的模型,优先使用宽度剪枝而不是深度剪枝。
  • 重新训练:仅使用蒸馏损失,在必要时结合 logit 和中间状态蒸馏。

Llama-3.1-Minitron:应用最佳实践

NVIDIA将这些实践应用于Llama 3.1 8B模型,最终得到高效的Llama-3.1-Minitron 4B模型。该模型在与类似大小的开源模型(如Phi-2 2.7B 和 Gemma2 2.6B)的竞争中表现不俗。

教师微调

在特定数据集上对未剪枝的8B模型进行微调,以纠正分布偏移,确保在蒸馏期间提供最佳指导。

仅深度剪枝和仅宽度剪枝

对于仅深度剪枝,NVIDIA从8B模型中剪枝了16层,重点是最不影响下游任务性能的层。对于仅宽度剪枝,他们减少了MLP的中介维度和隐藏大小,然后重新训练注意力头数和层数。

准确性和性能基准

表1比较了Llama-3.1-Minitron 4B变体与其他模型的性能,显示了显著的准确性和资源效率改进。性能基准表明,Llama-3.1-Minitron 4B模型相对于原始8B模型达到了约2.7倍的平均吞吐量。

结论

结合剪枝和经典知识蒸馏提供了一种具有成本效益的方法,可以创建比从头训练更小且具有更高准确性的LLM。NVIDIA的Llama-3.1-Minitron 4B模型正是这种方法的典范,提供了一种高效语言模型部署的可靠解决方案。



Read More