NVFP4 量化流程说明

NVFP4 量化与导出流程说明本文以 pi05LLMGemma 2B 解码器为主例说明 NVFP4 从标定、set_dynamic_quant、ONNX 导出到_nvfp4_post_processing的完整链路。目录总览端到端流程阶段 1quantize_model标定阶段 2set_dynamic_quant动态激活配置阶段 3exportONNX 导出阶段 4_nvfp4_post_processing后处理完整数据流一层 Linear 从标定到 TRT一、总览端到端流程以LLM.quantize()为例NVFP4 完整流水线如下阶段 4NVFP4 后处理is_nvfp4_quantized(quant_cfg)?_nvfp4_post_processing(llm.onnx)fp4qdq_to_2dq权重 TRT_FP4QDQ → 双 DQ FP4 打包外部权重文件 onnx_model.data阶段 3ONNX 导出self.export(export_dir, dynamoFalse)torch.onnx.export(self, ...)激活侧TRT_FP4DynamicQuantize权重侧TRT_FP4QDQ 占位阶段 2导出前配置set_dynamic_quant(self, dynamic_quant)激活 quantizer → dynamic权重 quantizer → static设置 _trt_high_precision_dtype BF16/FP16阶段 1标定PyTorch Fake Quantget_calibrate_dataset(calib_data)quantize_model(self.model, quant_cfg, loader)mtq.quantize插入 QuantLinear TensorQuantizercalibrate_loop多轮 forward 收集 amax固化 scalefake QDQ 就绪对应代码入口src/model_optimizer/models/pi05/llm.pydefquantize(self,quant_cfg,calib_data,export_dir,*,measure_quant_errorFalse):calib_dataloaderself.get_calibrate_dataset(calib_data)# AWQ 族标定可能临时转 fp32见下文quantize_model(self.model,quant_cfg,calib_dataloader,...)self.is_quantizedTruedynamic_quantself.feature_config.quantize.get(dynamic_quant,bf16)set_dynamic_quant(self,dynamic_quant)dynamobool(self.feature_config.export.get(dynamo,False))self.export(export_dir,dynamodynamo)ifis_nvfp4_quantized(quant_cfg):self._nvfp4_post_processing(f{export_dir}/llm.onnx,export_dir)一句话结论标定定 scale → 导出前配置 dynamic/static → 导出生成 TRT 自定义 Q/DQ 图 → 后处理仅压缩权重为真实 FP4激活在推理时仍走动态量化路径。二、阶段 1quantize_model标定2.1 为何 LLM 传入self.model而不是self对象含义self.modelHFGemmaModelpaligemma.get_decoder()selfLLM包装类额外做 KV cache 拼接ModelOpt 的register_hf_attentions_on_the_fly要求根模块是HF PreTrainedModel才会注册*_bmm_quantizer。若配置了FP8_KV_CFG/NVFP4_KV_CFG等 KV cache 量化必须直接量化GemmaModel。标定前向走的是qm(**batch)# GemmaModel.forward(inputs_embeds..., attention_mask..., position_ids...)不经过LLM.forward的 KV 聚合逻辑。标定数据来自Pi05LLMCalibCollector在language_model.forward上 hook 采集inputs_embeds[B, seq_len, 2048]attention_mask[B, 1, seq_len, seq_len]position_ids[B, seq_len]2.2quantize_model内部步骤入口src/model_optimizer/quantization/quantization_utils.py解析quant_cfgNVFP4 使用mtq.NVFP4_DEFAULT_CFGcfg.py里nvfp4典型特征num_bits (2, 1)→ FP4 E2M1block_sizes含scale_bits: (4, 3)→ NVFP4 block 量化is_nvfp4_quantized()通过num_bits[0] 2识别图改造每个nn.Linearq_proj、k_proj、v_proj、o_proj、MLP 等变为QuantLinearx → [input_quantizer] → x̂ (fake QDQ) → Linear(Ŵ) → y W 侧有 weight_quantizer标定循环Fake Quant对每个 calib batch 跑完整 decoder forward各层input_quantizer统计激活 amaxblock-wise各层weight_quantizer统计权重 amax误差在层间传播上一层 QDQ 输出作为下一层输入标定结束scale 固化PyTorch 推理仍是浮点 MatMul fake QDQ不是整数 GEMM。2.3 LLM 特有的 AWQ fp32 标定保护nvfp4_awq等 AWQ 族配置会临时把 decoder 转 fp32 标定避免 bf16 上 CUDA device-side assert标定完再恢复 bf16 供导出awq_fp32_calib_quant_cfg_uses_awq_family(quant_cfg)ifawq_fp32_calib:self.model.to(torch.float32)try:quantize_model(...)finally:ifawq_fp32_calib:self.model.to(saved_dtype)2.4 以q_proj为例的单层标定标定前: x [B,L,2048] → Linear(W [2048,2048]) → y 标定后 fake quant 前向: x → input_quantizer (NVFP4 block fake QDQ, 收集 amax) → weight_quantizer (NVFP4 block fake QDQ, 固化 scale) → matmul(x̂, Ŵ^T) ← 仍是 bf16 GEMM → y → 送入 k_proj 的 input_quantizer ...三、阶段 2set_dynamic_quant动态激活配置3.1 调用时机与参数在标定完成之后、export之前调用只改 ONNX 导出属性不改已固化的 scaledynamic_quantself.feature_config.quantize.get(dynamic_quant,bf16)# 默认 bf16set_dynamic_quant(self,dynamic_quant)Vit 等未走feature_config的 stage 通常写死set_dynamic_quant(self, bf16)。3.2 函数实现src/model_optimizer/utils/utils.pydefset_dynamic_quant(model:nn.Module,dtype:str)-None:formoduleinmodel.modules():ifis_nvfp4_linear(module):module.input_quantizer._trt_high_precision_dtype(Halfifdtypefp16elseBFloat16)module.input_quantizer._onnx_quantizer_typedynamicmodule.weight_quantizer._onnx_quantizer_typestatic识别 NVFP4 层is_nvfp4_linearQuantLinear且input_quantizer.block_sizes[scale_bits] (4, 3)。传入selfLLM 包装时self.modules()会遍历到self.model下所有QuantLinear因此能正确设置。3.3 三个属性的含义属性设置值作用_onnx_quantizer_typeinputdynamic激活导出为运行时动态量化_onnx_quantizer_typeweightstatic权重导出为静态 FP4 打包_trt_high_precision_dtypeBFloat16/HalfQDQ 周围 Cast 与 MatMul 的高精度 dtype3.4 导出时 dynamic vs static 在 ONNX 中的差异ModelOptexport_fp4()modelopt/torch/quantization/export_onnx.py激活dynamic— 推理时按 block 动态算 scaleifonnx_quantizer_typedynamic:sx_f32_per_tensorfloat(amax)/6.0/448.0# 用标定 amax 算 per-tensor 上界x_f4,sx_f8TRT_FP4DynamicQuantize(...)# → DequantizeLinear × 2 → BF16 激活ONNX 图里会出现激活 x (BF16) → Cast (若需要) → TRT_FP4DynamicQuantize(x, per_tensor_scale) → (x_f4, sx_f8) → DequantizeLinear(sx_f8, per_tensor_scale) → dyn_scale → DequantizeLinear(x_f4, dyn_scale, block_size) → x_dequant (BF16) → MatMul/Gemm要点标定 amax 变成per-tensor scale 因子不是把激活固定成常数每个 token/block 在推理时仍动态算 block scale。权重static— 导出占位后处理再压成 FP4returng.op(trt::TRT_FP4QDQ,inputs,block_size_iblock_size)3.5 为何 NVFP4 要「动态激活 静态权重」对象策略原因权重 Wstatic离线固定可预打包 FP4省显存/带宽激活 xdynamic分布随输入/prefix 变化大固定 scale 精度损失大Blackwell TRT 支持 runtime block quant标定阶段用 fake quant 定 amax是为了给 dynamic 路径提供per-tensor 上界推理时 block scale 仍按实际激活动态计算。3.6dynamic_quant配置项通过feature_config.quantize.dynamic_quant控制{quantize:{dynamic_quant:bf16}}{quantize:{dynamic_quant:fp16}}影响 ONNX 中 Cast 与 DQ 输出的高精度类型BF16 vs FP16需与 TensorRT engine 精度策略一致。四、阶段 3exportONNX 导出NVFP4 量化后默认dynamoFalse含 Q/DQ 的图用 legacy tracer 更稳。LLM 默认modeself_forward直接导出LLM包装模块输入: inputs_embeds [B, seq, 2048], attention_mask, position_ids 输出: past_keys, past_values, last_hidden_state导出的是已插入 QuantLinear 的量化图。每层 Linear 在 ONNX 中体现为[激活 Dynamic FP4 QDQ 子图] [权重 TRT_FP4QDQ 占位] → MatMul/Gemm若还配置了 KV cache 量化nvfp4kv_cache_qformatnvfp4attention 的*_bmm_quantizer也会出现在 ONNX 中依赖self.model作为 quantize root。五、阶段 4_nvfp4_post_processing后处理LLM 在llm.py中 override 了基类Model._nvfp4_post_processing逻辑与 Vit 等 stage 一致。5.1 触发条件ifis_nvfp4_quantized(quant_cfg):# 配置层num_bits[0] 2self._nvfp4_post_processing(onnx_path,export_dir)# 函数内部再检查ifis_fp4_quantized(self):# 模型层存在 scale_bits (4,3) 的 quantizer...5.2 步骤self.model.save_pretrained(export_dir)— 保留config.json等元数据onnx.shape_inference.infer_shapes_path(onnx_path)— 补全 tensor shapefp4qdq_to_2dq(onnx_model)— 核心转换见下清理 export 目录— 删除除.json外的所有文件onnx.save_model(..., save_as_external_dataTrue, locationonnx_model.data, convert_attributeFalse)— 重新保存 ONNX 外部权重convert_attributeFalse刻意关闭 attribute 转 external data避免 TensorRT 解析 external weights 时出现Expected size: … Actual size: 0 bytes错误。5.3 只处理权重不动激活算子后处理最终形态TRT_FP4QDQ权重fp4qdq_to_2dq替换双DequantizeLinear FP4 打包权重TRT_FP4DynamicQuantize激活保留TensorRT 运行时动态量化5.4 权重 2DQ 转换细节对每个TRT_FP4QDQ节点modelopt/onnx/quantization/qdq_utils.py读 FP16/BF16 权重 initializer计算 per-tensor scale per-block scale量化并_cast_fp4打包两个 FP4 值压进 1 byte替换为双 DQ 子图DQ1: (sw_f8_per_block, sw_f32_per_tensor) → sw_f32 DQ2: (w_f4_packed, sw_f32, block_size) → w_dequant在 MatMul/Gemm 输入前插入 Cast对齐 BF16/FP165.5 最终产物export_dir/ ├── config.json # Gemma / SigLIP config ├── llm.onnx # 2DQ 权重 DynamicQuantize 激活 └── onnx_model.data # FP4 打包权重外部存储六、完整数据流一层 Linear 从标定到 TRT以layers.0.self_attn.q_proj为例┌─────────────────────────────────────────────────────────────────┐ │ 标定 (quantize_model) │ │ calib batch → GemmaModel.forward │ │ x → input_quantizer.fake_qdq (block NVFP4, 收集 amax) │ │ W → weight_quantizer.fake_qdq (block NVFP4, 固化 scale) │ │ y matmul(x̂, Ŵ^T) │ └─────────────────────────────────────────────────────────────────┘ ↓ ┌─────────────────────────────────────────────────────────────────┐ │ set_dynamic_quant │ │ input_quantizer._onnx_quantizer_type dynamic │ │ weight_quantizer._onnx_quantizer_type static │ │ _trt_high_precision_dtype BFloat16 │ └─────────────────────────────────────────────────────────────────┘ ↓ ┌─────────────────────────────────────────────────────────────────┐ │ export → llm.onnx │ │ 激活: TRT_FP4DynamicQuantize → 2× DequantizeLinear → BF16 x │ │ 权重: TRT_FP4QDQ (占位, 仍是全精度权重 tensor) │ │ → MatMul/Gemm │ └─────────────────────────────────────────────────────────────────┘ ↓ ┌─────────────────────────────────────────────────────────────────┐ │ _nvfp4_post_processing → fp4qdq_to_2dq │ │ 权重: TRT_FP4QDQ → 双 DQ FP4 packed weights │ │ 激活: TRT_FP4DynamicQuantize 保持不变 │ │ 保存 llm.onnx onnx_model.data │ └─────────────────────────────────────────────────────────────────┘ ↓ ┌─────────────────────────────────────────────────────────────────┐ │ TensorRT build │ │ 权重: 真 FP4 GEMM读 packed weights │ │ 激活: 运行时 TRT_FP4DynamicQuantize FP4 MatMul │ └─────────────────────────────────────────────────────────────────┘