论文学习——How to Prune and Distill Llama-3.1 8B to an NVIDIA Llama-3.1-Minitron 4B Model

引言

• 压缩模型，知道我的模型是怎么生成的就行了，权当是了解，并不是直接使用这种压缩技术去生成压缩模型！
• 文章连接

正文

Abstract

• 目前大语言模型在各个任务中都表现出色，但是步数需要消耗大量的资源，所以开发小型语言模型就十分重要了。小型语言模型在特定任务中有足够的能力，并且部署成本更低。
• 英伟达通过结构化权重剪枝和知识蒸馏，从初始大语言模型逐步获得了更小的语言模型。
  • Minitron 4B和Minitron 8B这两个模型就是对Nemontron系列中较大模型15B模型的剪枝获得的小型模型。
• 剪枝和蒸馏的带来的好处
  • 在多方面显著提升语言模型的表现和资源效率
  • 提升MMLU分数
  • 减少了训练所需的数据量和计算成本
• 本文主要贡献如下
  • 提出了一套结构化的压缩方法，并应用在Llama 3.1 8B上，生成Llama-3.1-Minitron 4B上，效果和正常的Llama-3.1-8B的效果是一致的。

Pruning and distillation剪枝和蒸馏

• Pruning剪枝
  • 通过删除层或者移除神经元、注意力头和嵌入通道，使模型更小、更精简。
  • 剪枝常常伴随着模型的在训练，恢复模型的准确性。
• Distillation蒸馏
  • 是一种将大而复杂的模型的知识转移到更小、更简单的学生模型中的技术
  • 目标是保存原有模型的预测能力的同时，提高运行速度并减少资源消耗。
    Classical knowledge distillation vs. SDG finetuning传统的知识蒸馏 vs SDG微调
• 目前来说主要有两种类型的蒸馏
• SDG finetuning微调
  • 利用更大的教师模型生成的合成数据进一步微调较小的预训练学生模型
  • 学生模型，仅仅模仿教师模型最终预测生成的token
• 经典知识蒸馏
  • 学生模型模仿教师模型在训练数据集上的logits和其他中间状态，不仅仅是学习要预测的token。
  • 该方法提供了更加丰富的标签，即使是相同的数据，梯度中包含了更加丰富的反馈，提升了训练的准确性和效率。
  • 需要框架支持，因为中间的logits没有办法保存，太大了。

Pruning and distillation procedure剪枝和蒸馏过程

• 我们提出的方法是将剪枝和传统的知识蒸馏结合，是一种资源高效的再训练技术。具体技术如下
  • 重要性排名并剪枝
    • 我们从一个15B大小的模型开始，先估计每一个组件的重要性，然后对其进行排名，目标是将这个15B的模型，剪裁到8B大小。
  • 教师模型蒸馏
    • 我们以原型模型作为教师，以剪枝之后的模型作为学生，然后进行轻量的蒸馏和再训练过程。
    • 将15B大小的模型剪裁为8B
  • 蒸馏和剪枝混合缩小为4B
    • 最后，再以8B模型作为起点，进一步剪裁和蒸馏为一个更小的4B模型
      [图片]

Importance analysis重要性分析

• 该方法提出了一种基于激活的单次重要性估计策略，只需前向传播便可估计模型组件的重要性，从而实现更简便、低成本的剪枝。实验证明，单次估计已足够有效，迭代估计不再需要。
  • 使用一个小型的校准数据集，仅通过前向传播计算所有被考虑的轴的敏感信息
    • 深度、神经元、头和嵌入通道
      Retraining with classical knowledge distillation使用传统的知识蒸馏进行再训练
• 下图展示了将一个N层的模型，从教师模型中蒸馏到学生模型中的过程。
  [图片]

Pruning and distillation best practices剪枝和蒸馏的最佳实践

• Sizing尺寸
  • 首先训练最大模型，再通过剪枝和蒸馏获得较小模型
• Pruning剪枝
  • 优先使用宽度剪枝和单次重要性估计
• Retraining再训练
  • 在再训练中应仅用蒸馏损失，并根据深度变化决定是否使用中间状态和嵌入蒸馏

Llama-3.1-Minitron：putting best practices to work

• Meta推出了强大的Llama-3.1系列的模型，效果和很多闭源的模型相当，有405B、70B和8B三个不同的版本。在本文中，我们使用Llama3.1进行实验，将之从8B大小的模型变成4B大小的模型。实验内容包括
  • 监视模型微调
  • 仅仅深度剪枝
  • 仅仅宽度剪枝
  • 准确性基准测试
  • 性能基准测试

Teacher fine-tuning教师微调

• 为了校正模型训练数据集与原始数据集之间的分布偏移，我们首先在新的数据集（940亿个tokens）上微调了未剪枝的8B模型。实验表明，如果不进行分布偏移校正，教师模型在蒸馏过程中提供的指导效果较差。

Depth-only pruning仅仅深度剪枝

• 从8B变成4B，我们减去了16层，去除了50%的参数。

• 首先，通过移除模型的特定层或连续层组，并观察语言模型损失（LM loss）或下游任务的准确性降低情况，评估每层或层组的重要性。

  • 这里是直接通过评估去除当前层的模型的性能指标进行评判

• 图5展示了在验证集上移除1、2、8或16层后的LM损失值。例如，图中红色曲线在第16层表示移除前16层后的LM损失
  

• LM损失并不是直接和下游任务的性能是相关的，下图展示了每个剪枝模型在Winogrande任务上的准确性。

  • 结果表明，去除第16到第31层（其中第31层为倒数第二层）效果最佳，剪枝模型的5-shot准确性显著高于随机水平（0.5）。因此，我们采用此结果，移除了第16至第31层。
    

Width-only pruning宽度剪枝

• 我们沿宽度轴剪枝了Llama 3.1 8B的嵌入（隐藏）和MLP中间维度。具体地，我们使用之前提到的基于激活的策略计算了每个注意力头、嵌入通道和MLP隐藏维度的重要性分数。重要性估计后，我们进行了以下操作：
  • 将MLP中间维度从14336剪枝到9216。
  • 将隐藏尺寸从4096剪枝到3072。
  • 重新训练了注意力头数和层数。
• 需要指出的是，宽度剪枝后的初始LM损失高于深度剪枝，但经过短暂的再训练后，趋势逆转。

Accuracy benchmarks基准测试的准确性

• 我们使用以下参数对模型进行了蒸馏：

  • 峰值学习率=1e-4
  • 最小学习率=1e-5
  • 线性预热40步
  • 余弦衰减调度
  • 全局批量大小=1152

• 表1展示了Llama-3.1-Minitron 4B模型的不同变体（宽度剪枝和深度剪枝）与原始Llama 3.1 8B模型及同类大小的其他模型在各个基准测试中的性能对比。

• 从下图，进一步确定了，使用宽度剪枝的策略是由于基于深度剪枝的策略
  

• 为了验证蒸馏后的模型是否能够成为强大的指令模型，我们使用NeMo-Aligner对Llama-3.1-Minitron 4B模型进行了微调。训练数据来自Nemotron-4 340B，模型在IFEval、MT-Bench、ChatRAG-Bench和Berkeley Function Calling Leaderboard (BFCL)上进行了测试，以评估其指令跟随、角色扮演、RAG（检索增强生成）和函数调用的能力。

• 我们确认Llama-3.1-Minitron 4B模型在这些方面表现优异，超过了其他基准小型语言模型（表2）。

  • NeMo-Aligner：NVIDIA的一个工具，用于对模型进行指令跟随和对齐微调，使其更好地理解和执行指令。
  • IFEval：测试指令跟随能力的基准。
  • MT-Bench：用于评估模型的多任务处理能力。
  • ChatRAG-Bench：评估模型在检索增强生成（RAG）场景下的性能。
  • Berkeley Function Calling Leaderboard (BFCL)：一个评估模型函数调用和API对接能力的基准。
    

补充——各个任务的说明

• Winogrande
  • 用于测试模型的常识推理能力，通常包含复杂的填空题，要求理解上下文以选择正确答案。
• ARC Challenge：
  • 学术考试问题挑战，测试模型在科学常识、逻辑推理等方面的能力。
• MMLU（Massive Multitask Language Understanding）：
  • 跨学科的多任务语言理解测试，评估模型在各种知识领域的表现。
• HellaSwag：
  • 多选测试，涉及常识推理，通常包含较复杂的情景描述。
• GSM8k：
  • 数学推理数据集，专注于数学和算术题的解决。
• TruthfulQA：
  • 测试模型是否能够生成符合事实的回答，主要是关于常识和事实性问题。
• XLSum en：
  • 多语言文本摘要任务，评估模型的摘要生成能力。
• MBPP（Multiple-Choice Programming Problems）：
  • 编程题测试，评估模型的代码生成和编程理解能力。

补充，目前仅仅使用到语言分析的相关内容，这里列举一些各有关语言内容的相关检测

• CoLA（Corpus of Linguistic Acceptability）：主要用于检测句子是否符合语言学上的语法规范。
• GEC（Grammatical Error Correction）数据集：如FCE、Lang-8、NUCLE等，这些数据集包含标注的语法错误，可用于训练和评估语法纠正模型。
• Winograd Schema Challenge：部分场景下适合指称消解检测，判断代词指代是否合理。
  Performance benchmarks基准性能==》缩小模型的同时增加吞吐量
• 我们使用NVIDIA的TensorRT-LLM对Llama 3.1 8B和Llama-3.1-Minitron 4B模型进行了优化，这是一款开源工具包，专为大型语言模型推理优化。
• 图7和图8显示了在不同精度（FP8和FP16）下，8B和4B模型在不同任务中每秒的吞吐请求数，任务使用不同的输入序列长度/输出序列长度（ISL/OSL）组合，8B模型批量大小为32，4B模型为64，利用较小权重实现了更大的批次，测试在一块NVIDIA H100 80GB 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还提升了所有三个模型约1.3倍的性能，相较于BF16。
  

Conclusion总结

• 剪枝和经典知识蒸馏是一种非常高效的方法，能够逐步获得更小的LLM，在各个领域上实现优于从零开始训练的精度。相比于合成数据微调或从头预训练，这种方法更有效且更节省数据。
• Llama-3.1-Minitron 4B是我们基于最新的开源Llama 3.1系列的首个作品。想了解如何在NVIDIA NeMo中对Llama-3.1进行SDG微调，请参阅GitHub上的/sdg-law-title-generation notebook。
• Llama-3.1-width基于宽度的剪裁模型连接
• Llama-3.1-width基于深度的剪裁模型连接

总结

• 通过这个文章，我也就只是了解了我所使用的Minitron是如何通过大语言模型进行压缩的，如果要掌握大模型压缩技术，从硬件条件来说，根本不可能，目前所拥有的条件，只能跑一个4B的模型，所以暂时不用想了。而且从时间上来说，也不可能，目前就是从压缩转向直接使用压缩模型，同时涉及prompt，是压缩模型的检测效果和常规模型的检测效果相同。