“我”变小了但更强了！英伟达发布最新大语言模型压缩技术，无损性能且提升数倍！

1. 摘要

英伟达研究团队提出了一份全面报告，详细介绍了如何使用剪枝和蒸馏技术将Llama 3.1 8B和Mistral NeMo 12B模型分别压缩至4B和8B参数。他们探索了两种不同的剪枝策略：深度剪枝和联合隐藏/注意力/MLP（宽度）剪枝。

研究人员使用LM Evaluation Harness中的常用基准测试对压缩结果进行了评估。这些模型通过NeMo Aligner进行对齐，并在指令调优版本中进行了测试。

这种方法成功地从Llama 3.1 8B产生了一个令人信服的4B模型，并从Mistral NeMo 12B产生了一个最先进的Mistral-NeMo-Minitron-8B（简称MN-Minitron-8B）模型。

研究团队发现，在没有访问原始数据的情况下，对教师模型在蒸馏数据集上进行轻微的微调是有益的。这一发现为在有限资源条件下优化大型语言模型提供了新的思路。

为了促进社区合作和进一步研究，研究团队在Hugging Face平台上开源了他们的基础模型权重，并采用了宽松的许可证。这一举措将有助于加速相关研究的进展，并推动大型语言模型在更广泛的应用场景中的使用。

可在Hugging Face上获取的模型包括：

●Mistral-NeMo-Minitron-8B-Base

●Llama-3.1-Minitron-4B-Width-Base

●Llama-3.1-Minitron-4B-Depth-Base

2. 引言

在人工智能领域，大语言模型（LLMs）的发展日新月异。同时，这些模型的规模和计算需求也在不断增加，为实际应用带来了挑战。最新的研究成果为我们带来了一个好消息：通过创新的压缩技术，论文中方法可以在不损失性能的情况下，显著减小模型规模并提高运行速度。让我们一起来探索这项突破性的研究。

■2.1 传统方法的局限性

传统上，LLM提供商通常会从头开始训练一整个系列的模型，每个模型具有不同的规模（参数数量）。例如，Llama 3.1系列包括8B、70B和405B等不同规模的模型。这种做法的目的是为了满足用户在不同部署规模、大小和计算预算下的需求。

然而，这种方法存在一个重大缺陷：从头训练多个数十亿参数的模型极其耗时、耗数据和耗资源。在当前对AI效率和可持续性日益关注的背景下，这种做法显然不够理想。

■2.2 创新的压缩策略：Minitron方法

英伟达在7月份发布了一篇压缩大语言模型的方法，其展示了一种极具前景的方法，将权重剪枝（weight pruning）与知识蒸馏（knowledge distillation）相结合，大幅降低了训练LLM模型系列的成本。这种被称为Minitron压缩策略的方法为高效生产LLM模型系列开辟了一条新路。

◆剪枝技术

在本研究中，团队专注于结构化剪枝，即一次性从模型权重中移除非零元素的块（或通道）。图2展示了LLMs的剪枝过程，包括重要性估计、重要性排序和模型修剪等步骤。

值得注意的是，研究团队考虑了多种指标来评估层的重要性，最终选择使用Winogrande基准来评估层的重要性。

◆重训练与知识蒸馏

剪枝后，需要进行重训练以恢复准确性。知识蒸馏（KD）是一种将知识从较大模型（教师）转移到较小模型（学生）的技术。图3对蒸馏过程进行了简要说明，团队使用教师和学生logits上的前向KL散度损失进行蒸馏。

■2.3 关键改进：教师模型校正

在遵循Compact Language Models的基础上，研究团队做出了一个关键的改进：教师模型校正（teacher correction）。由于无法访问原始训练数据，团队在剪枝和蒸馏之前，先在自己的数据集上对教师模型进行微调。

这一步骤的重要性通过实验得到了证实：如果省略教师模型校正，会导致数据分布不匹配，从而对蒸馏过程产生负面影响。这一发现为未来的模型压缩研究提供了有价值的洞察。

3. 训练细节

研究团队在训练过程中采用了一系列精心设计的策略，以确保压缩模型的性能。以下是主要的训练细节：

■3.1 预训练模型

研究团队选择了两个预训练模型作为起点：Llama 3.1 8B和Mistral NeMo 12B。这些模型都是在专有数据集上预训练的，研究团队使用了Hugging Face上公开可用的Base模型版本。值得注意的是，Llama 3.1 8B模型在15T个token上进行了预训练，为后续的压缩实验奠定了坚实的基础。

■3.2 数据集

所有剪枝和蒸馏实验都使用了Nemotron-4精选的持续训练（CT）数据集。这个数据集的选择对于确保压缩模型能够保持原有模型的性能至关重要。

■3.3 剪枝策略

研究团队采用了基于Minitron论文的简化剪枝方法，针对不同的目标压缩模型制定了特定的剪枝策略。对于Llama-3.1-Minitron-4B-Width模型，他们减少了隐藏维度和MLP隐藏维度，同时保持注意力头和深度不变。Llama-3.1-Minitron-4B-Depth模型则将深度减半，保持其他维度不变。MN-Minitron-8B模型也进行了类似的维度调整，以达到预期的压缩效果。

■3.4 知识蒸馏

知识蒸馏过程包括两个关键步骤。首先，研究团队在约127B个token上微调教师模型，这一步骤对于处理原始数据集与蒸馏数据集之间的分布差异至关重要。如图4所示，如果在蒸馏过程中无法获得原始数据集，这种校正就变得尤为重要。这一发现强调了在知识蒸馏过程中数据一致性的重要性，同时也为那些无法访问原始预训练数据的研究者提供了一种可行的解决方案。

随后，研究团队进行了重训练。遵循Minitron工作中的经验，他们选择了仅logit蒸馏的方法，最小化教师和学生概率之间的前向KL散度损失，完全忽略LM交叉熵损失。在这个过程中，未剪枝的模型作为教师，剪枝后的模型作为学生。研究团队在蒸馏过程中使用了表4中列出的超参数。

表4展示了在基于蒸馏的重训练过程中使用的超参数。对于Llama-3.1-Minitron-4B和MN-Minitron-8B两个模型，研究团队分别采用了不同的设置。两个模型的峰值学习率都设置为1e-4，但最小学习率有所不同。Llama-3.1-Minitron-4B的预热步骤为40步，全局批量大小为1152，总共训练了94B个token。而MN-Minitron-8B的预热步骤为60步，全局批量大小为768，总共训练了380B个token。两个模型都使用了余弦学习率衰减计划，上下文长度均为8192。

整个训练过程使用了32个NVIDIA DGX H100节点，充分利用了高性能计算资源，以确保蒸馏过程的效率和效果。这种精心设计的知识蒸馏策略不仅解决了数据分布差异的问题，还确保了压缩模型能够最大程度地保留原始大模型的性能。通过这种方法，研究团队成功地在模型大小和性能之间取得了令人印象深刻的平衡。

■3.5 指令微调

为评估压缩模型的指令跟随能力，研究团队使用NeMo-Aligner对Llama-3.1-Minitron 4B模型进行了监督微调（SFT）。评估涵盖了多个方面，包括指令跟随和角色扮演能力（使用IFEval和MT-Bench）、RAG问答能力（使用ChatRAG-Bench）以及函数调用能力（使用BFCL）。这种全面的评估方法确保了压缩模型在各种实际应用场景中的表现。

这些详细的训练策略不仅为我们深入理解大语言模型压缩过程提供了宝贵的见解，同时也展示了研究团队在优化模型性能方面的深思熟虑。通过这些精心设计的步骤，研究团队成功地在保持模型性能的同时，显著减少了模型的大小和计算需求。

4. 评估结果：压缩模型的实际表现

研究团队对压缩后的模型进行了全面的评估，以验证其在各种任务中的性能。这些评估不仅展示了模型的能力，也为我们提供了宝贵的洞察。

■4.1 基准测试

研究团队选择了一系列具有代表性的下游任务来评估模型性能，包括语言理解、代码生成、常识推理问答和摘要生成等。对于指令调优后的模型，还使用了专门的基准测试来评估其指令遵循、角色扮演和函数调用能力。

■4.2 基础模型评估结果

表1展示了基础模型的评估结果。研究发现：

●MN-Minitron-8B和Llama-3.1-Minitron 4B模型在各项任务中表现出色，且仅使用了原模型极小部分的训练数据。

●宽度剪枝变体的性能普遍优于深度剪枝变体。

这些结果充分证明了研究团队方法的优势：在显著提高训练效率的同时，还能达到甚至超越原模型的性能。

■4.3 指令模型评估结果

表2展示了指令调优后的Llama-3.1-Minitron 4B变体的性能。与其他同等规模的基准模型相比，Minitron模型在多个任务上表现强劲，特别是在基于检索的问答和函数调用任务上达到了最先进的性能。

■4.4 运行时性能分析

研究团队使用NVIDIA TensorRT-LLM对模型进行了优化，图10展示了不同模型在FP8精度下的吞吐量。

关键发现：

●深度剪枝变体（Llama-3.1-Minitron-4B-Depth）是最快的，平均吞吐量比原模型提高了2.7倍。

●宽度剪枝变体也显著提升了性能，平均吞吐量提高了1.8倍。

●FP8精度相比BF16提供了1.4倍的性能提升。

这些评估结果证明了压缩模型不仅保持了强大的性能，还大幅提升了运行效率。通过显著减小模型规模同时保持或提升性能，这项研究为大语言模型在资源受限环境中的广泛应用铺平了道路。

5. 研究洞察

本节将深入探讨研究团队通过一系列消融实验所获得的重要发现和观察结果。这些洞察不仅深化了我们对大语言模型压缩过程的理解，也为未来的研究和应用提供了宝贵的指导。

■5.1 宽度与深度剪枝的比较

图5展示了Llama-3.1-Minitron-4B模型在宽度和深度剪枝下的训练曲线。研究发现，尽管两种变体的参数数量相同，但宽度剪枝的初始损失更小，并且在整个训练过程中始终优于深度剪枝模型。这一发现为模型压缩策略的选择提供了重要参考。

具体到Llama 3.1 8B到Llama-3.1-Minitron 4B的压缩过程中：

●宽度剪枝在MMLU基准测试上表现更好，达到60.5%的准确率，而深度剪枝为58.7%。

●然而，推理能力受到更显著的影响，GSM8K准确率在宽度剪枝中为41.24%，深度剪枝中仅为16.8%。

●深度剪枝在吞吐量方面表现出色，相比Llama-3.1 8B提高了2.7倍，而宽度剪枝提供了1.7倍的加速。

这些结果突显了不同剪枝策略在性能、准确率和效率方面的权衡。

■5.2 枝和蒸馏的协同效应

研究团队比较了四种不同的方法：随机权重初始化、随机剪枝、提议的剪枝方法配合传统训练，以及提议的剪枝方法配合蒸馏训练。图6清晰地展示了这些方法的效果。结果表明：

●剪枝相比随机初始化提供了更好的起点；

●基于蒸馏的训练优于传统训练方法，同时显著减少了所需的训练token数量（在本研究中减少了高达50倍）。

这些发现强调了剪枝和蒸馏结合的重要性，为高效模型压缩提供了新的思路。值得注意的是，研究团队发现只需要380B个token就可以在剪枝和蒸馏后达到最先进的准确率，这一发现与Minitron论文的观察结果一致，突显了该方法在训练效率方面的显著优势。

■5.3 教师模型校正的重要性

研究比较了两种教师校正方法：

●对校正后的教师进行剪枝和蒸馏；

●对原始教师进行剪枝，然后从持续校正的教师进行蒸馏。

图7的结果表明，教师校正对剪枝的最优性没有影响，但从校正后的教师进行蒸馏至关重要。

具体而言，在新的、未见过的数据集上进行蒸馏时，教师模型的校正可以将语言模型验证损失降低超过6%。这种校正不仅不会影响剪枝的最优性，甚至可以与蒸馏过程并行进行，进一步提高效率。

■5.4 深度剪枝策略的洞察

在研究连续移除层块如何影响LM验证损失时，图8显示模型的开始和结束层最为重要。然而，图9展示了基于每层重要性选择层进行剪枝与连续移除层之间的性能差异。这一结果表明，在评估下游任务性能时，连续层剪枝可能更为有效。

研究发现，连续删除模型的层比使用基于重要性的非连续剪枝更有效。这一发现为深度剪枝策略提供了重要指导。

6. 总结与展望

本研究通过创新的压缩技术，成功地将大型语言模型缩小到更加实用的规模，同时提升了模型性能。研究团队将Llama 3.1 8B压缩至4B，Mistral NeMo 12B压缩至8B，不仅在多项基准测试中表现出色，部分任务甚至超越了原始大模型。压缩后的模型显著提高了运行速度，深度剪枝变体平均加速2.7倍，宽度剪枝变体加速1.8倍。研究还揭示了多个关键洞察：宽度剪枝在准确性上略胜一筹，而深度剪枝在运行速度上更具优势；教师模型校正对提升压缩模型性能至关重要；剪枝与蒸馏的结合使用能够在保持性能的同时大幅提高训练效率。

这项突破性研究为AI领域带来了广泛的应用前景和未来展望。压缩模型为在资源受限环境下部署大语言模型开辟了新的可能性，如移动设备和边缘计算。它为进一步优化压缩技术指明了方向，如探索更高效的剪枝策略和蒸馏方法。同时，这种方法为高效开发不同规模的模型系列提供了新思路，有助于满足多样化的应用需求。通过减少模型规模和提高效率，研究也为构建更加环保和可持续的AI系统做出了贡献。研究团队开源模型权重的举措，进一步促进了社区协作和深入研究。随着压缩技术的不断进步，我们可以期待看到更多小型但强大的AI模型在各个领域发挥作用，推动人工智能技术向着更加高效、普及和可持续的方向发展。

参考：
[1]Compact language models via pruning and knowledge distillation
[2]A framework for few-shot language model evaluation
[3]Distilling the Knowledge in a Neural Network
[4]Winogrande: An adversarial winograd schema challenge at scale
[5]On information and sufficiency. The annals of mathematical statistics, 22(1), 79-86.
©️【深蓝AI】

Ref:

LLM Pruning and Distillation in Practice: The Minitron Approach

编译｜Jeffery

审核｜Los