【免费下载链接】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点击查看免费下载导读本文聚焦 NVIDIA Model-Optimizer 中「自定义 Hugging Face 模型量化」这一主题当模型的线性层并非标准torch.nn.Linear时ModelOpt 的默认量化流程无法自动识别这些算子需要开发者按本文给出的模式实现一个 Hugging Face 插件plugin将自定义算子改写为标准nn.Linear并注册到QuantModuleRegistry随后即可走通mtq.quantize量化、校准与export_hf_checkpoint导出最终部署到 TensorRT-LLM。文中以 DBRX 这一 MoE 模型作为完整案例逐行解读官方插件源码并给出可复用的注册机制、执行流程与实战命令读者学完后即可为自己的自定义模型编写同类量化插件。为什么自定义模型需要插件从nn.Linear搜索谈起ModelOpt 通常可以直接量化来自 Hugging Face 的 PyTorch 模型。其默认行为是遍历搜索PyTorch 模型结构把torch.nn.Linear模块原位替换为量化线性模块。这一替换发生在mtq.quantize()内部——它首先调用 conversion.py 中的replace_quant_module()递归地将模型中的模块替换为量化模块再由 model_quant.py 中的quantize()完成替换与校准。问题在于如果模型碰巧没有用nn.Linear实现线性层例如把权重存储为nn.Parameter并直接做matmul默认流程就无法感知这些算子量化配置匹配不到任何weight_quantizer/input_quantizer校准后这些层保持原样。此时就需要实现一个自定义的 Hugging Face 插件把模型转换为使用nn.Linear的形式从而让 ModelOpt 的量化替换、校准与导出机制重新生效。从源码结构看这类「开箱即用」的量化替换依托于注册表QuantModuleRegistry。在 quant_module.py 中QuantModuleRegistry _DMRegistryCls(Quant, QuantModule)而QuantModule继承自DynamicModule见 dynamic.py。也就是说被注册的量化模块本质上都是「动态模块」——通过_setup()方法在转换时装配量化器TensorQuantizer并通过重写的forward在正确位置调用量化器。案例背景DBRX 的定制 MoE 线性实现文档选取的案例是databricks/dbrx-instruct。DBRX 是一个 MoE 模型其 MoE 层使用了定制实现MoE 层以DbrxExperts模块形式出现其中三个线性算子w1、v1与w2被表示为nn.Parameter而不是nn.Linear子模块线性前向通过纯matmul完成而非nn.Linear.forward。由于 ModelOpt 无法在开箱即用的条件下检测到这些线性算子因此需要实现一个 HuggingFace 插件。核心思路分四步定义一个自定义的_QuantDbrxExpertGLU保持与原模块相同的forward签名把线性算子w1、v1、w2重写为标准nn.Linear并重新实现forward方法注册新的动态模块_QuantDbrxExperts来替换transformers库modeling_dbrx.py中的DbrxExperts插件实现完成后按 hf_ptq 示例 的流程量化 DBRX 模型使用export_hf_checkpoint导出量化模型。如果 TensorRT-LLM 尚不支持该自定义模型需要在其 PyTorch 后端补充支持并适配 HF 导出器详见统一 HF 导出指南或联系我们。插件源码逐段解读以下代码片段摘自 modelopt/torch/quantization/plugins/huggingface.py该文件即 ModelOpt 内置的 HuggingFace 插件实现包含 DBRX、Falcon、各类 MoE、注意力 KV Cache 量化等大量注册逻辑。文档原文展示的是继承DynamicModule的版本当前仓库中的同类实现已演化为继承QuantModule见同文件 L786-L884但核心模式一致通过_setup()改写子模块、重写forward()、再注册进QuantModuleRegistry。1._QuantDbrxExperts路由外层的量化包装class _QuantDbrxExpertGLU(DynamicModule): def _setup(self): Modify the DbrxExpertGLU by using nn.Linear layers. dtype, device self.w1.dtype, self.w1.device def _copy_weights(modules, weights): modules.to(dtypedtype, devicedevice) for expert_idx, module in enumerate(modules): with torch.no_grad(): module.weight.copy_(weights[expert_idx].detach()) self.w1_linear nn.ModuleList( [ nn.Linear(self.hidden_size, self.ffn_hidden_size, biasFalse) for _ in range(self.moe_num_experts) ] ) _copy_weights( self.w1_linear, self.w1.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size), ) delattr(self, w1) self.v1_linear nn.ModuleList( [ nn.Linear(self.hidden_size, self.ffn_hidden_size, biasFalse) for _ in range(self.moe_num_experts) ] ) _copy_weights( self.v1_linear, self.v1.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size), ) delattr(self, v1) self.w2_linear nn.ModuleList( [ nn.Linear(self.ffn_hidden_size, self.hidden_size, biasFalse) for _ in range(self.moe_num_experts) ] ) _copy_weights( self.w2_linear, self.w2.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size).transpose( 1, 2 ), ) delattr(self, w2) def forward(self, x: torch.Tensor, expert_idx: int) - torch.Tensor: x1 self.w1_linearexpert_idx x2 self.v1_linearexpert_idx x1 self.activation_fn(x1) x1 x1 * x2 return self.w2_linearexpert_idx_setup()是DynamicModule转换时的初始化钩子这里完成了三件关键工作逐专家展开为nn.LinearDBRX 的权重是按专家堆叠的三维参数形状为(moe_num_experts, ffn_hidden_size, hidden_size)_setup为每个专家各创建一个nn.Linear(..., biasFalse)组成nn.ModuleList权重搬运_copy_weights在torch.no_grad()下把原始nn.Parameter对应切片detach()后逐份拷贝进新的nn.Linear.weight随后delattr删除原始w1/v1/w2参数避免重复计算语义保持重写的forward保持与原始DbrxExpertGLU相同的输入输出语义——先分别过w1、v1经激活函数后逐元素相乘GLU 结构再过w2输出。需要特别注意的是w2在拷贝时需要先transpose(1, 2)。这是因为nn.Linear.forward内部执行的是x W.T而原始实现是intermediate w2[i].T见 huggingface.py 中针对 transformers 5.0 布局的注释w1[i]形状为(ffn_hidden_size, hidden_size)要匹配F.linear(x, W)需存W w1[i].T。这也说明插件作者必须精确对照原始实现的前向数学形态才能保证量化前后数值语义完全等价。2._QuantDbrxExperts路由层包装class _QuantDbrxExperts(DynamicModule): def _setup(self): Modify the DbrxExpert. # No setup is needed for DbrxExpert, we only need to update DbrxExpertGLU pass # forward method copied from the original dbrx repo def forward( self, x: torch.Tensor, weights: torch.Tensor, top_weights: torch.Tensor, top_experts: torch.LongTensor, ) - torch.Tensor: bsz, q_len, hidden_size x.shape x x.view(-1, hidden_size) out torch.zeros_like(x) expert_mask nn.functional.one_hot(top_experts, num_classesself.moe_num_experts).permute( 2, 1, 0 ) for expert_idx in range(0, self.moe_num_experts): topk_idx, token_idx torch.where(expert_mask[expert_idx]) if token_idx.shape[0] 0: continue token_list token_idx.tolist() topk_list topk_idx.tolist() expert_tokens x[None, token_list].reshape(-1, hidden_size) expert_out ( self.mlp(expert_tokens, expert_idx) * top_weights[token_list, topk_list, None] ) out.index_add_(0, token_idx, expert_out) out out.reshape(bsz, q_len, hidden_size) return out_QuantDbrxExperts本身无需改写结构其forward与原始实现保持一致其作用是把「专家路由 汇聚」的外层逻辑纳入量化体系——它逐专家取出路由到的 token调用已量化的self.mlp(expert_tokens, expert_idx)再按top_weights加权并用index_add_写回。这里的self.mlp在转换后就是_QuantDbrxExpertGLU因此外层路由不变、内层专家计算被量化。当前仓库中的实现huggingface.py沿用了同一逻辑仅将参数命名对齐为top_experts/top_weights并把moe_num_experts改为num_experts属性。3. 注册进QuantModuleRegistryif transformers.models.dbrx.modeling_dbrx.DbrxExperts not in QuantModuleRegistry: QuantModuleRegistry.register( {transformers.models.dbrx.modeling_dbrx.DbrxExperts: hf.DbrxExperts} )(_QuantDbrxExperts) if transformers.models.dbrx.modeling_dbrx.DbrxExpertGLU not in QuantModuleRegistry: QuantModuleRegistry.register( {transformers.models.dbrx.modeling_dbrx.DbrxExpertGLU: hf.DbrxExpertGLU} )(_QuantDbrxExpertGLU)注册机制的原理是QuantModuleRegistry.register({原始类: 注册名})(量化包装类)把 transformers 中的原始类DbrxExperts/DbrxExpertGLU映射到量化包装类。量化时replace_quant_module()会在注册表中查找模型的模块类型命中后通过registry.convert(model)把原模块转换成动态量化模块——即「新建一个继承自包装类、动态基类与原始类按此顺序的动态类」见 dynamic.py 中_create_new_dynamic_class的实现。注册表中的条目同时会记录字符串 key如hf.DbrxExperts、hf.DbrxExpertGLU这些 key 可以直接作为 ptq 配置 中quantizer_name的通配匹配对象用于控制是否对这些层启用量化。4. 对动态加载代码的「现场注册」register_dbrx_moe_on_the_flydef register_dbrx_moe_on_the_fly(model): Register DBRX MoE modules as QUANT_MODULE. The MoE class in DBRX is transformers_modules.modeling_dbrx.DbrxExpertGLU, which loads dynamically. if type(model).__name__ in [DbrxForCausalLM]: moe_type type(model.transformer.blocks[0].ffn.experts.mlp) # Create a QuantDbrxExpertGLU class on the fly if QuantModuleRegistry.get(moe_type) is None: QuantModuleRegistry.register({moe_type: moe_type.__name__})(_QuantDbrxExpertGLU)许多 HF 模型通过trust_remote_code动态加载自定义代码transformers_modules.*其 MoE 类并不来自 transformers 主包因此静态的if ... not in QuantModuleRegistry分支无法命中。register_dbrx_moe_on_the_fly解决了这个问题它在量化真正开始时从模型实例中取出实际的moe_typemodel.transformer.blocks[0].ffn.experts.mlp的运行时类型再把它注册到QuantModuleRegistry——这就是「on the fly现场注册」。插件机制如何接入量化主流程这些register_*_on_the_fly函数并非散兵游勇而是被统一挂载到插件调度集合中。在 huggingface.py 末尾CUSTOM_MODEL_PLUGINS.update( [ register_falcon_linears_on_the_fly, register_dbrx_moe_on_the_fly, register_moe_linear_on_the_fly, register_fused_experts_on_the_fly, force_eager_experts_impl_on_the_fly, register_sparse_moe_on_the_fly, register_hf_attentions_on_the_fly, convert_hf_parallel_linears_on_the_fly, ] )CUSTOM_MODEL_PLUGINS定义在 plugins/custom.pyregister_custom_model_plugins_on_the_fly(model)会依次调用集合内每个回调。而该函数正是replace_quant_module()conversion.py的第一步量化模块替换前先运行所有自定义插件完成现场注册再执行registry.convert(model)与递归替换_replace_quant_module。因此整条链路是用户调用mtq.quantize(model, quant_cfg, forward_loop)见 model_quant.py内部触发replace_quant_module→ 先运行CUSTOM_MODEL_PLUGINS中的各注册回调包括 DBRXQuantModuleRegistry依据注册表把 DBRX 的DbrxExperts/DbrxExpertGLU原位转换为量化包装模块依据quant_cfg中的量化器配置weight_quantizer/input_quantizer通配匹配装配各TensorQuantizerforward_loop驱动校准数据收集amax等缩放因子导出export_hf_checkpoint(model, export_dir)生成统一 HF 检查点供 TensorRT-LLM / vLLM / SGLang 部署。实战量化并导出 DBRX插件实现好后量化流程与普通 HF 模型完全一致。以仓库中的 hf_ptq 示例 为参照核心代码为import modelopt.torch.quantization as mtq from transformers import AutoModelForCausalLM # 1. 加载模型若 DBRX 使用 remote code需 trust_remote_codeTrue model AutoModelForCausalLM.from_pretrained(databricks/dbrx-instruct, trust_remote_codeTrue) # 2. 准备校准数据与 forward loop典型 128-512 个样本 calib_set get_dataloader(num_samplescalib_size) def forward_loop(model): for batch in calib_set: model(batch) # 3. 原位替换为量化模块并校准 model mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, forward_loop)其中mtq.NVFP4_DEFAULT_CFG是 ModelOpt 内置的 NVFP4 量化配置。对于 MoE 模型hf_ptq README 还特别建议优先使用mtq.NVFP4_MLP_ONLY_CFG量化 MLP/MoE 层、保留注意力层、mtq.NVFP4_EXPERTS_ONLY_CFG只量化*mlp.experts*与*block_sparse_moe*专家层或mtq.NVFP4_OMLP_ONLY_CFG额外量化o_proj以在压缩率与精度间取得更好平衡。这与本文主题直接相关DBRX 这类 MoE 模型的敏感算力集中在注意力 QKV 投影保留它们不动、只量化专家层能显著改善精度。校准完成后导出统一 HF 检查点from modelopt.torch.export import export_hf_checkpoint with torch.inference_mode(): export_hf_checkpoint( model, # 量化后的模型 export_dir, # 导出文件存放目录 )导出的检查点即为 TensorRT-LLMPyTorch 与 C 后端、vLLM、SGLang 可消费的统一格式。如果 TensorRT-LLM 尚不支持 DBRX 这种自定义架构则需要在其 PyTorch 后端添加该模型支持并视需要适配 unified_export_hf.py 中的 HF 导出器——即原文档所述「add support in its PyTorch backend and adapt the HF exporter if needed」。命令行方式可参考 huggingface_example.shexport HF_PATHdatabricks/dbrx-instruct scripts/huggingface_example.sh --model $HF_PATH --quant nvfp4 --tp 1从 DBRX 到通用模式仓库内置的同类插件DBRX 插件只是 ModelOpt「自定义 HF 模型量化」模式的样板。当前仓库的 huggingface.py 中还内置了大量同构插件可作为自定义开发的横向参照插件回调适用模型解决的核心问题register_dbrx_moe_on_the_flyDBRX远程代码动态加载的 MoE GLU 层改为nn.Linearregister_falcon_linears_on_the_flyFalcon 40B 等远程代码动态加载的线性层现场注册L1602-L1612register_fused_experts_on_the_flyMixtral、Qwen2/3-MoE、DeepSeek-V3、Jamba、OLMoE 等3-D 融合专家权重gate_up_projdown_proj的通用量化包装_QuantFusedExpertsregister_sparse_moe_on_the_fly标准 TopKRouter 结构的 MoE 块通过结构特征gateexperts自动识别 MoE 块并注册_QuantSparseSequentialMoeregister_moe_linear_on_the_flyStep-3.5 / Step-3.7专家索引式 3-D 权重MoELinear展开为逐专家nn.LinearL1907-L1998register_hf_attentions_on_the_fly各类 HF 注意力KV Cache 量化所需的注意力接口打补丁transformers ≥ 4.48new_attention_interfaceconvert_hf_parallel_linears_on_the_fly张量并行切分后的 HF 模型把被 HF 动态切分为 rowwise / colwise 的nn.Linear识别为并行线性层这些插件的共同点值得总结为「自定义插件四条黄金法则」_setup()只做结构改写把非nn.Linear的算子展开/改写为标准nn.Linear并正确搬运权重必要时delattr删除原参数forward()保持数学语义不变严格对齐原始实现的前向形态注意F.linear内部转置、matmul顺序、逐专家索引等细节注册优先做静态注册兜底做现场注册模型类在 transformers 主包内时静态注册通过trust_remote_code动态加载时用register_*_on_the_fly从模型实例反查运行时类型再注册挂载进CUSTOM_MODEL_PLUGINS只有加入该集合mtq.quantize才会在替换前自动执行你的注册逻辑。注意事项与排错要点mtq.register是更轻量的通用入口若你的自定义层只是「类名不在注册表」可以直接调用 conversion.py 中的mtq.register(original_cls, quantized_cls)完成静态注册无需走插件集合插件模式适用于需要从模型实例动态反查类型、或需要复杂结构改写的场景。权重量化未命中的自检model_quant.py 中内置了一道检查当配置要求权重量化但模型里没有任何启用的weight_quantizer时会抛出RuntimeError并列出未匹配的模式与可能原因包括「unsupported custom module, e.g. a trust_remote_code MoE layout」——这正是你的插件没生效时最常见的报错检查注册是否命中即可。transformers 版本差异DBRX 的权重布局与forward语义在不同 transformers 版本间有差异例如 transformers 5.0 中DbrxExpertGLU.forward使用原始matmul权重形状为(ffn_hidden_size, hidden_size)插件实现必须针对目标版本逐一核对必要时像仓库那样在代码注释中记录每个版本差异。校准数据PTQ 通常用 128-512 个样本校准默认校准数据混合了cnn_dailymail与Nemotron-Post-Training-Dataset-v2见 hf_ptq README可通过修改calib_set替换。导出兼容性export_hf_checkpoint产出的统一 HF 检查点可直接被 TensorRT-LLM、vLLM、SGLang 消费若架构不在后端支持矩阵内需按支持矩阵与统一 HF 导出指南评估工作量。总结自定义 Hugging Face 模型量化插件的本质是「把不规则的算子形态还原为 ModelOpt 能识别的标准形态并确保它在量化替换、校准、导出三个阶段全程可追踪」。DBRX 案例展示了从nn.Parametermatmul到nn.ModuleListnn.Linear的完整改写路径以及静态注册与现场注册两种挂载方式仓库中十余个register_*_on_the_fly插件则证明这套模式已经覆盖了从 MoE 专家层到注意力 KV Cache 量化、张量并行线性层的广泛场景。对开发者而言只要遵循「改写_setup、保持forward语义、注册进QuantModuleRegistry、挂载进CUSTOM_MODEL_PLUGINS」这四步就能让任意自定义 HF 模型接入 ModelOpt 的量化-校准-导出流水线最终部署到 TensorRT-LLM 等推理框架。赞分享【免费下载链接】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点击查看免费下载相关推荐Model Optimizer 实战从 NVIDIA Hugging Face Model Hub 一键部署 FP8 量化模型到 TensorRT-LLM、vLLM 与 SGLangModel Optimizer 实战从 NVIDIA Hugging Face Model Hub 一键部署 FP8 量化模型到 TensorRT LLM、vModel-Optimizer 统一 Hugging Face 检查点导出指南从 PTQ 量化到 TensorRT-LLM / vLLM / SGLang 一键部署Model Optimizer 统一 Hugging Face 检查点导出指南从 PTQ 量化到 TensorRT LLM / vLLM / SGLang 一hugging-face-model-trainer 实战将 Hugging Face 上微调好的 LoRA 模型转换为 GGUF并部署到 Ollama、llama.cpp 与 LM Studiohugging face model trainer 实战将 Hugging Face 上微调好的 LoRA 模型转换为 GGUF并部署到 Ollama、l示例工程上一篇AMD Ryzen性能调试完全指南免费解锁CPU隐藏潜能的终极工具下一篇免费开源直播录制神器一键搞定抖音等40平台自动化录制创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
