NVIDIA 中国开发者日活动 中国・苏州 | 2025 年 11 月 14 日 了解详情
开发与优化

使用 NVIDIA TensorRT Model Optimizer 剪枝和蒸 LLM

大语言模型(LLM)在自然语言处理(NLP)任务,如代码生成、推理和数学计算等方面,展现出卓越的性能,树立了新的标杆。然而,这些模型的部署通常需要大量计算资源,成本较高,因此人们日益关注能够在显著降低资源消耗的同时,仍保持强劲性能的小语言模型(SLM)

NVIDIA 的研究人员和工程师展示了一种结合结构化权重剪枝与知识蒸馏的方法。知识蒸馏是一种高效策略,能够在几乎不损失模型性能的前提下,将大型模型压缩为更小、更高效的版本。更多详细信息请参阅通过剪枝和知识蒸馏实现的紧凑型语言模型

本文将介绍模型剪枝与知识蒸馏的基本原理及其工作机制,并讲解如何利用 NVIDIA TensorRT Model Optimizer 将这些优化技术便捷地应用于您自己的模型,以获得卓越的性能表现。

什么是模型剪枝? 

剪枝是一种模型优化技术,旨在利用训练过程中神经网络常见的过参数化现象。这类模型通常具备足够的能力来学习复杂的特征并实现稳定的收敛。剪枝通过系统性地识别并移除已训练模型中不重要的参数,例如权重、神经元甚至整个层,从而减少模型的冗余,提升效率。

该过程通常能够去除模型的大量权重,同时对准确率的影响极小,从而直接实现模型的精简,提升推理速度并降低计算成本。类似于园艺师修剪树木以促进其健康生长,模型剪枝也能使模型变得更小、更高效。

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

深度剪枝会移除神经网络中的整个层,从而降低网络的深度与复杂性(图 1)。

This figure illustrates depth-pruning in a neural network. Since all the neurons in the first hidden layer are identified as unimportant or redundant, the entire layer is targeted for removal, which will significantly reduce the network's overall depth and complexity.
图1:深度剪枝的神经网络能够有效降低整体深度与复杂性。

宽度剪枝能够移除模型中的单个神经元、注意力头或嵌入通道等内部结构,从而缩减模型的宽度(图 2)。

The original neural network, with a highlighted neuron identified as redundant. Its subsequent removal will reduce the width of this layer, creating a smaller, more computationally efficient model.
图2:通过剪枝不重要的神经元来缩减图层宽度

其核心理念是识别并移除对大语言模型整体性能影响最小的组件。可通过多种方法评估各组成部分的重要性,例如:

  • 幅度剪枝:将绝对值较小的权重置为接近零的值。
  • 基于激活的剪枝: 利用校准数据集,通过分析模型各部分的激活情况来评估其重要性。
  • 结构剪枝: 移除完整的结构单元,例如网络层或注意力头。

研究表明,宽度剪枝通常比深度剪枝在准确性上更具优势,而深度剪枝则在相同参数量的情况下,往往能更显著地降低推理延迟。选择深度剪枝、宽度剪枝或二者结合的方式,应基于对准确性和延迟之间权衡的综合考量。更多细节可参考LLM 剪枝和蒸馏实践:Minitron 方法

什么是知识蒸馏? 

知识蒸馏是一种模型压缩技术,旨在将知识从较大的“教师”模型迁移至更小、更高效的“学生”模型(图 3)。该方法的目标是构建一个结构紧凑的模型,在保持较高性能的同时,降低资源消耗,从而更便于实际部署。

This diagram shows the successful outcome of knowledge distillation by comparing the teacher network to the smaller, trained student network. The student model, despite being more compact, produces an output probability vector that closely mimics the teacher's vector.
图3。经过训练的学生模型与教师模型在知识蒸馏中的输出

知识蒸馏通过训练一个紧凑的学生模型来模仿更大规模的教师模型,其方法不仅依赖于硬标签,还利用教师模型输出的软标签进行指导。这种方式能够传递更丰富、更具泛化性的行为特征,使学生模型在保持接近教师模型准确率的同时,实现更高的运行效率。

两种常见的蒸馏方法——基于响应和基于特征——在将知识从教师模型传递给学生模型时具有不同的特点。

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

基于响应的知识蒸馏通过训练学生模型拟合教师模型的输出概率(而不仅仅是标签),将教师模型的知识迁移给学生模型。这些软目标能够反映类别之间的相似性关系,例如,“猫”与“老虎”的语义距离比与“汽车”更近。学生模型通过最小化与软目标之间的KL散度来实现对齐。

该方法易于实现,无需访问教师模型的内部结构,且在分类任务中表现出良好的效果。在实际应用中,通常将蒸馏损失与标准交叉熵损失结合使用,并通过调整损失权重来平衡模型的稳定性和预测准确性,从而获得能够保留教师模型大部分性能的紧凑型学生模型。

This figure illustrates the training loop of a student neural network using response-based distillation. The student network's softened output probabilities, its "soft predictions," are compared against the teacher model's "soft targets." The resulting loss is then used to update the student network's weights via backpropagation, as indicated by the "Update Weights" arrow.
图4:学生通过输出结果对比,从教师的软目标中进行学习

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

基于特征的知识蒸馏通过传递教师模型的中间表示,如隐藏激活或特征图,来指导学生模型学习相似的内部结构,而不仅限于输出结果的模仿。在训练过程中,教师网络的特定层与学生网络的对应层进行配对和对齐;当两者维度不一致时,通常引入投影层以实现匹配。

与基于响应的蒸馏方法相比,这种更深层次的层级监控能够提供更丰富的学习信号,并已在视觉领域(如 CNN 特征图)和自然语言处理领域(如 Transformer 的隐藏状态与注意力机制)中得到验证。由于该方法依赖于模型内部激活,因此需要访问教师模型的中间层,同时需谨慎选择匹配的网络层以及权衡任务损失的权重,以兼顾训练的稳定性与最终性能。

This figure illustrates the training loop of a student neural network using output-based distillation. The student network's hidden layer output feature map is compared against the teacher model's hidden layer’s feature map using a loss function. This is then used to update the student network's weights through backpropagation, as indicated by the "Update Weights" arrow.
图5。学生模型通过对比教师模型隐藏层的特征图进行学习

剪枝与知识蒸馏共同构成了高效的模型压缩流程,能够生成非常适合在生产环境和边缘设备中部署的小型语言模型。TensorRT 模型优化器可简化这些技术的大规模应用,将先进的大语言模型转化为可高效部署且成本优化的解决方案。

如何使用 TensorRT Model Optimizer 对模型进行剪枝

本节将介绍如何使用 TensorRT Model Optimizer 构建优化工作流,涵盖数据集准备、在 WikiText 数据集上对教师模型进行微调,以及应用剪枝与知识蒸馏技术,从 Qwen3-8B 模型生成一个 6B 参数的轻量化模型。更多详细信息,请参考支持 NeMo 2.0 框架的 Qwen3-8B 剪枝与蒸馏 Notebook

在进行剪枝和蒸馏之前,需先将 Hugging Face 模型转换为 NVIDIA NeMo 检查点格式,并对数据集完成预处理。具体操作步骤请参考模型转换与数据准备的相关说明

在此,我们将演示如何结合使用深度剪枝和宽度剪枝方法进行模型剪枝。所提供的脚本可在 NVIDIA NeMo 框架 容器 nvcr.io/nvidia/nemo:25.09 中运行。

如何通过深度剪枝模型来构建学生模型

初始方法采用自动方式筛选出最佳的 24 层,并利用包含 1024 个样本的小型校准数据集,将 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)从 12288 减少到 9216,嵌入层(hidden_size)从 4096 减少到 3584,从而得到一个 6B 规模的模型。

可以根据需要减少注意力头(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 模型优化器进行蒸馏

该示例结合知识蒸馏技术,利用 Model Optimizer 和 NeMo 2.0 框架对 Qwen3 模型进行深度和宽度剪枝。

在将知识从教师模型蒸馏到深度剪枝的学生模型时,学生模型的路径为 Qwen3-8B-nemo-depth-pruned。该路径对应深度剪枝步骤的输出结果,具体细节可参考 NeMo 蒸馏 notebook

以下提供了该过程的脚本,展示了如何在单节点 8-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

在将教师模型的知识蒸馏到宽度剪枝模型时,student_model_path 模型将作为由 NeMo 剪枝 Notebook 中的宽度剪枝步骤生成的 Qwen3-8B-nemo-width-pruned 模型。更多详细信息请参考 NeMo 蒸馏 Notebook

以下提供了该过程的脚本,展示了如何在单节点 8-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 Framework 的蒸馏文档。这些资源将有助于您轻松启用蒸馏功能,并将其顺利集成到工作流程中。

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

实验结果表明,经过模型优化器(Model Optimizer)对剪枝与蒸馏技术的优化,Qwen3 8B 在深度剪枝后得到的 6B 模型,相比原生的 Qwen3 4B 模型,推理速度提升了 30%,同时在 MMLU(大规模多任务语言理解)基准测试中表现更优。该剪枝方法将模型从原有的 36 层压缩至 24 层,并利用单张 NVIDIA H100 80GB HBM3 显卡成功生成了 6B 规模的模型。

使用从 nvidia/ClimbMix 预训练数据集处理得到的 OptimalScale/ClimbMix 数据,基于 Qwen3-8B 提取 Pruned 模型。该实验采用 25% 的数据,约为 900 亿个 token。模型蒸馏过程耗时 8 小时,使用了 96 个节点,每个节点配备 8 块 NVIDIA H100 GPU,总计消耗约 6000 GPU 小时。

Plot showing that Qwen3-8B-DPruned-6B model is 30% faster than the Qwen3-4B model, and also performs better on the MMLU benchmark (72.5 versus 70).
图6显示,Qwen3深度剪枝6B模型在速度和准确率上均优于4B模型,准确率接近8B模型,同时运行速度更快。

与对应的 4B 模型相比,6B 剪枝模型在性能上实现了显著提升。值得注意的是,经过剪枝的 60 亿参数模型推理速度提升了 30%,显著增强了处理各类计算任务的效率。为公平比较吞吐量,所有模型均采用 Model Optimizer 量化至 FP8 精度,并通过 TensorRT-LLM 进行推理运行。

除了速度优势外,60亿参数的剪枝模型还展现出极高的准确性,这在其MMLU基准测试的优异表现中得到了印证。该模型取得了72.5分的成绩,超过了40亿参数模型的70.0分,显示出其在各类语言相关任务中更强的理解能力。

速度与准确性的双重提升,使包含60亿参数的剪枝模型成为更可靠、更高效的解决方案,适用于对处理速度和结果质量要求较高的应用场景。

经过剪枝的模型基于预训练数据集进行知识提炼,因此属于基础模型变体。对于此类基础模型,我们仅在基础模型的基准测试(例如 MMLU)上进行模型间的比较。实际上,在将这些模型应用于推理任务时,通常还需在其基础上进一步开展后训练。

开始学习模型剪枝与知识蒸馏

剪枝和知识蒸馏是一种经济高效的方法,能够逐步压缩大语言模型,在多个领域中达到甚至超越基准准确率,且通常在数据利用效率上优于合成数据微调或从头开始的完整预训练。

准备开始了吗?请查看支持 Qwen3 8B 剪枝与知识蒸馏 Notebook。访问 NVIDIA/TensorRT-Model-Optimizer GitHub 仓库,深入了解剪枝蒸馏技术。如需了解更多关于使用 TensorRT Model Optimizer 进行模型优化的方法,可参考有关后训练量化量化感知训练以及预测解码的相关文章。

 

标签