由于大型语言模型（LLM）的有效性和通用性，它们现在已经成为自然语言处理和理解领域的主导力量。LLM（例如 Llama 3.1 405B 和 NVIDIA Nemotron-4 340B）在编码、推理和数学等许多具有挑战性的任务中表现出色。但是，它们的部署需要大量资源。因此，业内出现了另一种趋势，即开发小型语言模型（SLM），这些模型在许多语言任务中足够精通，但部署给大众的成本要便宜得多。
最近，NVIDIA 研究人员表明，结构化权重剪枝与知识提炼相结合，形成了一种有效且高效的策略，可以从初始较大的同级产品中逐步获得较小的语言模型。NVIDIA Minitron 8B 和 4B 是如此小的模型，通过在 NVIDIA Nemotron 系列中剪枝和提炼其较大的 15B 同级产品来获得。
剪枝和提炼可带来以下优势：
- 与从头开始训练相比，MMLU 分数提高了 16%。
- 每个新增模型所需的训练令牌更少，约 ~100B 个令牌，减少高达 40 倍。
- 与从头开始训练所有模型相比，训练一系列模型可节省高达 1.8x 的计算成本。
- 性能可与使用高达 15T 令牌训练的 Mistral 7B、Gemma 7B 和 Llama-3 8B 相媲。
本文还介绍了一套适用于 LLM 的实用且有效的结构化压缩最佳实践，这些实践将深度、宽度、注意力和 MLP 剪枝与基于知识提炼的再训练相结合。
在本文中，我们首先讨论这些最佳实践，然后展示其在应用于 Llama 3.1 8B 模型以获得 Llama-3.1-Minitron 4B 模型时的有效性。Llama-3.1-Minitron 4B 与类似大小的先进开源模型（包括 Minitron 4B、Phi-2 2.7B、Gemma2 2.6B 和 Qwen2-1.5B）相比性能较好。Llama-3.1-Minitron 4B 将很快发布到 NVIDIA Hugging Face 集合中，等待批准。
修剪和提炼
剪枝是通过丢弃图层（深度剪枝）或丢弃神经元、注意力头和嵌入通道（宽度剪枝）来缩小模型并使其更加精简的过程。通常情况下，剪枝会伴随一定数量的重新训练，以恢复模型的准确性。
模型提炼是一种技术，用于将知识从大型、复杂的模型（通常称为教师模型）转移到更小、更简单的学生模型。其目标是创建更高效的模型，在更快、更低资源消耗的情况下，保留大型原始模型的大部分预测能力。
经典知识提炼与 SDG 微调
主要有两种distillation:
- SDG 微调：较大的教师模型生成的合成数据用于进一步微调较小的预训练学生模型。在这里，学生仅模仿教师预测的最终令牌。例如 Azure AI Studio 中的 Llama 3.1 Azure Distillation 和 AWS 使用 Llama 3.1 405B 生成合成数据和进行提炼以微调较小的模型教程。
- 经典知识提炼：学生在训练数据集上模仿教师的 logits 和其他中间状态，而不仅仅是学习必须预测的令牌。这可以被视为提供更好的标签（与一次性标签相比的分布）。即使使用相同的数据，梯度也包含更丰富的反馈，从而提高训练准确性和效率。但是，由于 logits 太大，无法存储，因此必须为这种提炼方式提供训练框架支持。
这两种提炼方式是相互补充的，而不是相互排斥的。本文主要介绍经典知识提炼方法。
修剪和提炼程序
我们建议将剪枝与经典知识提炼相结合作为一种资源节约型再训练技术（图 1）。
- 我们从一个 15B 的模型开始。我们估计每个组件（层、神经元、头部和嵌入通道）的重要性，然后对模型进行排名并将其裁剪为目标大小：一个 8B 的模型。
- 我们使用模型提炼执行轻度再训练程序，原始模型作为教师，经过剪枝的模型作为学生。
- 训练完成后，小型模型 (8B) 将作为起点，用于裁剪和提取较小的 4B 模型。
图 1 显示了单个模型的剪枝和提炼过程（上）以及模型剪枝和提炼链（下）。在后者中，前一个阶段的输出模型将作为下一个阶段的输入模型。
重要性分析
要对模型进行剪枝，务必要了解模型的哪些部分很重要。我们建议使用纯粹基于激活的重要性估计策略，该策略使用小型（1024 个样本）的校正数据集，同时计算考虑的所有轴（深度、神经元、头部和嵌入通道）的敏感度信息，并且仅使用前向传播通道。与依赖梯度信息并需要反向传播通道的策略相比，此策略的实施更直接、更经济高效。
在剪枝时，您可以在给定轴或轴组合的剪枝和重要性估计之间进行迭代交替。但是，我们的经验工作表明，使用单次重要性估计已经足够，而迭代估计并没有任何好处。
通过经典知识提炼进行再训练
图 2 展示了学生模型 (剪枝模型) 的提炼过程，该模型从具有 M 层的教师模型 (原始未剪枝模型) 中提炼 N 层。学生通过最小化嵌入输出损失、对数损失以及跨学生块 S 和教师块 T 映射的 Transformer 特定编码器损失的组合来学习。
剪枝和提炼最佳实践
基于我们在紧凑语言模型中通过剪枝和知识提炼进行的大量消融研究，我们将所学到的经验总结为几种结构化压缩最佳实践：
- Sizing:
- 要训练一系列LLMs，请先训练最大的LLM，然后进行迭代剪枝和提炼，以获得较小的LLMs。
- 如果使用多阶段训练策略训练最大的模型，则最好对从训练的最后阶段获得的模型进行剪枝和重新训练。
- 剪枝最接近目标大小的可用源模型。
- Pruning:
- 首选宽度而非深度剪枝。这在所考虑的模型比例 (约 150B) 方面效果很好。
- 使用单次重要性估计。迭代重要性估计没有任何益处。
- Retraining:
- 仅使用蒸损失进行重新训练，而不是使用传统训练。
- 当深度显著降低时，使用logit加中间状态加嵌入提炼。
- 当深度未显著降低时，使用 Logit 纯提炼。
Llama-3.1-Minitron：将最佳实践付诸实践
Meta 最近推出了功能强大的 Llama 3.1 模型系列，这是第一波开源模型，在许多基准测试中可与闭源模型相媲美。Llama 3.1 的范围从庞大的 405B 模型到 70B 和 8B。
凭借 Nemotron 提炼经验，我们着手将 Llama 3.1 8B 模型提炼成更小、更高效的 4B 同级产品：
- 教师微调
- 仅深度剪枝
- 仅宽度的剪枝
- 准确性基准测试
- 性能基准测试
教师微调
为了修正训练模型所用的原始数据集上的分布偏移，我们首先在数据集上微调了未删减的 8B 模型（94 亿个令牌）。实验表明，在不修正分布偏移的情况下，教师模型在提炼数据集时提供次优指导。
仅深度剪枝
为了从 8B 扩展到 4B，我们剪枝了 16 层（50%）。我们首先评估了每一层或连续子组的重要性，方法是将其从模型中删除，并观察下游任务的 LM 损失增加或准确性降低。
图 5 显示了在验证集上删除 1、2、8 或 16 层后 Language Model 的损失值。例如，如果我们删除前 16 层，则第 16 层的红色图表示 Language Model 的损失。如果我们留下第一层并删除第 2 至 17 层，则第 17 层表示 Language Model 的损失。我们观察到开始和结束的层最为重要。
但是，我们发现此 LM 损失不一定与下游性能直接相关。
图 6 显示了每个剪枝模型的 Winogrande 精度。这表明最好删除第 16 层到第 31 层，其中 31 层是倒数第二层，其中剪枝模型的 5 次射击精度明显高于随机（0.5）。我们根据这个见解删除了第 16 层到第 31 层。
仅宽度的剪枝
我们沿着宽度轴剪枝了嵌入（hidden）和 MLP 中间维度，以压缩 Llama 3.1 8B。具体来说，我们使用前面介绍的基于激活的策略计算了每个注意力头、嵌入通道和 MLP 隐藏维度的重要性分数。根据重要性估计，我们：
- 将 MLP 中间维度从 14336 剪枝 (修剪) 到 9216、
- 将隐藏大小从 4096 缩减到 3072、
- 重新训练注意力头数和层数。
值得一提的是，在一次性剪枝后，宽度剪枝的 LM 损失要高于深度剪枝。但是，经过简短的重新训练后，趋势就会逆转。
准确性基准测试
我们使用以下参数对模型进行提炼：
- 峰值学习率 = 1e-4
- 最低学习率 = 1e-5
- 40 个步骤的线性预热
- 余弦衰减计划
- 全局批量大小 = 1152
表 1 显示了 Llama-3.1-Minitron 4B 模型变体（宽度剪枝和深度剪枝）在跨多个领域的基准测试中的比较性能，与原始 Llama 3.1 8B 模型和其他类似大小的模型进行比较。
总体而言，我们再次确认了宽度剪枝策略相比遵循最佳实践的深度剪枝策略的有效性。
| 基准测试 | 射击次数 | 指标 | Llama-3.1 8B | Minitron 4B | Llama-3.1-Minitron 4B | Phi-2 27 亿 | Gemma2 2.6 B† | Qwen2-1.5 B† | |
| 宽度剪枝 | 深度剪枝 | 宽度剪枝 | |||||||
| Winogrande | 5 | acc | 0.7272 | 0.7403 以上 | 0.7214 | 0.7348 | 0.7400++ | 0.709 | 0.662 |
| arc_challenge | 25 | acc_norm | 0.5794 | 0.5085 | 0.5256 | 0.5555% | 0.6100% | 0.554 | 0.439 |
| MMLU | 5 | acc | 0.6528 | 0.5860%* | 0.5871 | 0.6053% | 0.5749 | 0.513 | 0.565 |
| 希腊 | 10 | acc_norm | 0.8180 | 0.7496 | 0.7321 | 0.7606% | 0.7524++ | 0.73 | 0.666 |
| GSM8K | 5 | acc | 0.4860 | 0.2411 | 0.1676 | 0.4124 | 0.5500%* | 0.239 | 0.585% |
| 真实 | 0 | mc2 | 0.4506 | 0.4288 | 0.3817 | 0.4289 | 0.4400++ | – | 0.459% |
| XLSum en (20%) | 3 | rougeL | 0.3005 | 0.2954% | 0.2722 | 0.2867%* | 0.0100 | – | – |
| MBPP | 0 | pass = 1 | 0.4227 | 0.2817 | 0.3067 | 0.324 | 0.700 以上 | 0.29 | 0.374%* |
| 训练令牌 | 15T | 94 亿 | 1.4 T | 3T | 7T | ||||
*最佳模型**次佳模型 – 不可用结果† 模型发行商在模型报告中报告的结果。
为了验证提炼的模型是否可以成为强指令模型，我们使用 NeMo-Aligner 微调了 Llama-3.1-Minitron 4B 模型。我们使用了 Nemotron-4 340B 的训练数据，并在 IFEval、MT-Bench、ChatRAG-Bench 和 Berkeley 函数调用排行榜（BFCL）上评估了模型，以测试指令遵循、角色扮演、RAG 和函数调用功能。我们确认，Llama-3.1-Minitron 4B 模型可以成为强指令模型，其性能优于其他基准 SLM（表 2）。
| Minitron 4B | Llama-3.1-Minitron 4B | Gemma 2B | Phi-2 27 亿 | Gemma2 2.6 B | Qwen2-1.5 B | ||
| 基准测试 | 宽度剪枝 | 深度剪枝 | 宽度剪枝 | ||||
| IFEval 测试 | 0.4484 | 0.4257 | 0.5239++ | 0.4050 | 0.4400 | 0.6451% | 0.3981 |
| MT 工作台 | 5.61 | 5.64 | 6.34%* | 5.19 | 4.29 | 7.73% | 5.22 |
| 聊天 AG† | 0.4111++ | 0.4013 | 0.4399% | 0.3331 | 0.3760 | 0.3745 | 0.2908 |
| BFCL | 0.6423 | 0.6680% | 0.6493++ | 0.700 | 0.2305 | 0.3562 | 0.3275 |
| 训练令牌 | 94 亿 | 3T | 1.4 T | 2T | 7T | ||
*最佳模型**次优模型*基于 ChatRAG 的代表性子集，而非整个基准测试。
性能基准测试
我们使用 NVIDIA TensorRT-LLM（一个用于优化 LLM 推理的开源工具包）优化了 Llama 3.1 8B 和 Llama-3.1-Minitron 4B 模型。
图 7 和图 8 显示了不同用例下不同模型在 FP8 和 FP16 精度下的每秒吞吐量请求，这些请求以输入序列长度/输出序列长度 (ISL/OSL) 组合表示，在单个 NVIDIA H100 80GB GPU 上，以 8B 模型的批量大小为 32、4B 模型的批量大小为 64，这要归功于单个 GPU 上较小的权重允许批量更大。
Llama-3.1-Minitron-4B-Depth-Base 变体的运行速度最快，平均吞吐量约为 Llama 3.1 8B 的 2.7 倍，而 Llama-3.1-Minitron-4B-Width-Base 变体的平均吞吐量约为 Llama 3.1 8B 的 1.8 倍。在 FP8 中部署所有三种模型时，与 BF16 相比，性能也提升了约 1.3 倍。
组合：对于 Llama 3.1 8B，BS = 32；对于 Llama-3.1-Minitron 4B 型号，BS = 64.1 块 H100 80GB GPU。
结束语
剪枝和传统知识提炼是一种极具成本效益的方法，可以逐步获得较小尺寸的LLM，与跨所有领域的从头开始训练相比，可以实现更高的准确性。相比于合成数据式微调或从头开始预训练，它是一种更有效、更高效的方法。
Llama-3.1-Minitron 4B 是我们首次使用先进的开源 Llama 3.1 系列。要在 NVIDIA NeMo 中使用 Llama-3.1 的 SDG 微调，请参阅 GitHub 上的/sdg-law-title-generation notebook。
有关更多信息，请参阅以下资源：
- 通过剪枝和知识提炼实现紧凑语言模型
- /NVlabs/MinitronGitHub 库
- Hugging Face 的 Llama-3.1-Minitron 模型:
致谢
如果没有 NVIDIA 的许多人的贡献，这项工作是不可能完成的。例如，核心团队：Sharath Turuvekere Sreenivas、Saurav Muralidharan、Marcin Chochowski、Raviraj Joshi；顾问：顾问：顾问：Mostofa Patwary、Mohammad Shoeybi、Bryan Catanzaro、Jan Kautz、Pavlo Molchanov；教学调整：Ameya Sunil Mahabaleshwarkar、Hayley Ross、Brandon Rowlett、Oluwatobi Olabiyi、Shizhe Diao、Yoshi Suhara；数据集：Sanjeev Satheesh、Shengyang Sun、Jiaqi Zeng、Zhilin Wang、Yi Dong、Zihan Liu、Rajarshi Roy、Wei Ping、Makesh Narsimhan Sreedhar、Oleksii Kuchaiev；TRT-LLM：TRT-LLM：TRT-LLM：Bobby Chen、James Shen；HF 支持：Ao Tang、Greg Heinrich；模型优化：Chenhan Yu；讨论和反馈：Daniel Korzekwa；博客后期准备：Vinh Nguyen、Sharath Turuvekere Sreenivas。
