LLM Pruning and Distillation in Practice: The Minitron Approach

📄 arXiv: 2408.11796v4 📥 PDF

作者: Sharath Turuvekere Sreenivas, Saurav Muralidharan, Raviraj Joshi, Marcin Chochowski, Ameya Sunil Mahabaleshwarkar, Gerald Shen, Jiaqi Zeng, Zijia Chen, Yoshi Suhara, Shizhe Diao, Chenhan Yu, Wei-Chun Chen, Hayley Ross, Oluwatobi Olabiyi, Ashwath Aithal, Oleksii Kuchaiev, Daniel Korzekwa, Pavlo Molchanov, Mostofa Patwary, Mohammad Shoeybi, Jan Kautz, Bryan Catanzaro

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-08-21 (更新: 2024-12-09)

备注: v4: Update author order

💡 一句话要点

Minitron：通过剪枝和蒸馏将Llama 3.1和Mistral NeMo模型压缩至更小规模。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 大语言模型压缩 剪枝 蒸馏 模型优化 知识迁移

📋 核心要点

  1. 现有大语言模型参数量巨大，部署和推理成本高昂，需要有效的压缩方法。
  2. 论文提出Minitron方法，通过剪枝和蒸馏技术，在不显著降低性能的前提下，压缩模型规模。
  3. 实验表明，该方法成功将Llama 3.1 8B压缩到4B，Mistral NeMo 12B压缩到8B，并在基准测试中表现良好。

📝 摘要（中文）

本文全面报告了使用剪枝和蒸馏技术将Llama 3.1 8B和Mistral NeMo 12B模型分别压缩到4B和8B参数的实践。研究探索了两种不同的剪枝策略：（1）深度剪枝；（2）联合隐藏层/注意力/MLP（宽度）剪枝，并在LM Evaluation Harness的常见基准上评估了结果。然后，使用NeMo Aligner对模型进行对齐，并在指令调优版本中进行测试。这种方法从Llama 3.1 8B生成了一个引人注目的4B模型，并从Mistral NeMo 12B生成了一个最先进的Mistral-NeMo-Minitron-8B（简称MN-Minitron-8B）模型。研究发现，在无法访问原始数据的情况下，在蒸馏数据集上对教师模型进行轻微微调是有益的。论文在Hugging Face上以宽松的许可开源了基础模型权重。

🔬 方法详解

问题定义：现有的大型语言模型（LLM）虽然性能强大，但其巨大的参数量带来了高昂的计算和存储成本，限制了它们在资源受限环境中的部署和应用。现有的模型压缩方法，如量化、剪枝和蒸馏，各有优缺点，如何在保持模型性能的同时，有效地压缩模型规模是一个重要的研究问题。

核心思路：Minitron方法的核心思路是结合剪枝和蒸馏两种技术，首先通过剪枝减少模型的参数量，然后利用蒸馏技术将大型教师模型的知识迁移到小型学生模型，从而在压缩模型规模的同时，尽可能地保留原始模型的性能。论文还探索了不同的剪枝策略，并研究了教师模型微调对蒸馏效果的影响。

技术框架：Minitron方法的整体框架包括以下几个主要阶段：1) 教师模型选择：选择Llama 3.1 8B和Mistral NeMo 12B作为教师模型。2) 模型剪枝：采用深度剪枝和联合宽度剪枝两种策略，减少模型参数量。3) 蒸馏训练：使用NeMo Aligner对剪枝后的模型进行对齐，并进行蒸馏训练，将知识从教师模型迁移到学生模型。4) 指令调优：对蒸馏后的模型进行指令调优，提高模型在特定任务上的性能。5) 模型评估：在LM Evaluation Harness的常见基准上评估模型的性能。

关键创新：Minitron方法的关键创新在于：1) 探索了深度剪枝和联合宽度剪枝两种不同的剪枝策略，并比较了它们的效果。2) 研究了在无法访问原始数据的情况下，对教师模型进行微调对蒸馏效果的影响，发现轻微微调可以提高蒸馏效果。3) 提出了一个完整的模型压缩流程，包括剪枝、蒸馏和指令调优，并开源了压缩后的模型权重。

关键设计：在剪枝方面，论文探索了深度剪枝（移除整个Transformer层）和联合宽度剪枝（同时剪枝隐藏层、注意力机制和MLP层）两种策略。在蒸馏方面，使用了NeMo Aligner进行模型对齐，并采用了标准的蒸馏损失函数。论文还发现，在蒸馏之前，对教师模型在蒸馏数据集上进行轻微微调可以提高蒸馏效果。具体的微调策略和超参数设置在论文中有详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Minitron方法成功将Llama 3.1 8B模型压缩到4B，同时保持了良好的性能。压缩后的Mistral-NeMo-Minitron-8B模型在多个基准测试中取得了state-of-the-art的结果。论文还发现，在蒸馏之前对教师模型进行轻微微调可以显著提高蒸馏效果。

🎯 应用场景

该研究成果可广泛应用于自然语言处理领域，尤其是在资源受限的场景下，如移动设备、边缘计算等。压缩后的模型可以降低部署成本，提高推理速度，从而加速大语言模型在实际应用中的落地。此外，该方法还可以用于构建更高效的对话系统、智能助手等。

📄 摘要（原文）

We present a comprehensive report on compressing the Llama 3.1 8B and Mistral NeMo 12B models to 4B and 8B parameters, respectively, using pruning and distillation. We explore two distinct pruning strategies: (1) depth pruning and (2) joint hidden/attention/MLP (width) pruning, and evaluate the results on common benchmarks from the LM Evaluation Harness. The models are then aligned with NeMo Aligner and tested in instruct-tuned versions. This approach produces a compelling 4B model from Llama 3.1 8B and a state-of-the-art Mistral-NeMo-Minitron-8B (MN-Minitron-8B for brevity) model from Mistral NeMo 12B. We found that with no access to the original data, it is beneficial to slightly fine-tune teacher models on the distillation dataset. We open-source our base model weights on Hugging Face with a permissive license.