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

程序猿李巡天

于 2024-09-26 19:52:58 发布

阅读量602

点赞数 6

文章标签：剪枝人工智能算法语言模型知识图谱分布式机器学习

本文链接：https://blog.csdn.net/m0_59235945/article/details/142574643

版权

一、背景

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

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

二、摘要

本文中，作者探索了两种不同的 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 来增加吞吐：

如何学习大模型 AI ？

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

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

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

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

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

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

在这里插入图片描述