LLM 剪枝+蒸馏：NVIDIA 的最佳实践 精华

一、背景

模型剪枝、蒸馏是传统 AI 模型压缩常用的方案，尤其是模型要用于端侧部署的场景，相比模型量化，模型剪枝和蒸馏的代价比较高，而且往往在指标上也没有特别明显的优势，因此真正落地的场景要少得多。当然，也传言有些模型会蒸馏 OpenAI 的 ChatGPT，不过主要是用其生成高质量数据。本文中，我们介绍 NVIDIA 最近发布的 LLM 剪枝和蒸馏的最佳实践。

对应的论文为：[2408.11796] LLM Pruning and Distillation in Practice: The Minitron Approach

我们之前也介绍过 NVIDIA 的两篇相关文章，可以参考：

• NVIDIA LLM 持续预训练的最佳实践
• NVIDIA 开源 LLM：迷你杯 Nemotron-4 15B & 超大杯 Nemotron-340B

二、摘要

本文中，作者探索了两种不同的 LLM 剪枝策略：（1）深度剪枝和（2）联合 Hidden、Attention、MLP 的宽度剪枝，并使用 LM 评估框架 Harness 进行常见基准评估。然后将模型与 NeMo Aligner 对齐，并在 Instruct-tuned 版本中进行测试。使用提出的剪枝和蒸馏将 LLaMA 3.1 8B 和 Mistral NeMo 12B 模型分别压缩为 4B 和 8B 参数的模型。此外，作者发现，在无法访问原始数据的情况下，在蒸馏数据集上适当微调 Teacher 模型也有帮助。

PS：基于当前论文的工作还有许多可以尝试的方案，比如：

• 论文中宽度剪枝后每个 Transformer Block 层的参数还保持一致，而根据之前的研究（本文中也提到了），模型的最开始 2 层和最后 2 层可能更加的重要，是否可以通过某种方式实现只剪枝中间层，而保留最前 2 层和最后 2 层，效果是否会更好？
• 如果使用 LLaMA 3.1 70B 作为 Teacher 模型，是否能将 4B 参数量规模的模型精度与 8B 模型对齐？

三、方案

3.1 概览

如下图 Figure 1 所示为本文方法的概览，其包含 3 个阶段：

• Teacher 校正：在目标数据集（127B）上对预训练模型进行微调，生成校正的 Teacher 模型，以便用于蒸馏。
• 剪枝：应用剪枝技术压缩模型，生成 Student 模型。（PS：需要说明的是，非结构化剪枝往往导致模型无法充分发挥 GPU 算力，所以在使用 GPU 推理的场景中相对较少，这里作者主要是使用的结构化剪枝，非常适合 GPU 运算）
• 蒸馏：使用 Teacher 模型蒸馏 Student 模型，以恢复剪枝损失的模型准确性。​

3.2 剪枝

权重剪枝是一种强大且众所周知的模型压缩技术。本文中，作者重点介绍结构化剪枝（并不是为了适配 GPU 稀疏算力的 2:4 稀疏结构化剪枝，其可以参考附录），也就是从模型中删除 Block 或 Channel（PS：不是将其置为 0），包括 Neuron、Attention Head、Convolutional Filter 和深度剪枝。如下图 Figure 2 所示，对于 LLM 而言，首先计算每个层、神经元、Head 和 Embedding 维度的重要性；然后对这些重要性分数进行排序；最后进行剪枝操作，并多次迭代。

3.2.1 重要性预估

作者采用纯粹基于激活的重要性评估策略，该策略使用小型校准数据集，通过前向推理来计算所有轴（深度、神经元、Head，嵌入通道）的灵敏度信息。此外，作者将深度修剪作为一种特殊情况，不会与其他压缩维度结合使用。

具体来说，作者使用一个 1024 个 Sample 的小型校准数据集，通过分别检查 MHA、MLP 和 LayerNorm 层产生的激活来计算每个 Head、神经元和 Channel 的重要性。

对于深度剪枝，作者使用 3 个指标评估 Layer 的重要性：

• LM 验证损失。
• Block 重要性（BI）。
• 下游任务的准确性。

对于基于 Loss 的排序，只需删除单个或连续的 Block，并计算其对 LM Loss 的影响，这可以作为层的“重要性”或“敏感度”。BI 使用 Layer 或 Layer Blocks 的输入和输出之间的余弦距离来计算。作者注意到 BI 和 LM 损失指标高度相关，但并没有在下游任务上生成最准确的剪枝模型。因此，作者使用 Winogrande 基准来评估 Layer 的重要性。

3.2.2 模型修剪

对于给定的模型，获得每个轴的重要性排名之后，可以直接对相应的权重矩阵进行修剪。对于神经元和 Head 修剪，分别修剪 MLP 和 MHA 层权重；对于 Embedding Channel，修剪 MLP、MHA 和 LayerNorm 中权重矩阵的 Embedding 维度。

3.2.3 蒸馏训练

对修剪后的模型进行 ReTraining 以恢复准确性。本文中，作者探索了两种 ReTraining 策略：

• 利用 Ground truth 标签的常规训练。
• 使用未修剪模型（Teacher）进行监督知识蒸馏。

蒸馏的过程如下图 Figure 3 所示，作者只在最后的 Logits 上添加 KL 散度损失。

3.3 训练详情

3.3.1 预训练

使用预训练的 LLaMA 3.1 8B（[2407.21783] The Llama 3 Herd of Models） 和 Mistral Nemo 12B 模型（Mistral NeMo | Mistral AI | Frontier AI in your hands）。

3.3.2 数据集

所有实验使用 Nemotron-4 （[2402.16819] Nemotron-4 15B Technical Report 和 [2407.07263] Reuse, Don't Retrain: A Recipe for Continued Pretraining of Language Models）的 Continuous Training（CT） 数据集。

3.3.3 剪枝

作者采用的简化剪枝方案来自 Minitron 论文（[2407.14679] Compact Language Models via Pruning and Knowledge Distillation）中的最佳实践。具体来说：

• 宽度剪枝：

分别使用 l2-norm 和 mean 作为跨 Batch 和 Sequence 维度的聚合函数。

执行单次修剪，避免迭代方案。

• 深度剪枝：
• 遵循[2403.17887] The Unreasonable Ineffectiveness of the Deeper Layers 中的观察结果，删除一个连续 subgroup，该 subgroup 会使 Winogrande 的准确性下降最小。
• 没有采用 NAS 搜索的结构。

修剪后的 3 个模型参数如下图所示，这里也可以看出，宽度剪枝后各层的超参数一致：

3.3.4 蒸馏

Teacher 校正：直接使用 Mistral Nemo 12B 模型在作者自己的数据集上表现不佳，这是由于 Teacher 模型训练的原始数据集与蒸馏数据集的分布不一致。为了解决这个问题，作者首先使用数据集中 >= 127B 的 Token 微调 Teacher 模型。如下图 Figure 4 所示，使用经校正的 Teacher 模型蒸馏，Student 模型在验证集上的 Loss 明显低于使用原始 Teacher 模型。因此，作者将这种方案应用到了 Mistral-Nemo 和 LLaMA 3.1 Teacher 模型。当然，微调 Teacher 模型也会导致其在一些指标上有所提升，而在一些指标上有所下降：

Retrining：根据 Minitron 中的方案，作者选择仅 Logit 蒸馏（PS：之前很多工作也会蒸馏 Feature Map），最大限度减少 Teacher 和 Student 的 KL 散度损失，并完全忽略 LM 交叉熵损失。蒸馏的超参数如下图 Table 4 所示，在 32 个 DGX H100 节点上训练。

指令微调：为了评估蒸馏模型的指令跟随能力，作者使用 Nemo-Aligner 和用于 Nemotron-4 340B 的指令微调数据集对 LLaMA 3.1 Minitron 4B 模型进行 SFT，结果如下图 Table 2 所示：

四、分析和评估

4.1 分析

4.1.1 宽度和深度剪枝

如下图 Figure 5 所示为根据宽度和深度剪枝的 LLaMA-3.1-Minitron-4B 的训练曲线，可以看出，两者具有相同的参数量，但是宽度剪枝对应的初始损失更小，并且始终优于深度剪枝。

4.1.2 剪枝和蒸馏

如下图 Figure 6 展示了剪枝和蒸馏方法的正交优势。作者比较了下述 4 种方案，可以看出，与随机初始化相比，剪枝的起点明显更好，而基于蒸馏的训练优于传统的训练方法，同时需要训练的 Token 明显减少：

• Random Init + Distillation：随机权重初始化和蒸馏。
• Random Pruning + Distillation：随机剪枝和蒸馏。其中的组件被随机修剪而不是依赖重要性分数。
• Pruning + LM Loss：使用本文的修剪方案，但使用基于交叉熵的 LM Loss 训练。
• Pruning + Distillation：本文的剪枝和蒸馏方案。LM 验证损失最低。​

4.1.3 Teacher 校正

如下图 Figure 7 所示，作者对比了两种 Teacher 校正方法，结果表明，Teacher 校正并不影响剪枝的最优性，用校正后的 Teacher 至关重要。Teacher 校正也可以与蒸馏同时进行，以弥合差距：

• Prune corrected teacher + distill corrected teacher：剪枝和蒸馏校正的 Teacher 模型。
• Prune original teacher + distill continuously corrected teacher：剪枝原始的 Teacher 模型，并使用不断校正的 Teacher 模型来蒸馏。​

4.1.4 深度剪枝度量

在检查 LM 验证损失如何随着连续 Layer Block 的删除而增加时，如下图 Figure 8 所示，作者观察到开始和结尾的层是最重要的。删除非连续层可能导致更好的 LM 验证损失。

但是在评估下游任务时，上述结论不一定成立。如下图 Figure 9 所示，根据每层重要性删除 16 层，Winogrande 精度为 0.5，而连续删除 16-31 层的精度为0.595。在基于蒸馏的 Retraining 中，差距仍然存在，作者选择了后一种方法。

4.2 评估

4.2.1 Base 模型评估

如下图 Table 1 所示为 Base 模型的评估，与类似大小的模型相比，MN-Minitron8B 在各方面都表现出卓越的准确性，并且训练 Token 数小 40x（380B vs 15T）。同样，与 Teacher LLaMA 3.1 8B 模型相比，LLaMA-3.1 4B 模型表现良好，并且使用的训练 Token 减少 150x（94B vs 15T）。剪枝后的 LLaMA-3.1 4B 也优于之前的 Minitron 4B。此外，从中也可以看出，基于宽度剪枝的变体优于基于深度剪枝的变体。这些结果充分表明了方案的有效性。

PS：不过从 8B -> 4B 的损失依然比较大，甚至和直接进行 AWQ（W4A16）的量化损失差不多。当然 8B AWQ 的推理效率可能不如 4B，然而一些 W8A8 的方案也能获得相当的精度，详情可以参见后文“附录”的量化部分。此外，量化的成本可能远低于剪枝+蒸馏。如下图所示为 neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w4a16 · Hugging Face 中使用 W4A16 的 Instruct 模型精度：

4.2.2 Instruct 模型评估

如下图 Table 2 所示为 LLaMA-3.1 Minitron 4B 指令微调的性能。可以看出，其宽度剪枝变体在所有指标上优于原始的 Minitron 4B：

五、附录

5.1 模型压缩

模型压缩有 4 种常见的方案：量化，剪枝，蒸馏，低秩分解。在 LLM 场景中，模型量化的方案非常多，比如 llm.int8()、AWQ、GPTQ、SmoothQuant 等等，其实现简单，代价小，是最常见的方案。而其他几种方案应用相对比较少，它们的区别如下图 Figure 2 所示，图片来自 [2308.07633] A Survey on Model Compression for Large Language Models：

5.2 量化

量化是模型压缩中最常用的技术，不同的量化方案会针对不同的精度、Tensor 类型等，比如有常见的 KV Cache Only 量化，Weight Only 量化，以及几种方案的结合，具体如下图所示：

不同的量化方案在不同模型上的量化损失也会有所不同，但是大体上来说，压缩后的 Bit 数越低损失越大。如下图 Table 1 所示为 [2404.14047] An Empirical Study of LLaMA3 Quantization: From LLMs to MLLMs 中对 LLaMA3-8B 模型的量化评估（都使用 INT 类型，未使用 FP8/FP6/FP4），可以看出：

• W8A16 的量化损失都比较小，几乎无损
• W4A16 和 W8A8 的损失相比 W8A16 会大一些，大部分情况也基本可以接受，但也和量化方法有关，比如，GPTQ 和 QuIP 的 W4A16 相比 AWQ 的损失会更大一些。
• 更低 Bit 的量化损失会比较大，比如 W3A16 和 W2A16 的精度会明显下降。​

NVIDIA 的 GPU 从 Hopper 架构开始可以支持 FP8 计算，使用 FP8 后其精度相比 SmoothQuant 的 INT8 以及其他的 W4A16 损失更小，更具有使用价值（数据来自 https://github.com/NVIDIA/TensorRT-LLM/blob/v0.9.0/docs/source/blogs/quantization-in-TRT-LLM.md）：

那么这些量化方案的速度怎么样呢，如下图所示，在 [2404.14294] A Survey on Efficient Inference for Large Language Models 中作者评估了 TensorRT-LLM 和 LMDeploy 推理框架在不同场景的 W4A16 推理性能，使用的 GPU 为 NVIDIA A100，图中的数据分别为 Prefill/Decoding/End2End 的加速比，可以看出，基本都可以实现 2 倍左右加速，当序列比较长或者 Batch size 比较大时会略低一些，当然两个框架也各有千秋：

• TensorRT-LLM 在 Batch size 比较小时优势更明显。
• LMDeploy 在 Batch size 比较大时优势更明显。​

如下图所示为使用 FP8 相比 FP16 可以加速 1.4-1.5 倍，这是比较早的数据，进一步优化后可以到 1.6-1.7 倍：

如下表所示为 TensorRT-LLM 中作者对不同量化方案的总结，可以看出，如果 GPU 支持 FP8，则使用 FP8 是最理想的选择，如果允许少量精度损失，则可以进一步使用 INT4-FP8 AWQ 的方案：

5.3 剪枝

NVIDIA 在 Ampere 架构的 Tensor Core 中引入了稀疏矩阵乘法支持，理论最多可以提升 2 倍性能，实际上可能只有 1.5 倍，而且对稀疏化的方式有要求，如下图所示，每 4 个 Weight 中需要有 2 个是 0 值，也就是可以剪掉的值：

基本上稀疏化都会带来一定的精度损失，如下图 Table 2 所示，论文 [2310.15929] E-Sparse: Boosting the Large Language Model Inference through Entropy-based N:M Sparsity 中作者评估了 2:4 稀疏化对模型精度的影响，可以看出（PS：论文图片中的部分数字有误，比如 13B 模型上 Magnitude 平均精度达到了 57.71，实际平均只有 49.36）：

• 每种稀疏化方法都有一定的精度损失，即使最优的方案损失也比较大。
• 不同规模的模型上精度损失差异比较大，小模型损失比较大。
• 13B 模型稀疏化后的精度还不如原始的 7B 模型，那么 13B 模型的稀疏化基本没有太大的意义。​

在这样的损失下能带来什么收益呢？其收益主要体现在速度的提升和显存的节约，如下图 Table 5 所示，其矩阵乘法可以加速 20%-30%，而端到端延时只能加速 16% 左右：

当然，显存的节约还是比较客观的，基本可以节约 43% 左右的显存空间，也许可以通过增加 Batch Size 来增加吞吐：

六、参考链接

1. https://arxiv.org/abs/2408.11796
2. https://arxiv.org/abs/2407.21783
3. https://mistral.ai/news/mistral-nemo/
4. https://arxiv.org/abs/2402.16819
5. https://arxiv.org/abs/2407.07263
6. https://www.arxiv.org/abs/2407.14679
7. https://arxiv.org/abs/2403.17887
8. https://arxiv.org/abs/2308.07633
9. https://arxiv.org/abs/2404.14047
10. https://arxiv.org/abs/2404.14294
11. https://arxiv.org/abs/2310.15929

本文转载自 ​​AI闲谈​​，作者： AI闲谈