知识蒸馏实战：Model Optimizer教会小模型模仿大模型的完整流程

知识蒸馏实战Model Optimizer教会小模型模仿大模型的完整流程【免费下载链接】Model-OptimizerA unified library of SOTA model optimization techniques like quantization, distillation, pruning, neural architecture search, speculative decoding, etc. It compresses deep learning models for downstream deployment frameworks like TensorRT-LLM, TensorRT, vLLM, etc. to optimize inference speed.项目地址: https://gitcode.com/GitHub_Trending/te/Model-OptimizerModel Optimizer是 NVIDIA 开源的统一模型优化工具库内置量化、剪枝、知识蒸馏、神经架构搜索等 SOTA 技术可将大模型压缩后部署到 TensorRT-LLM、TensorRT、vLLM 等推理框架。本文以知识蒸馏为主线带你在最小改动下完成大模型教小模型的完整流程从教师/学生模型准备、KD 损失计算到训练与导出并附上官方现成示例脚本 examples/llm_distill/main.py 的实战解读 为什么选择 Model Optimizer 做知识蒸馏知识蒸馏Knowledge Distillation让小模型模仿大模型的行为分布用更低的推理成本获得接近大模型的效果。它适合三类场景场景说明 模型瘦身用剪枝后的小学生模型恢复大教师模型的精度⚡ 替代从零训练从已有大模型蒸馏 微调比从头训练更快 模块替换用高效模块替换模型中的某一层靠蒸馏重新融合回整体Model Optimizer 的蒸馏 API 位于 modelopt/torch/distill/把教师前向传播、KD 损失计算、损失平衡等细节全部封装进元模型你的训练脚本通常只需加一行代码。官方指南见 docs/source/guides/4_distillation.rst。知识蒸馏快速开始安装与模型准备安装蒸馏依赖HuggingFace 模型路线只需安装带hf附加依赖的版本再装上示例的依赖文件pip install -U nvidia-modelopt[hf] pip install -r requirements.txt示例级依赖清单见 examples/llm_distill/requirements.txt。教师与学生模型怎么选官方端到端示例的思路是用已指令微调的Llama-3.2-3B-Instruct当教师蒸馏出Llama-3.2-1B学生模型数据选用 smol-smoltalk-Interaction-SFT 对话集。8 张 RTX 6000 上开启 FSDP 即可装下师生双模型完整说明见 examples/llm_distill/README.md。已验证兼容的模型包括Llama 3/4、Qwen 2/3、Gemma 2/3、Phi 3、Nemotron、Mamba 等——不限于此清单只要师生都是 PyTorch/HF 模型即可。核心工作流程convert 转换、训练、export 导出Model Optimizer 的通用蒸馏 API 分三步核心实现在 modelopt/torch/distill/distillation.py转换convertmtd.convert()把师生打包成DistillationModel元模型屏蔽两者交互细节训练用元模型替代原模型跑原有训练脚本损失计算只需多调用一次compute_kd_loss()导出export训练结束用mtd.export()摘掉教师还原出纯净的学生模型。import modelopt.torch.distill as mtd config { teacher_model: teacher_model, criterion: mtd.LogitsDistillationLoss(), loss_balancer: mtd.StaticLossBalancer(), } distill_model mtd.convert(student_model, mode[(kd_loss, config)]) # 训练时kd_loss distill_model.compute_kd_loss(student_loss) exported_student mtd.export(distill_model)配置项教师模型、criterion、loss_balancer由 modelopt/torch/distill/config.py 中的KDLossConfig校验。criterion还支持按层映射例如{(classifier, layers.18): mtd.LogitsDistillationLoss()}对学生和教师的任意中间层输出做逐层蒸馏 KDTrainerHuggingFace 用户的最小改动方案如果你熟悉 HFTrainer可以直接用 modelopt/torch/distill/plugins/huggingface.py 里的KDTrainer——它是Trainer的无缝替换品内部自动处理教师前向传播与 KD 损失计算学生模型始终保持原生 HF 格式from modelopt.torch.distill.plugins.huggingface import KDTrainer class KDSFTTrainer(KDTrainer, SFTTrainer): # 与普通 SFT 训练器组合 pass trainer KDSFTTrainer(student, training_args, distill_args{teacher_model: teacher}, train_datasettrain_ds, eval_dataseteval_ds) trainer.train()跑通官方 LLM 蒸馏示例只需一条命令8 卡 FSDP2配置见 examples/llm_distill/accelerate_config/fsdp2.yamlaccelerate launch --config-file ./accelerate_config/fsdp2.yaml main.py \ --teacher_name_or_path meta-llama/Llama-3.2-3B-Instruct \ --student_name_or_path meta-llama/Llama-3.2-1B \ --output_dir ./llama3.2-distill --max_steps 200由于学生从未被包装成DistillationModel训练结束后trainer.save_model()直接以原始 HF 格式保存无需额外导出步骤 蒸馏损失函数怎么选内置损失函数定义在 modelopt/torch/distill/losses.py开箱即用的有三种LogitsDistillationLoss对输出 logits 计算 KL 散度支持temperature参数软化分布、突出暗知识是 LLM 蒸馏的默认选择MFTLossMinifinetuning在教师分布上做阈值校正适合小数据微调场景——让小数据集提升能力的同时不冲掉教师的通用知识MGDLoss面向二维卷积输出的 Masked Generative Distillation多用于视觉生成任务。多个层对同时蒸馏时返回的损失字典会由Loss Balancer归并为单一标量官方提供静态加权求和的StaticLossBalancer也可按接口自定义动态权重策略。进阶玩法量化感知蒸馏QAD恢复精度蒸馏不只用于以大打小。Model Optimizer 的 Megatron-Bridge 蒸馏脚本 examples/megatron_bridge/distill.py 支持QADQuantization-Aware Distillation先对大模型做激进量化如 W4A4 NVFP4再用 BF16 原始模型当教师做蒸馏把量化损失的精度蒸回来。官方 Qwen3.6-35B-A3B 教程中W4A4 量化模型经 500 步 QAD 后IFBench 从 −2.6 分恢复到基本无损吞吐相比 BF16 提升 12%~30%权重体积缩小 3.1 倍常见问题 FAQ问KDTrainer 和 mtd.convert() 我该用哪个HuggingFace 生态、只做 logits 级蒸馏 → 选 KDTrainer改动最小需要隐状态/中间层蒸馏 → 用mtd.convert()DistillationModel。问分布式训练要注意什么KDTrainer 开启 FSDP 时要求 FSDP2不支持 FSDP1HF 默认的 DataParallel 会破坏师生前向请使用 FSDP2、DeepSpeed 或 DDP。问保存的检查点里会混入教师模型吗KDLossConfig默认expose_minimal_state_dictTrue检查点只保存学生权重避免把教师重复存储使用 FSDP 时需设为False。问训练完怎么部署蒸馏出的学生模型就是标准 HF 模型可直接走 Model Optimizer 的量化/导出流程部署到 vLLM、TensorRT-LLM实现小 快 省显存的落地。小结用 Model Optimizer 做知识蒸馏的路径非常清晰选师生模型 → KDTrainer/mtd.convert 接入损失 → 正常训练 → 导出学生。两条路线HF 插件与通用 API覆盖了从快速验证到逐层深度蒸馏的全部需求配合 QAD 还能把蒸馏用在量化精度的恢复上。完整概念说明可阅读 docs/source/guides/4_distillation.rst动手代码看 examples/llm_distill/ 目录 【免费下载链接】Model-OptimizerA unified library of SOTA model optimization techniques like quantization, distillation, pruning, neural architecture search, speculative decoding, etc. It compresses deep learning models for downstream deployment frameworks like TensorRT-LLM, TensorRT, vLLM, etc. to optimize inference speed.项目地址: https://gitcode.com/GitHub_Trending/te/Model-Optimizer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考