14-LLM Pruning and Distillation in Practice: The Minitron Approach

一、前言

我们将llama3 8B和Mistral NeMo12B的大模型压缩到了4B和8B.

我们探索了两种不同的修剪策略：
（1）深度修剪
指的是从模型的深度维度进行剪枝，即减少模型的层数。这种方法通过去除不重要的层，减少模型的整体计算量和复杂度。
（2）联合隐藏层/注意力/MLP（宽度）剪枝
在这个策略中，模型的隐藏层、注意力头、以及MLP层的宽度（通道数或神经元数）被一起剪枝。这种方法通过减少这些组件的维度来压缩模型。

我们发现，当没有访问原始数据时，对教师模型在蒸馏数据集上进行轻微的微调是有益的。

二、引言

最近的工作证明了将权重剪枝与知识蒸馏相结合可以显着降低LLM模型族训练成本的有效性。只有模型族里面最大的模型会从头开始训练，其他的模型是通过修剪较大的模型得到的，然后再通过知识蒸馏来恢复剪枝模型的准确性。
本论文成功将Minitron压缩策略应用到了llama3 8B和Mistral NeMo 12B。
关键修改：由于无法访问到原始训练数据，我们在修剪和蒸馏之前，将原始模型在我们的数据集上进行了微调。

三、方法论

在这里，教师模型首先在目标数据集上进行轻微微调以用于蒸馏。然后，应用剪枝来压缩模型。最后，使用蒸馏来提升被剪枝模型的精度。

剪枝

权重剪枝是一种有效且出名的减小模型大小的技术（权重剪枝：通过在训练过程中逐渐将模型权重归零，以实现模型稀疏化。包括但不限于全局剪枝、层剪枝、结构化剪枝等）。但是，本论文聚焦于结构化剪枝（非零元素块（或通道）立即从模型权重中删除）。
结构化剪枝技术的示例包括神经元、注意力头、卷积滤波器和深度剪枝

结构化剪枝的步骤

结构化剪枝的过程通常包括以下步骤：

1、计算重要性：根据特定标准计算每一层、神经元、注意力头或嵌入维度的重要性。
2、排序重要性分数：对计算出的重要性分数进行排序，从而为模型的组件（如神经元、头等）生成一个重要性排名。
3、剪枝：根据排名，移除重要性较低的组件。

重要性估计

我们使用纯基于激活的重要性估计策略，该策略使用小型数据集同时计算我们考虑的所有轴（深度、神经元、头部和嵌入通道）的灵敏度信息（仅在前向传播中）。
我们通过分别检查多头注意力（MHA）、多层感知器（MLP）和 LayerNorm 层产生的激活来计算每个头、神经元和嵌入通道的重要性。

评估重要性有三个不同的指标：
（1）LM validation loss（LM验证损失）
对于基于损失的排名，我们只需删除单个或一组连续层并计算其对 LM 损失的影响； 这充当该层的“重要性”或敏感性
（2）Block Importance (BI) （块重要性）
BI 使用层或层块的输入和输出之间的余弦距离。 我们注意到 BI 和 LM 损失指标高度相关，但不会在下游任务上产生最准确的修剪模型
（3）accuracy on the downstream task （下游任务的准确性）

综上，因此，我们使用 Winogrande 基准评估层重要性

模型修剪

模型修剪是一种模型压缩技术，旨在通过减少模型中的参数数量来降低模型的复杂性和计算成本，同时尽量保持模型性能。与原始方法通过 Neural Architecture Search (NAS) 来寻找最佳架构不同，论文中所描述的方法直接基于给定的架构配置，通过评估元素的重要性并相应地修剪（或重塑）权重矩阵来实现模型精简。
步骤如下：
（1）元素重要性评估：对于模型中的每个层（如MLP多层感知机、MHA多头注意力层等），首先根据一定的标准（如权重的大小、梯度信息等）计算每个元素（或特征、神经元、注意力头等）的重要性。
（2）权重矩阵修剪：基于计算出的重要性排名，直接对相应的权重矩阵进行修剪（reshaping）。这意味着删除或合并一些不重要的元素，从而减少矩阵的大小。在神经元和头部修剪的情况下，会分别针对MLP层和MHA层中的权重进行操作。
（3）嵌入通道修剪：对于嵌入层（embedding layers），则是对嵌入矩阵的维度进行修剪。这意味着减少嵌入向量中的特征数量，从而影响MLP、MHA和LayerNorm层中的权重矩阵大小。
（4）利用已有架构知识：不同于原始方法通过NAS来寻找最优架构，论文方法利用了从原始论文中学到的与网络架构相关的知识，直接在这些固定架构上进行修剪。这大大简化了流程，并可能提高修剪过程的效率和效果。

蒸馏再训练

蒸馏再训练是一种在模型剪枝后常用的策略，旨在通过蒸馏的方式将剪枝后的模型性能恢复或提升。

两种重训练策略：

1、常规训练（Conventional Training）：
这种方法利用真实的标签（ground truth labels）对剪枝后的模型进行训练。常规训练是机器学习中最直接的方法，它通过最小化预测输出与真实标签之间的误差来优化模型参数。
2、知识蒸馏（Knowledge Distillation, KD）：
知识蒸馏是一种模型压缩技术，它将一个大型或复杂的模型（教师模型）的知识转移到一个小型或简单的模型（学生模型）中。未压缩的模型和剪枝后的模型分别对应于教师模型和学生模型。

对于知识蒸馏，我们使用了前向KL散度损失（forward KL Divergence loss）来度量教师模型和学生模型输出之间的差异。KL散度是一种衡量两个概率分布之间差异的方法，通过最小化这种差异，可以使学生模型的输出更加接近教师模型的输出。

四、训练细节

预训练

由于我们没有权限访问llama3 8B和Mistral NeMo 12B的预训练的专有数据集。因此我们使用了Nemotron-4 curated continued training (CT) dataset来实现预训练。

剪枝

宽度剪枝

（1）使用l2-norm和均值作为聚合函数

• 在批量（batch）维度上，选择了l2-norm（也称为欧几里得范数）作为聚合函数。l2-norm可以计算每个神经元（或特征）在所有样本上的权重向量的欧几里得长度，从而评估其重要性。较高的l2-norm值通常意味着该神经元对模型的贡献更大，因此在剪枝时更有可能被保留。
• 在序列（sequence）维度上，选择了均值作为聚合函数。这意味着对于序列数据（如文本、时间序列等），您会计算每个神经元在序列所有位置上的权重的平均值，以此来评估其重要性。这种方法有助于捕获神经元在整个序列上的总体影响。

深度剪枝

在深度剪枝方面，遵循了Gromov等人的观察结果，并采用了在Winograd Schema Challenge上导致精度下降最小的连续层子组进行剪枝。与宽度剪枝不同，深度剪枝关注的是移除模型中的整个层或层组，以简化模型结构并减少计算量。
跳过了轻量级神经网络架构搜索（NAS, Neural Architecture Search）阶段，而是为Llama-3.1-Minitron-4B和MN-Minitron-8B模型选择了手动架构配置。

蒸馏

Teacher Correction 的应用是为了解决在使用预训练模型（如 Mistral NeMo 12B）作为教师模型时，由于数据集分布差异导致的性能次优问题。这种情况通常发生在教师模型训练时所用的原始数据集与当前蒸馏过程中使用的数据集在子词标记（sub-word tokens）的分布上存在显著差异。

为了克服这一挑战，采取了先对教师模型进行微调（fine-tuning）的策略，具体是在自己的数据集上使用大约1270亿个标记（tokens）来重新训练教师模型。这一步骤的目的是使教师模型更好地适应当前数据集的特性，从而提高蒸馏过程的效率和效果。

当原始数据集在蒸馏过程中不可用时，这种教师模型的微调是至关重要的。

微调过程对教师模型在下游任务上的准确性影响较小，有些任务的性能得到了提升，而有些则略有下降。推测这种差异可能是由微调时使用的数据集的特性所导致的。不同的数据集可能包含不同的词汇、语法和语义信息，这些信息在微调过程中会影响到模型的参数分布和特征表示，进而对下游任务的性能产生影响。

总的来说，教师模型的微调是一种有效的策略，可以帮助克服数据集分布差异带来的问题，提高蒸馏过程的性能和效率。然而，需要注意的是，微调的效果也受到所用数据集特性的影响，因此在选择微调数据集时需要谨慎考虑。

再训练/重训练

重训练（Retraining）过程中，采用了Mini-tron工作中的经验，选择了仅对数（logit-only）蒸馏的方法。这种方法专注于最小化教师模型和学生模型在预测概率分布上的前向KL散度（KL Divergence）损失，而完全忽略了语言模型（LM）的交叉熵损失。
在蒸馏过程中，您使用了下表中列出的超参数配置。这些超参数对于控制蒸馏过程的优化方向、学习速率、正则化强度等方面至关重要，有助于确保蒸馏后的学生模型能够继承教师模型的有效知识，同时保持较小的模型尺寸和计算复杂度。

指令调优

为了评估蒸馏后模型的指令遵循能力，您采用了监督微调（SFT, Supervised Fine-Tuning）的方法，在Llama-3.1-Minitron 4B模型上使用了NeMo-Aligner]工具，并结合了为Nemotron-4 340B 设计的指令微调数据集。这种微调过程旨在使模型更好地理解和执行自然语言指令。
论文中展示了经过对齐（aligned）的模型在多个基准测试上的性能评估结果，包括指令遵循和角色扮演、基于检索的问答以及函数调用能力

• 指令遵循和角色扮演（IFEval 和 MT-Bench）：这些基准测试用于评估模型理解和执行自然语言指令的能力，包括文本生成、逻辑推理、数学计算等多种任务。通过在这些测试上的表现，可以直观地了解模型在指令遵循方面的能力。
• 基于检索的问答（RAG QA）：ChatRAG-Bench是一个评估基于检索的增强型生成模型（如结合了检索机制的LLMs）在问答任务上性能的基准测试。通过这一测试，可以评估模型在利用外部知识库进行问答时的能力。
• 函数调用能力（BFCL）：BFCL是一个评估大型语言模型在函数调用和代码生成方面能力的基准测试。它要求模型根据给定的自然语言指令生成相应的程序代码或函数调用序列。