Puzzletron算法揭秘：Model Optimizer如何为LLM/VLM做异构剪枝与NAS

Puzzletron算法揭秘Model Optimizer如何为LLM/VLM做异构剪枝与NAS【免费下载链接】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-OptimizerPuzzletron 是 NVIDIA Model Optimizer 内置的实验性模型压缩算法基于神经架构搜索NAS为 LLM/VLM 做异构剪枝它逐层评估 FFN 宽度缩减、注意力移除等候选方案再用混合整数规划MIP在内存/延迟预算下拼出精度最优的每层结构都不同的模型最后通过知识蒸馏找回损失。 它解决什么问题从一刀切到外科手术传统剪枝如 Minitron是同质剪枝——所有层统一减宽、统一删层产出的模型结构标准、部署方便但压缩一激进就容易把关键推理通路剪没。Puzzletron 反过来每一层都可以有自己的结构。有的层保留完整注意力有的层直接删掉注意力FFN 宽度从原始的 14336 到 3072 不等。这种外科手术式压缩在 30% 的激进压缩下能显著保住精度代价是产出的是异构架构需要 vLLM 的 AnyModel 后端来部署。理论关系Minitron 是 Puzzletron 的特例——任何 Minitron 能产出的架构Puzzletron 在足够大的搜索空间里也都能找到。区别在于搜索粒度和优化目标Puzzletron 可以直接以内存预算为硬约束。⚙️ 8步自动化流水线一次命令跑完 NAS整条流水线由 examples/puzzletron/main.py 驱动核心实现在 modelopt/torch/puzzletron/puzzletron_nas_plugin.py。一次完整运行包含 8 个阶段阶段内容执行方式1启动 Puzzletron 流水线-2HF 模型转换为 PuzzletronAnyModel格式单卡3剪枝打分用激活钩子评估各通道/注意力贡献多卡4执行剪枝并保存各候选检查点单卡5构建替换库每层的候选块 子块统计单卡6计算单块得分精度损失量化多卡7运行 MIP 求解器并拼装出目标模型多卡8流水线完成输出逐层架构与评估指标多卡其中第 3、4 步的打分逻辑见 modelopt/torch/puzzletron/pruning/ 目录FFN 中间维度剪枝ffn_intermediate_pruning_mixin.py、KV 头剪枝、MoE 专家移除等均以 Mixin 形式注册配合激活钩子完成重要性评分。 MIP 优化把选结构变成数学规划MIP 求解入口在 modelopt/torch/puzzletron/mip/run_puzzle.py。它拿到每个层的候选块清单及其质量/成本分数后求解一个混合整数规划问题在满足约束的前提下让总精度损失最小。支持的约束很直白配置在 YAML 的human_constraints里target_memory显存预算MiBnum_params参数量上限target_latency_seconds端到端延迟上限需开启 runtime 统计通过 vLLM 实测target_throughput、stats.has_attention等以 llama-3_1-8B_pruneffn_memory.yaml 为例只需指定三样东西mip: human_constraints: target_memory: 78_000 # 78 GiB 显存预算 pruning: intermediate_size_list: [3072, 5888, 8704, 11520] # 候选 FFN 宽度求解后的输出是逐层架构清单例如中间 12 层注意力被置为no_op直接删掉其余层保留gqa_4FFN 保持或收窄——模型因此从 113 GiB 降到 75.8 GiB。两个省时技巧--mip-only模式剪枝和打分完成后改约束比如把target_memory从 78000 改成 96000只需重跑 MIP不用重新做昂贵的打分MIP 扫描模式在配置里开启sweep.enabled: true并给出一组压缩率一次跑出完整的精度-内存权衡曲线结果导出为 CSV 并绘图即文首那张图。 实战效果Qwen3-8B 与 Llama-3.1-8B 的公开数据官方对比指南examples/pruning/minitron_vs_puzzletron/README.md在 Qwen3-8B 上的关键结论压缩目标方法蒸馏后 MMLU相对教师~14% 参数7BMinitron95.6%✅~14% 参数7BPuzzletron91.1%~38% 内存78k MiBPuzzletron74.9%✅~38% 内存78k MiBMinitron61.7%一句话决策规则温和压缩20%选 Minitron激进压缩或有硬内存/延迟预算选 Puzzletron。两者交叉点大约在 20%~38% 压缩区间可参考完整的内存扫描实验图 蒸馏不是可选项是精度恢复的必选项剪枝后的模型失忆程度与压缩强度正相关温和压缩只掉几个点但 38% 压缩下 MMLU 会掉到接近随机猜测25% 基线附近。知识蒸馏教师 logits KL 散度损失把 Qwen3-8B 剪枝模型拉回28.6 个百分点仅用 WikiText-103 跑 100 次迭代就实现。一个有意思的观察蒸馏前 Puzzletron 在 80% 内存档位领先 Minitron 8.3pp蒸馏后却被反超 3.8pp——蒸馏前的架构排名不保证蒸馏后成立官方提到的 Blockwise Local DistillationBLD正是为此准备的改进方向。蒸馏流程本身基于 Megatron-Bridge参考 examples/megatron_bridge/README.md。 部署到 vLLM 并实测加速Puzzletron 产出的是标准 HuggingFace 检查点但需要 vLLM 的AnyModel后端来加载异构架构部署步骤见 examples/puzzletron/README.md 的 Deploy compressed model in vLLM 一节另见 examples/puzzletron/evaluation/ 的评估脚本。在单张 H200 上的推理基准并发 64模型吞吐量 (tok/s)相对基线Qwen3-8B 基线218.9-Puzzletron 78k压缩38%370.569%同时别忘了压缩组合拳剪枝减架构、量化FP8/NVFP4减精度二者互补可叠加使用。️ 支持模型与上手路径已内置的模型支持见 examples/puzzletron/configs/Llama-3.1-8B、Llama-3.2-3B、Qwen2.5-7B、Qwen3-8B、Mistral-Small-24B、Nemotron-Nano-12B-v2、Nemotron-Nano-30B-A3B含 MoE 专家剪枝、GPT-OSS-20B专家移除以及 VLM 方向的 Qwen3-VL。接入新模型只需两步写一个ModelDescriptor定义权重分块正则 一个Converter生成逐层 BlockConfig完整指引在 modelopt/torch/puzzletron/anymodel/README.md。最小启动命令Llama-3.1-8B2 卡torchrun --nproc_per_node 2 examples/puzzletron/main.py \ --config examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml 延伸阅读资料路径Puzzletron 官方教程examples/puzzletron/README.mdMinitron vs Puzzletron 完整对比指南examples/pruning/minitron_vs_puzzletron/README.md算法核心插件modelopt/torch/puzzletron/puzzletron_nas_plugin.pyMIP 求解器modelopt/torch/puzzletron/mip/剪枝 Mixin 与打分modelopt/torch/puzzletron/pruning/替换库构建modelopt/torch/puzzletron/replacement_library/蒸馏结果examples/pruning/puzzletron/README.md小结Puzzletron 把剪哪些层、每层剪多狠这个架构设计问题交给了 MIP 求解器让压缩结果精确贴合你的显存或延迟预算温和压缩场景下 Minitron 仍是更简单可靠的选择两者同属 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-Optimizer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考