使用TensorRT模型优化器进行LLM剪枝与蒸馏

使用NVIDIA TensorRT模型优化器进行LLM剪枝与蒸馏

大型语言模型（LLMs）在代码生成、推理和数学等自然语言处理任务上设定了高标准。然而，其部署仍然需要大量资源，这促使业界对能以较低成本提供强劲性能的小型语言模型（SLMs）兴趣日增。

某中心的研究人员和工程师展示了一种将结构化权重剪枝与知识蒸馏相结合的方法，这是一种强大的策略，可在不显著损失质量的情况下将大型模型压缩为更小、高效的变体。更多细节，请参阅《通过剪枝和知识蒸馏实现紧凑语言模型》。

本文解释了模型剪枝和知识蒸馏、它们的工作原理，以及如何轻松地将它们应用到自己的模型中，以使用某中心的TensorRT模型优化器实现最佳性能。

什么是模型剪枝？

剪枝是一种模型优化技术，它利用神经网络常见的过度参数化现象。过度参数化源自训练具有足够容量的模型以学习复杂特征并确保平滑收敛。剪枝系统性地识别并移除已训练模型中的不重要参数，如权重、神经元，甚至整个层。

这个过程通常可以消除模型的大量权重，而对准确性影响最小，直接转化为模型更紧凑、推理速度加快、计算成本降低。类似于园丁修剪树木以改善其健康与生长，模型剪枝使模型变得更小、更高效。

深度剪枝和宽度剪枝是两种主要方法。

深度剪枝从神经网络中移除整个层，减少整体深度和复杂性（图1）。

图1. 深度剪枝神经网络降低了整体深度和复杂性

宽度剪枝消除内部结构，如单个神经元、注意力头或嵌入通道，缩减模型的宽度（图2）。

图2. 通过剪除不重要神经元来减少层宽度

核心思想是识别并移除对LLM整体性能贡献最小的部分。使用不同的方法来评估不同组件的重要性，例如：

• 幅度剪枝：将绝对值较小的权重设置为接近零。
• 基于激活的剪枝：使用校准数据集根据激活来估计模型不同部分的重要性。
• 结构剪枝：移除整个结构，如层或注意力头。

研究表明，宽度剪枝通常比深度剪枝获得更好的准确性，尽管深度剪枝在相同参数数量下通常更能减少推理延迟。深度剪枝、宽度剪枝或两者结合的选择应取决于准确性与延迟之间的期望平衡。更多信息，请参阅《LLM剪枝与蒸馏实践：Minitron方法》。

什么是知识蒸馏？

知识蒸馏是一种模型压缩技术，它将知识从更大的“教师”模型转移到更小、更高效的“学生”模型（图3）。目标是创建一个保持较大模型高性能的紧凑模型，使其适合在较低资源成本下部署。

图3. 知识蒸馏训练学生和教师模型输出

知识蒸馏训练一个紧凑的学生模型来模仿更大的教师模型，不是仅仅依赖硬标签，而是从教师的指导中学习。这转移了丰富、可泛化的行为，使学生模型在运行效率高得多的同时，逼近教师的准确性。

两种常见的蒸馏风格——基于响应的和基于特征的——在将知识从教师传递给学生的方式上有所不同。

什么是基于响应的知识蒸馏？

基于响应的知识蒸馏通过训练学生模型匹配教师模型的软输出概率（而非仅硬标签）来转移知识。这些软目标传达了类间相似性，例如“猫”更接近“老虎”而非“汽车”，并且使用KL散度优化学生模型以与它们对齐。

该方法实现简单，无需访问教师的内部特征，并且对于分类任务非常有效。在实践中，通常将蒸馏损失与标准交叉熵损失结合使用，并调整损失权重以平衡稳定性和保真度，从而产生能保留教师大部分准确性的紧凑模型。

图4. 学生通过输出比较从教师的软目标中学习

什么是基于特征的知识蒸馏？

基于特征的知识蒸馏转移教师的中间表示（隐藏激活或特征图），以引导学生学习相似的内部结构，而不仅仅是相似的输出。在训练期间，选定的教师和学生层被配对并对齐，当维度不同时经常使用投影层。

这种更深层的、层级监督提供了比基于响应的KD更丰富的信号，并已在视觉（例如CNN特征图）和NLP（例如Transformer隐藏状态和注意力）领域证明有效。由于它依赖于内部激活，该技术需要访问教师的中间层，并且需要仔细选择层和权重，与标准任务损失一起以平衡稳定性和准确性。

图5. 学生从教师的隐藏层特征图比较中学习

剪枝和蒸馏形成了一个强大的模型压缩流程，使得创建适合生产环境和边缘应用部署的SLMs成为可能。TensorRT模型优化器简化了大规模应用这些技术的过程，将最先进的LLMs转变为可部署的、具有成本效益的解决方案。

如何使用TensorRT模型优化器剪枝模型

本节将指导如何使用TensorRT模型优化器构建一个流程。包括数据集准备、在WikiText数据集上微调教师模型，以及应用剪枝和蒸馏技术从Qwen3-8B生成一个60亿参数的模型。更多信息，请参阅《使用NeMo 2.0框架进行Qwen3-8B剪枝与蒸馏》笔记。

在剪枝和蒸馏之前，需要将Hugging Face模型转换为某中心的NeMo检查点格式并预处理数据集。详细说明请参考模型转换和数据准备步骤。

这里，将演示如何使用深度剪枝和宽度剪枝两种方法进行剪枝。提供的脚本可以在某中心的NeMo框架容器 nvcr.io/nvidia/nemo:25.09 内运行。

如何深度剪枝模型以创建学生模型

初始方法是通过使用1,024个样本的小型校准数据集自动选择要保留的最佳24层，将Qwen3 8B模型从36层修剪到24层（约60亿参数）。

以下是此过程的脚本，展示了如何使用双GPU管道并行设置进行剪枝。

torchrun --nproc_per_node 2 /opt/NeMo/scripts/llm/gpt_prune.py \
    --devices 2 \
    --pp_size 2 \
    --restore_path Qwen3-8B-nemo \
    --legacy_ckpt \
    --save_path Qwen3-8B-nemo-depth-pruned \
    --seq_length 4096 \
    --num_train_samples 1024 \
    --mbs 4 \
    --data_paths wikitext-data/wikitext-train_text_document \
    --target_num_layers 24

如何宽度剪枝模型以创建学生模型

第二种减小模型尺寸的替代方法是宽度剪枝。这是通过缩小关键架构组件来实现的：MLP中间层（ffn_hidden_size）从12,288减少到9,216，嵌入层（hidden_size）从4,096减少到3,584，同样得到一个60亿参数的模型。

可以根据需要进一步减少注意力头的数量（num_attention_heads）和GQA查询组的数量（num_query_groups）。也可以调整层数（num_layers）以达到所需的模型尺寸。

以下是此过程的脚本，展示了如何使用双GPU管道并行设置进行剪枝。

torchrun --nproc_per_node 2 /opt/NeMo/scripts/llm/gpt_prune.py \
    --devices 2 \
    --pp_size 2 \
    --restore_path Qwen3-8B-nemo \
    --legacy_ckpt \
    --save_path Qwen3-8B-nemo-width-pruned \
    --seq_length 4096 \
    --num_train_samples 1024 \
    --mbs 4 \
    --data_paths wikitext-data/wikitext-train_text_document \
    --target_ffn_hidden_size 9216 \
    --target_hidden_size 3584

通过修剪冗余或低重要性权重，剪枝不仅缩小了模型的内存占用，还可以加速推理。然而，这个过程通常随后需要进行微调或重新训练，以恢复剪枝阶段可能损失的准确性，并确保剪枝后的模型在目标任务上保持高性能。这就是蒸馏的用武之地。

如何使用TensorRT模型优化器进行蒸馏

此示例使用模型优化器和NeMo 2.0框架，通过知识蒸馏来蒸馏Qwen3深度剪枝和宽度剪枝模型。

当从教师模型蒸馏知识到深度剪枝模型时，学生模型的路径将是 Qwen3-8B-nemo-depth-pruned。此路径对应于深度剪枝步骤的输出，如NeMo蒸馏笔记中所述。

以下是此过程的脚本，展示了如何使用单节点八GPU张量并行设置进行蒸馏。在实践中，建议使用多节点训练以获得更快的训练速度。

torchrun --nproc_per_node 8 /opt/NeMo/scripts/llm/gpt_train.py \
    --name Qwen3-8B-nemo-depth-pruned-distill \
    --devices 8 \
    --num_nodes 1 \
    --tp_size 8 \
    --model_path Qwen3-8B-nemo-depth-pruned \
    --teacher_path Qwen3-8B-nemo \
    --legacy_ckpt \
    --max_steps 40 \
    --warmup_steps 1 \
    --gbs 768 \
    --mbs 8 \
    --lr 1e-4 \
    --min_lr 1e-5 \
    --seq_length 4096 \
    --log_dir . \
    --log_interval 5 \
    --val_check_interval 5 \
    --limit_val_batches 2 \
    --data_paths wikitext-data/wikitext-train_text_document

当从教师模型蒸馏知识到宽度剪枝模型时，学生模型路径将是 Qwen3-8B-nemo-width-pruned，由NeMo剪枝笔记中的宽度剪枝步骤生成。更多细节请参阅NeMo蒸馏笔记。

以下是此过程的脚本，展示了如何使用单节点八GPU张量并行设置进行蒸馏。在实践中，建议使用多节点训练以获得更快的训练速度。

torchrun --nproc_per_node 8 /opt/NeMo/scripts/llm/gpt_train.py \
    --name Qwen3-8B-nemo-width-pruned-distill \
    --devices 8 \
    --num_nodes 1 \
    --tp_size 8 \
    --model_path Qwen3-8B-nemo-width-pruned \
    --teacher_path Qwen3-8B-nemo \
    --legacy_ckpt \
    --max_steps 40 \
    --warmup_steps 1 \
    --gbs 768 \
    --mbs 8 \
    --lr 1e-4 \
    --min_lr 1e-5 \
    --seq_length 4096 \
    --log_dir . \
    --log_interval 5 \
    --val_check_interval 5 \
    --limit_val_batches 2 \
    --data_paths wikitext-data/wikitext-train_text_document

更全面的信息，请参阅NeMo框架蒸馏文档。这些资源将帮助您轻松启用并将蒸馏集成到您的工作流程中。

剪枝和蒸馏如何影响模型性能？

使用模型优化器从Qwen3 8B进行剪枝和蒸馏的实验结果表明，Qwen3深度剪枝60亿模型比Qwen3 40亿模型快30%，并且在MMLU（大规模多任务语言理解）基准测试中也表现更好。深度剪枝将模型从36层减少到24层，得到一个60亿模型，使用了单个某中心H100 80 GB HBM3 GPU。

剪枝后的模型使用经过处理的OptimalScale/ClimbMix数据（源自nvidia/ClimbMix预训练数据集）从Qwen3-8B蒸馏而来。实验使用了25%的数据，大约900亿个token。蒸馏耗时8小时，使用了96个节点，每个节点配备八块某中心H100 GPU（总计6K GPU小时）。

图6. Qwen3深度剪枝60亿模型在速度和准确性上均优于40亿模型，并在运行速度快得多的同时接近80亿模型的准确性

与40亿模型相比，60亿剪枝模型在性能上取得了显著进步。值得注意的是，60亿剪枝模型速度提升了30%，使其在各种计算任务中效率大大提高。对于吞吐量比较，所有模型均使用模型优化器量化为FP8精度，并使用TensorRT-LLM运行。

除了速度优势，60亿剪枝模型还表现出更高的准确性，这体现在其在MMLU基准测试中获得更高的分数。其72.5的分数超过了40亿模型的70.0分，表明其在广泛的与语言相关的任务中具有更好的理解和能力。

这种速度和准确性的双重提升，使得60亿剪枝模型成为需要快速处理和高品质结果的应用程序的更强大、更有效的解决方案。

剪枝后的模型是在预训练数据集上蒸馏的，因此该模型是基础变体。由于是基础模型，仅在基础模型基准测试（如MMLU）上比较了所有模型。在实际将这些模型用于推理任务时，可能还需要对模型进行后训练。

开始使用剪枝和知识蒸馏

剪枝和知识蒸馏是高度成本效益的方法，可以逐步缩小LLMs，同时匹配或超越跨领域的基线准确性，并且它们通常比合成数据微调或完整的预训练更数据高效。

准备好开始了吗？查看《使用NeMo 2.0框架进行Qwen3 8B剪枝与蒸馏》笔记。访问某中心/TensorRT-Model-Optimizer GitHub仓库以了解有关剪枝和蒸馏的更多信息。有关使用TensorRT模型优化器进行模型优化技术的更多信息，请参阅关于训练后量化、量化感知训练和推测解码的相关文章。
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）

公众号二维码

公众号二维码