极速搞定TRT大模型剪枝蒸馏：性能提升30%！

大模型（LLMs）在编码、推理、数学等自然语言处理（NLP）领域里，能力确实是杠杠的，把行业门槛一下子拉高了。但咱们做跨境实战的都清楚，把这些大家伙落地应用，那可真是个资源“无底洞”，投入巨大。所以，小模型（SLMs）这两年越来越受关注，它们能在保证不错性能的前提下，大大降低成本，这不正是我们追求的降本增效嘛。

新媒网跨境获悉，英伟达（NVIDIA）的研究人员和工程师们，就展示了一套“组合拳”——把结构化权重剪枝（structured weight pruning）和知识蒸馏（knowledge distillation）结合起来。这招厉害了，能把大模型“瘦身”，变成更小、更高效的版本，而且效果损失还不明显。想深入了解，可以看看他们那篇《通过剪枝和知识蒸馏实现紧凑型语言模型》的文章。今天，我就带大家一块儿聊聊模型剪枝和知识蒸馏到底是怎么回事，它们是怎么运作的，以及咱们怎么利用NVIDIA的TensorRT模型优化器，把这些技术轻松应用到自己的模型上，让你的跨境业务AI跑得又快又省！

一、模型剪枝，给AI“减肥瘦身”

想象一下，咱们训练一个神经网络，就像给它准备了一大桌子菜，生怕它吃不饱、学不好。结果呢，模型学完之后，很多“参数”其实是“多余”的，就像菜没吃完一样。剪枝（Pruning）就是一种模型优化技术，它专门干这事儿：系统性地找出并移除那些不那么重要的参数，比如权重、神经元，甚至整个层。

这个过程，往往能在几乎不影响模型准确率的前提下，砍掉大量的权重，直接让模型更“苗条”，推理速度更快，计算成本也更低。这跟咱们园丁给树修剪枝叶，为了让树长得更健康、更旺盛是一个道理。通过剪枝，模型变得更小、更高效。

剪枝主要有两种方式，咱们一起来看看：

• 深度剪枝（Depth Pruning）
  深度剪枝就是直接把神经网络的整层都砍掉，这样一来，模型的整体深度和复杂性就大大降低了。

  图1. 对神经网络进行深度剪枝，减少整体深度和复杂性。

• 宽度剪枝（Width Pruning）
  宽度剪枝则是针对模型内部结构，比如单个神经元、注意力头（attention heads）或者嵌入通道（embedding channels），把它们“瘦身”，让模型的宽度变窄。

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

核心思想，就是找出并移除那些对LLM整体性能贡献最小的部分。那么，我们怎么判断哪个部分不重要呢？通常有几种方法：

• 幅度剪枝（Magnitude pruning）：简单粗暴，把那些绝对值小的权重直接设为接近零。
• 基于激活的剪枝（Activation-based pruning）：用校准数据集来评估模型不同部分的重要性，看它们在激活时的表现。
• 结构化剪枝（Structural pruning）：直接移除整个结构，比如一整层或者多个注意力头。

实践中，外媒研究表明，宽度剪枝通常能达到更好的准确率，而深度剪枝在参数量相同的情况下，往往能更大幅度地降低推理延迟。具体选择哪种，或者两者结合，得看咱们对准确率和速度的侧重。

二、知识蒸馏，让“小徒弟”学到“师傅”的精髓

知识蒸馏（Knowledge Distillation），就像咱们常说的“师傅带徒弟”。它是一种模型压缩技术，把一个大模型（“师傅”模型）的“知识”传授给一个更小、更高效的模型（“学生”模型）。目标很明确：就是让这个小巧玲珑的“学生”模型，也能拥有“师傅”模型那样的高性能，这样咱们在实际部署的时候，就能用更少的资源搞定事情。

图3. 知识蒸馏训练后的“学生”和“师傅”模型输出对比。

知识蒸馏训练“学生”模型，不是只让它照着“正确答案”学，而是让它跟着“师傅”的指引学。这样，“学生”能学到更丰富、更具泛化能力的行为模式，虽然它更紧凑，但准确率却能接近“师傅”。

蒸馏的方式主要有两种：基于响应的蒸馏和基于特征的蒸馏。

1. 基于响应的知识蒸馏

这种方式是让“学生”模型去模仿“师傅”模型的“软输出概率”，而不是仅仅匹配硬标签。啥叫“软输出概率”？比如，对一张图，“师傅”模型可能说这是“猫”的概率是90%，“老虎”的概率是8%，而“汽车”的概率只有2%。这些软概率就包含了更丰富的类别间相似度信息，比如“猫”和“老虎”比“汽车”更接近。学生模型会通过KL散度（KL divergence）等方法，优化自己的输出，让它尽量跟“师傅”的软概率对齐。

这种方法实现起来简单，也不需要访问“师傅”模型内部的特征，对分类任务非常有效。实际操作中，我们常常会把蒸馏损失和基于真实标签的标准交叉熵损失结合起来，调整好损失权重，这样既能保证训练稳定，又能让“学生”模型在紧凑的同时，尽可能地保留“师傅”的准确率。

图4. “学生”通过对比输出与“师傅”的软目标学习。

2. 基于特征的知识蒸馏

基于特征的知识蒸馏就更深入一层了。它关注的是“师傅”模型的中间表示，比如隐藏层的激活值或特征图，以此来引导“学生”模型学习类似的内部结构。在训练时，我们会把“师傅”和“学生”模型中选定的层进行配对对齐。如果维度不一样，通常还会用到投影层来适配。

这种更深层次的、层级的监督，比基于响应的蒸馏提供了更丰富、更具体的信号。它在视觉（比如CNN的特征图）和自然语言处理（比如Transformer的隐藏状态和注意力）领域都非常有效。因为它依赖内部激活，所以需要能访问“师傅”模型的中间层，而且要仔细选择和加权这些层，同时兼顾标准任务损失，才能在稳定性和准确性之间找到最佳平衡。

图5. “学生”通过对比隐藏层特征图向“师傅”学习。

剪枝和蒸馏，这对“黄金搭档”，为模型压缩提供了一整套强大的解决方案，让咱们能轻松创建出适合生产环境和边缘设备部署的小模型。而TensorRT模型优化器，就是咱们把这些先进技术规模化应用的好帮手，能把那些尖端的大模型，变成真正能落地的、性价比高的解决方案。

三、手把手教你：用TensorRT模型优化器剪枝模型

好啦，理论讲完了，咱们来点儿实战的！这一节，我将手把手带大家搭建一个使用TensorRT模型优化器的流程。这里会用到WikiText数据集来准备数据，然后用Qwen3-8B模型做“师傅”，通过剪枝和蒸馏，最终得到一个大约60亿参数的“学生”模型。更详细的操作，可以参考《使用NeMo 2.0框架对Qwen3-8B进行剪枝和蒸馏》这个笔记本。

在正式剪枝和蒸馏之前，咱们得先做两件事：把Hugging Face格式的模型转换成NVIDIA NeMo的检查点格式，并且对数据集进行预处理。这些基础工作是万万不能马虎的，具体怎么操作，请参考模型转换和数据准备步骤的详细说明。

接下来，咱们就演示怎么进行深度剪枝和宽度剪枝。下面的脚本，都可以在NVIDIA NeMo框架容器 nvcr.io/nvidia/nemo:25.09 里运行。工欲善其事必先利其器，务必确保你的环境配置是最新且正确的。

1. 深度剪枝，创建“学生”模型

咱们的第一步，是把Qwen3 8B模型从36层“瘦身”到24层（大约60亿参数）。这里会用一个包含1024个样本的小型校准数据集，让系统自动选出要保留的最佳24层。

下面是剪枝的脚本，这里展示的是一个使用两块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

这里简单解释下几个关键参数：

• --nproc_per_node 2：表示每台机器使用2个进程，通常对应2个GPU。
• --devices 2：指定使用的GPU数量。
• --pp_size 2：流水线并行（Pipeline Parallelism）的大小，这里是2，意味着模型被分成2个部分，由2个GPU处理。
• --restore_path Qwen3-8B-nemo：指定“师傅”模型的路径，也就是我们剪枝的起点。
• --save_path Qwen3-8B-nemo-depth-pruned：剪枝后“学生”模型的保存路径。
• --seq_length 4096：序列长度，也就是模型能处理的文本最大长度。
• --num_train_samples 1024：用于校准的训练样本数量。
• --mbs 4：每个GPU上的微批次大小（micro batch size）。
• --data_paths wikitext-data/wikitext-train_text_document：校准数据集的路径。

2. 宽度剪枝，创建“学生”模型

除了深度剪枝，咱们还可以选择宽度剪枝来缩小模型。这种方法是通过缩减模型的关键架构组件实现的：比如把多层感知机（MLP）的中间层大小（ffn_hidden_size）从12288减少到9216，以及把嵌入层（hidden_size）从4096减少到3584，这样也能得到一个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

这里需要额外注意的参数是：

• --target_ffn_hidden_size 9216：指定目标MLP中间层的大小。
• --target_hidden_size 3584：指定目标嵌入层的大小。

剪枝之后，咱们的模型确实是“瘦身”了，内存占用小了，推理速度也可能更快。但通常，这只是第一步。为了弥补剪枝过程中可能损失的准确率，并确保模型在特定任务上依然表现出色，我们还需要进行微调或者再训练。这时候，知识蒸馏就该登场了！

四、TensorRT模型优化器，如何进行知识蒸馏

接下来，咱们就来看看怎么用Model Optimizer和NeMo 2.0框架，对刚才剪枝过的Qwen3模型进行知识蒸馏。当咱们对深度剪枝后的模型进行知识蒸馏时，“学生”模型的路径就是Qwen3-8B-nemo-depth-pruned，这正是深度剪枝那一步的输出。具体细节，大家可以参考NeMo的蒸馏笔记本。

下面是蒸馏的脚本，这里演示的是一个单节点八块GPU张量并行（Tensor Parallel）设置的例子。实际操作中，为了更快地完成训练，我们通常会推荐使用多节点训练。当然，咱们也得提前规划好计算资源和成本投入，毕竟“磨刀不误砍柴工”，但“烧钱”也不能过度。

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

这里几个重要的参数：

• --nproc_per_node 8：每台机器使用8个进程，对应8个GPU。
• --name Qwen3-8B-nemo-depth-pruned-distill：本次蒸馏任务的名称。
• --devices 8：使用的GPU数量。
• --tp_size 8：张量并行（Tensor Parallelism）的大小，这里是8，意味着模型参数在8个GPU上分布。
• --model_path Qwen3-8B-nemo-depth-pruned：指定“学生”模型的路径，也就是剪枝后的模型。
• --teacher_path Qwen3-8B-nemo：指定“师傅”模型的路径，用来传授知识。
• --max_steps 40：最大训练步数。
• --lr 1e-4：学习率。
• --gbs 768：全局批次大小（global batch size）。

同样地，如果咱们要从“师傅”模型蒸馏知识到宽度剪枝后的模型，“学生”模型的路径就会是Qwen3-8B-nemo-width-pruned，这个是宽度剪枝那一步的产物。更具体的说明，也可以参考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框架的蒸馏文档。这些资源能帮助你轻松地将蒸馏技术应用到你的工作流中。记住，技术更新迭代快，请务必关注官方文档，确保使用的工具和方法都是最新、最稳定的版本。

五、剪枝和蒸馏，对模型性能影响几何？

新媒网跨境认为，这组数据无疑为我们指明了方向。使用Model Optimizer对Qwen3 8B模型进行剪枝和蒸馏的实验结果表明，Qwen3深度剪枝后的60亿参数模型，比Qwen3的40亿参数模型快了30%！而且，在MMLU（大规模多任务语言理解）基准测试中，它的表现也更好。

具体来说，深度剪枝把模型从36层减少到24层，得到了一个60亿参数的模型，整个过程用了一块NVIDIA H100 80 GB HBM3显卡。这个剪枝后的模型，是从Qwen3-8B模型蒸馏而来的，使用了OptimalScale/ClimbMix数据，该数据来源于nvidia/ClimbMix预训练数据集。实验中只用了25%的数据，大约是900亿个Token。整个蒸馏过程用了96个节点，每个节点有8块NVIDIA H100 GPU，总计消耗了6000个GPU小时，可见高效训练的重要性。

图6. Qwen3深度剪枝后的60亿参数模型在速度和准确性上均优于40亿参数模型，且接近80亿参数模型的准确性，运行速度更快。

这个60亿参数的剪枝模型，相比40亿参数版本，性能提升非常显著。值得注意的是，它实现了30%的速度提升，这意味着在各种计算任务中，它的效率都更高。在吞吐量对比方面，所有模型都通过Model Optimizer量化为FP8精度，并使用TensorRT-LLM运行，确保了公平性。

除了速度优势，这个60亿参数的剪枝模型还展现出了卓越的准确性。在MMLU基准测试中，它取得了72.5分的好成绩，超过了40亿参数模型的70.0分。这说明它在广泛的语言相关任务中，理解和处理能力更强。

速度和准确性的双重提升，使得这个60亿参数的剪枝模型，成为了一个更强大、更高效的解决方案，非常适合那些既要求快速处理又追求高质量结果的应用场景。

需要提醒大家的是，这次剪枝后的模型是在预训练数据集上进行蒸馏的，所以它是一个基础版本。我们只在MMLU这样的基础模型基准上对所有模型进行了比较。如果咱们想在实际的推理任务中使用这些模型，通常还需要对模型进行后训练（post-training），比如针对特定领域的数据进行微调，才能发挥出它最大的潜能。

六、是时候行动了：开启你的剪枝和知识蒸馏之旅

剪枝和知识蒸馏，是咱们逐步“瘦身”大模型，同时保持甚至超越基线准确率的两种极其划算的方法。它们通常比合成数据微调或完全预训练更节省数据。想亲手试一试了吗？那就赶紧去看看《使用NeMo 2.0框架对Qwen3 8B进行剪枝和蒸馏》这个笔记本吧。想了解更多关于剪枝和蒸馏的信息，也可以访问NVIDIA/TensorRT-Model-Optimizer的GitHub仓库。

关于使用TensorRT模型优化器进行模型优化的其他技术，比如训练后量化（post-training quantization）、量化感知训练（quantization-aware training）和推测解码（speculative decoding），大家可以看看相关的文章。

希望今天的分享能为大家在跨境业务中应用大模型带来新的思路和启发，祝愿大家都能把AI这把“利剑”用得炉火纯青，为自己的事业添砖加瓦！

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/trt-llm-pruning-distillation-30-percent-speed-up.html