英伟达发表的Minitron方法：实际应用中的LLM剪枝与蒸馏

小天才学习机打游戏

于 2024-09-27 20:04:26 发布

阅读量403

点赞数 7

文章标签：剪枝人工智能深度学习 pdf 开源 javascript 算法

本文链接：https://blog.csdn.net/m0_59164520/article/details/142601374

版权

一、结论写在前面

论文标题：LLM Pruning and Distillation in Practice: The Minitron Approach

论文链接：https://arxiv.org/pdf/2408.11796

Mistral-NeMo-Minitron-8B-Base模型链接：https://huggingface.co/nvidia/Mistral-NeMo-Minitron-8B-Base

nvidia/Llama-3.1-Minitron-4B-Width-Base模型链接：https://huggingface.co/nvidia/Llama-3.1-Minitron-4B-Width-Base

nvidia/Llama-3.1-Minitron-4B-Depth-Base模型链接：https://huggingface.co/nvidia/Llama-3.1-Minitron-4B-Depth-Base

论文提供了一份关于使用剪枝和蒸馏技术将Llama 3.1 8B和Mistral NeMo 12B模型分别压缩至4B和8B参数的全面报告。论文探讨了两种不同的剪枝策略：（1）深度剪枝和（2）联合隐藏层/注意力/MLP（宽度）剪枝，并在LM评估工具套件的常见基准上评估了结果。

随后，这些模型通过NeMo Aligner进行了对齐，并在指令调优版本中进行了测试。这种方法从Llama 3.1 8B模型中产生了一个引人注目的4B模型，以及从Mistral NeMo 12B模型中产生了一个最新的Mistral-NeMo-Minitron-8B（简称MN-Minitron-8B）模型。

论文的压缩策略产生了一个最先进的8亿参数模型（MN-Minitron-8B）在常见的语言建模基准测试中，全面超越所有同等规模的模型。论文的Llama-3.1-Minitron-4B模型（深度和宽度剪枝变体）与教师模型Llama 3.1 8 B和上一代Minitron-4B模型相比，也表现出强大的准确性；

在这两个变体中，宽度剪枝变体优于深度剪枝变体。在通过TensorRT-LLM测量的运行时推理性能方面，MN-Minitron-8B模型相对于教师模型Mistral NeMo 12B提供了平均1.2倍的加速。同样，与教师模型Llama 3.1 8B相比，Llama-3.1-Minitron-4B模型的深度和宽度剪枝变体分别提供了平均2.7 ×和1.8×的加速。

二、论文的简单介绍

2.1 论文的背景

LLM提供商通常从头开始训练一整套模型，每个模型具有不同的规模（参数数量，例如Llama 3.1 8B、70B、405B）；这样做是为了帮助用户针对不同的部署规模、尺寸和计算预算。然而，从头开始训练多个数十亿参数的模型极其耗时、耗数据和资源密集。

近期研究[1]展示了将权重剪枝与知识蒸馏相结合，以显著降低训练大型语言模型（LLM）系列成本的有效性。在此方法中，仅对家族中最大的模型进行从头训练；其他模型则通过逐步剪枝较大模型并随后进行知识蒸馏以恢复剪枝模型的精度来获得。

论文成功地将Minitron压缩策略[1]应用于两个最先进的模型：Llama 3.1 8B和Mistral NeMo 12B，分别将它们压缩至40亿和8亿参数。图1提供了论文方法的高层次概览。

尽管遵循原始论文[1]，论文进行了关键修改：由于无法访问原始训练数据，论文在剪枝和蒸馏之前，先在论文的数据集上对教师模型进行微调。论文将此步骤称为教师校正。图4显示，省略教师校正会导致数据分布不匹配，从而对蒸馏产生负面影响。

图1 论文提出的剪枝和蒸馏方法的高层次概览。每个步骤使用的token总数显示在括号中。

2.2 论文的党发

论文方法的概述在图1中进行了说明。在这里，首先对教师模型在用于蒸馏的目标数据集上进行轻度微调 - 论文将这一步骤称为教师修正。接下来，应用剪枝来压缩模型，之后使用蒸馏来恢复任何损失的模型准确性。论文建议读者参考Minitron论文[1]以获取关于剪枝和蒸馏方法的完整描述。

表1 论文的MN-Minitron-8B和Llama-3.1-Minitron-4B模型的准确性数据。论文比较了论文的模型与同等规模的SoTA开放模型在各种常见语言建模基准测试上的表现。所有评估均由论文进行，除了标有的条目（取自相应论文）*

表 2 对齐的 Llama-3.1-Minitron 模型的准确率数据。论文在各种基准测试上将论文的模型与类似大小的最先进（SoTA）开放对齐模型进行比较。所有评估均由论文进行。表示在基准测试的代表性子集上获得的结果。最佳结果以粗体显示，次佳结果以下划线表示。MN-Minitron-8B 的对齐工作正在进行中，一旦准备就绪将发布*

在目标数据集上进行微调以用于蒸馏——论文称这一步骤为教师校正。接下来，应用剪枝来压缩模型，然后使用蒸馏来恢复任何丢失的模型准确性。

2.2.1 剪枝

权重剪枝是一种强大且广为人知的技术，用于减小模型大小。在本报告中，论文专注于结构化剪枝，即从模型权重中一次性移除非零元素的块（或通道）；结构化剪枝技术的例子包括神经元、注意力头、卷积滤波器和深度剪枝。对于大型语言模型（LLMs），如图2 所示，论文首先计算每个层、神经元、头和嵌入维度的重要性。然后论文对这些重要性分数进行排序，以计算相应的重要性排名。

重要性估计：论文使用一种纯粹基于激活的重要性估计策略，该策略使用一个小型校准数据集和仅前向传播，同时计算论文考虑的所有轴（深度、神经元、注意力头和嵌入通道）的敏感性信息。论文将深度剪枝视为一种特殊情况，不将其与压缩其他维度相结合。

图 2 原文 [1] 中概述的剪枝和蒸馏过程。论文在本工作中遵循相同的方法。

论文分别通过检查多头注意力（MHA）、多层感知器（MLP）和LayerNorm层产生的激活，计算每个头部、神经元和嵌入通道的重要性。为此，论文使用了一个小型校准数据集（1024个样本）。

对于深度剪枝，论文考虑了三个不同的指标来评估层的重要性：（1）语言模型验证损失，（2）块重要性（BI）5，以及（3）下游任务的准确性。对于基于损失的排名，论文简单地移除单个或连续的层块，并计算其对语言模型损失的影响；这作为该层的重要性或敏感度。BI使用层或层块的输入和输出之间的余弦距离。论文注意到BI和语言模型损失指标高度相关，但如图 8 和 9 所示，它们在下游任务上并不产生最准确的剪枝模型。因此，论文使用Winogrande基准 [6] 来评估层的重要性。

模型修剪：如图2所示，对于给定的架构配置，论文首先根据计算出的重要性对每个轴的元素进行排序，并直接对相应的权重矩阵进行修剪（重塑）。对于神经元和注意力头修剪，论文分别修剪MLP和MHA层的权重。在嵌入通道的情况下，论文修剪MLP、MHA和LayerNorm层权重矩阵的嵌入维度。原始方法( [ 1 ] )使用神经架构搜索( NAS )来寻找最佳架构；在这项工作中，论文跳过这一步，而是利用原始论文中的网络架构相关学习。

2.2.2 重新训练与知识蒸馏

论文使用“重新训练”一词来指代修剪后的精度恢复过程。在这项工作中，论文探索了两种重新训练策略：（1）利用地面真实标签的传统训练，以及（2）使用未修剪模型（教师）监督的知识蒸馏。知识蒸馏（KD）涉及从一个被称为教师的更大或更复杂的模型向一个被称为学生的更小/更简单的模型转移知识。知识转移通过让学生模型模仿教师模型的输出和/或中间状态来实现。在论文的案例中，未压缩和修剪的模型分别对应于教师和学生。对于蒸馏，论文遵循先前工作中的最佳实践[ 1 ]，仅在教师和学生对数上使用前向KL散度损失。这在图3中有所说明。

2.3 训练细节

2.3.1 预训练

Llama 3.1 8B和Mistral NeMo 12B在不同的专有数据集上进行了预训练，这些数据集论文无法访问。根据Llama 3.1技术报告，8B模型在15T token。论文使用在线可获得的相应基础模型，这些模型在Hugging Face平台上公开可用。

图3 蒸馏概览：若原始训练数据不可用，建议对教师模型进行轻微微调。随后通过最小化对数上的KL散度进行蒸馏，原始模型作为教师，剪枝模型作为学生。

数据集：论文使用的所有实验均基于Nemotron-4精心整理的持续训练数据集（CT）

2.3.2 剪枝

论文的简化剪枝方案基于Minitron论文中概述的最佳实践，并在方法论部分进行了描述。具体而言，对于宽度剪枝，论文（1）分别使用12-范数和均值作为批处理和序列维度的聚合函数，（2）执行一次性剪枝，避免迭代方法。对于深度剪枝，如方法论部分所述，论文遵循Gromov等人的观察结果，丢弃导致在Winogrande上精度下降最小的连续子组层。在此工作中，论文跳过了轻量级神经架构搜索（NAS）阶段，而是手动配置了Llama-3.1-Minitron-4B和MN-Minitron-8B的架构。论文提出的架构受到Minitron-4B和Minitron-8B模型的启发，详见表3。现在论文描述每个目标压缩模型的剪枝方案：

2.3.3 蒸馏

教师校正：直接使用Mistral NeMo 12B模型作为教师在论文的数据集上表现次优。这是由于教师模型训练所用的原始数据集与正在蒸馏的数据集之间子词token分布的变化所致。为了解决这个问题，论文首先使用约127Btoken在论文的数据集上对教师进行微调。如图所示，如果蒸馏过程中无法获得原始数据集，这种校正是必不可少的。因此，论文将此技术应用于Mistral-NeMo和Llama-3.1教师模型。微调过程对教师模型在下游任务上的精度影响较小，某些任务有所提升，而某些任务则有所下降，如表1所示。论文推测这可能是由于用于微调的数据集的人为因素。

表3 论文压缩模型的架构细节。

图 4 压缩学生模型的训练收敛图。论文比较了原始教师和修正教师的监督效果。

重新训练：根据 Mini-tron 工作中的经验，论文选择仅使用 logit 蒸馏，最小化教师和学生概率之间的前向 KL 散度损失，并完全忽略 LM 交叉熵损失。在这里，未剪枝和剪枝模型分别对应教师和学生。论文在蒸馏过程中使用表 4 中列出的超参数。论文使用 32 个 NVIDIA DGX 100 节点进行训练任务。

表 4 蒸馏训练期间使用的超参数。

2.3.4 指令调优

为了评估论文蒸馏模型的指令遵循能力，论文使用 NeMo-Aligner ll] 对 Llama-3.1-Minitron 4B 模型进行监督微调(SFT)，使用用于 Nemotron-4 340B 的指令调优数据集。如表 2 所示，论文评估了对齐模型在指令遵循和角色扮演（IFEval 和 MT-Bench）、RAG QA（ChatRAG-Bench）和函数调用能力（BFCL）方面的表现。

2.4 分析

论文进行了一系列消融研究，以更好地理解这些新模型的压缩特性。论文在本节中报告了论文的结果。

宽度与深度剪枝：图 5 显示了 Llama-3.1-Minitron-4B 在宽度与深度剪枝的训练曲线。论文注意到宽度剪枝导致较小的初始损失，并且尽管两种变体具有相同数量的参数，但始终优于深度剪枝模型。

剪枝与蒸馏：图6展示了论文提出的方法与剪枝和蒸馏的正交优势。论文比较了以下四种方法：（1）随机权重初始化和蒸馏，（2）随机剪枝和蒸馏，其中组件被随机剪枝，忽略重要性分数，（3）论文提出的剪枝方法与基于典型交叉熵的语言模型（LM）损失训练，以及（4）论文提出的剪枝方法与基于蒸馏的训练。论文注意到，与随机初始化相比，剪枝产生了一个显著更好的起点，并且基于蒸馏的训练在表现上优于传统的训练方法，同时所需训练token数量显著减少（在论文的案例中高达50倍）。

图 5 展示了宽度剪枝和深度剪枝的Llama 8B到4B模型的收敛情况。在给定的参数预算下，宽度剪枝始终优于深度剪枝。

教师校正：论文比较了两种教师校正方法：（1）剪枝并蒸馏校正后的教师模型，以及（2）剪枝原始教师模型并从持续校正的教师模型中蒸馏。图 7 的结果表明，教师校正不影响剪枝的最优性，而从校正后的教师模型中蒸馏是关键。教师校正可以与蒸馏并行进行以弥合差距。

图 6 Mistral Nemo 12B 压缩模型的训练收敛图。论文比较了 (a) 随机初始化与蒸馏，b) 随机剪枝权重与蒸馏，© 使用标准语言模型损失进行剪枝，以及 (d) 论文的剪枝与蒸馏流水线。

图 7 Mistral Nemo 12B 压缩模型的训练收敛图。论文比较了（1）修剪和蒸馏校正后的教师模型与（2）修剪原始教师模型并从持续校正的教师模型中蒸馏。

深度修剪指标：在检查语言模型验证损失随着连续层块被移除而增加的情况时（图8），论文观察到开始和结束处的层最为重要。移除非连续层可能会导致更好的语言模型验证损失（虚线）。然而，这一观察结果在评估下游任务性能时并不一定成立。图9显示，根据每层重要性选择并移除16层后，Wino-grande 准确率为0.5，而连续移除第16到31层后准确率为0.595。这一差距在基于蒸馏的再训练过程中保持不变，论文选择后者方法。

2.5评估

基准测试：遵循Touvron等人[18]的方法，论文在一系列下游任务上评估论文的压缩模型，包括MMLU、用于Python代码生成的HumanEval、几个用于常识推理的问答数据集：Arc-C、HellaSwag、TruthfulQA和WinoGrande，以及用于摘要生成的XL-Sum英文版。论文报告MMLU的5-shot性能，Winogrande的5-shot性能，ARC-Challenge的25-shot性能，HellaSwag的10-shot性能，XL-Sum 20%数据的0-shot性能，以及HumanEval和MBPP的平均pass@1分数。对于pass@1分数，论文使用温度为0.2的nucleus采样，top-p = 0.95。对于经过指令微调的模型，论文使用MT-Bench、Instruction-Following Eval (IFEval)、ChatRAG-Bench和Berkeley Function Calling Leaderboard (BFCL)。

图 8 | 在移除 Llama 3.1 8B 中连续的 1, 2、8 或 16 层后，验证集上的 LNM 损失值。例如，第 16 层的紫色线表示如果论文丢弃前 16 层的 LM 损失。第 17 层表示如果论文保留第一层而丢弃第 2 到 17 层的 LM 损失。虚线对应于移除 16 个非连续层（这些层对损失增加最少）时的 LM 损失值。

图 9 | 在移除Llama 3.1 8B模型中16个连续层时的Winogrande任务准确率。层号 1 7 表示如果论文保留第一层而丢弃第 2 到 1 7 层时的语言模型损失。虚线对应于移除16个非连续层时最少增加损失的准确率。

2.5.1 基础模型

基础模型评估结果如表1所示。与类似大小的模型相比，MN-Minitron- 8B 在各方面展现了更高的准确率，超过了最近的Llama 3.1 8B模型，使用 4 0 x 更少的训练token（380B vs. 15T）。同样，Llama-3.1-Minitron 4B模型在使用150倍更少的训练token（94B vs. 15T）的情况下，表现优于教师Llama 3.1 8B模型；论文的修剪Llama模型也优于前一代Minitron 4B模型。从表1中论文注意到，宽度修剪的变体优于深度修剪的变体。这些结果清楚地展示了论文方法的优势：最先进的准确率与训练效率的显著提升。

2.5.2 指令模型

指令调优的Llama-3.1-Minitron 4B变体的性能如表2所示。论文将Llama-3.1-Minitron 4B变体与其他类似大小的基线进行比较，并注意到论文的模型在遵循指令和角色扮演能力方面表现强劲，仅在IFEval和MT-Bench上落后于Gemma2。在基于检索的问答（ChatRAG-Bench 15）和函数调用（BFC）上，Minitron模型实现了最先进的性能。

2.6 洞察

论文总结了一些有趣且令人惊讶的观察结果。

通用

1.教师模型的校正对于在新、未见数据集上最优地进行蒸馏至关重要。以这种方式使用蒸馏数据集微调教师模型，可以使语言模型验证损失减少超过 6%。教师模型的校正不影响剪枝的最优性，甚至可以与蒸馏并行进行。

2.与 Minitron 论文的观察一致，论文仅需要 380B 个token就能在剪枝后通过蒸馏达到最先进的准确性。

3.对于宽度剪枝，论文通过保留注意力头并剪枝其他维度（MLP 中间维度，嵌入通道）来实现更高的准确性。

Mistral NeMo 12B 到 MN-Minitron-8B：

1.论文的压缩模型在剪枝和蒸馏后，在两个基准测试 GSM8k 和 HumanEval 上优于教师模型：GSM8k 从 5 5. 7% 提高到58.5%，HumanEval 从23.8% 提高到 36. 2%。这种改进可能受数据集影响。然而，重训练仅使用蒸馏损失进行。

Llama. 3.1 8B 到 Llama-3.1-Minitron 4B：

1.宽度剪枝在 MMLU 上达到更好的准确性，为 60.5%，而深度剪枝为58.7%，适用于 Llama-3.1 压缩。

2.推理能力受到进一步显著影响，GSM8K 准确性在宽度为 41.24 %，深度为16.8%。

3.深度剪枝提高吞吐量，实现2. 7×比 Llama-3.1 8B 加速，而宽度剪枝提供1. 7×加速。

4.对于深度剪枝，论文观察到从模型中丢弃连续层比使用非连续、基于重要性的剪枝更有效。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述