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

大型语言模型（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模型正是这种方法的典范，提供了一种高效语言模型部署的可靠解决方案。