Diffusers / DiT 模型量化对比工具:一行代码对比 FP16 和 FP8 模型在 Block 级与算子级(可选)精度差异。
- 流式执行:在 FP16 block 的 forward 中嵌入 FP8 block 调用,无需存储中间激活
- Block 级对比:精确定位量化退化的来源
- 算子级对比(可选):对 Top-K block 下一层全部算子做局部隔离误差分析,并对 FFN 追加 up/down projection 专项采集
- Timestep 分析:按 denoising step 维度分析误差变化
- 峰度分析:识别量化敏感的 blocks(高峰度权重/激活)
- 灵活的设备管理:支持 sequential_cpu_offload 和 FP16 常驻 GPU 两种模式
- 可视化报告:自动生成热力图、趋势图、峰度图,并在 HTML 报告中展示算子级统计
| 模型 | 架构/场景 | 脚本 | 状态 |
|---|---|---|---|
| Qwen-Image | Diffusers DiT,FP16 vs FP8 | examples/qwen_image_example.py |
已测试通过 |
| SoulX-FlashTalk | Wan2.2 衍生(self-forcing + KV cache),常用自检 FP16 vs FP16 | examples/soulx_flashtalk_example.py / examples/run_soulx_flashtalk_qc.sh |
已测试通过 |
| FLUX.1-dev | 通用 Diffusers DiT 对比示例 | examples/flux_example.py |
提供脚本(可复验) |
cd quantize_compare
pip install -e .或直接安装依赖:
pip install -r requirements.txtfrom qc import list_targets, print_targets
# 加载 pipeline
pipe_fp16 = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.float16)
# 查看可用组件
print_targets(pipe_fp16)
# 输出:
# ┌──────────────┬────────────────────────┬─────────┬────────────┐
# │ Name │ Type │ Blocks │ Params │
# ├──────────────┼────────────────────────┼─────────┼────────────┤
# │ transformer │ FluxTransformer2DModel │ 57 │ 12.0B │
# │ vae │ AutoencoderKL │ 12 │ 84.0M │
# └──────────────┴────────────────────────┴─────────┴────────────┘
# Recommended: targets=["transformer", "vae"]from qc import compare, DeviceConfig
# 加载 FP8 量化模型
pipe_fp8 = FluxPipeline.from_pretrained("your-fp8-model", torch_dtype=torch.float8_e4m3fn)
# 运行对比
report = compare(
pipe_fp16=pipe_fp16,
pipe_fp8=pipe_fp8,
inputs={
"prompt": "A beautiful sunset over the ocean",
"num_inference_steps": 20,
"guidance_scale": 7.5,
},
targets=["transformer"],
device_config=DeviceConfig.both_offload(), # 大模型用 offload
analyze_kurtosis=True,
)
# 查看结果
report.summary()# 导出到目录
report.export("./qc_report")
# 生成的文件:
# - summary.json # 运行元信息
# - block_metrics.json # 每个 Block 的指标
# - kurtosis_analysis.json # 峰度分析结果
# - report_data.json # 报告渲染所需的完整数据
# - timestep_trend.png # 误差趋势图
# - heatmap.png # Block × Step 热力图
# - weight_kurtosis.png # 权重峰度分布图
# - activation_kurtosis.png # 激活峰度图
# - report.html # HTML 报告
# - operator_analysis # 算子级结果已写入 block_metrics.json / report_data.jsonpython examples/qwen_image_example.py \
--fp16_path /path/to/qwen-image-fp16 \
--fp8_path /path/to/qwen-image-fp8 \
--output_dir ./qwen_image_qc_report \
--enable_operator_analysisbash examples/run_soulx_flashtalk_qc.sh如需开启算子级分析:
ENABLE_OPERATOR_ANALYSIS=1 bash examples/run_soulx_flashtalk_qc.shdevice_config = DeviceConfig.both_offload()双方都使用 sequential_cpu_offload,适合显存有限的场景。
device_config = DeviceConfig.fp16_resident()FP16 模型常驻 GPU,只对 FP8 blocks 进行 offload,可以减少 FP16 block 的加载时间。
主入口函数。
report = compare(
pipe_fp16, # FP16 pipeline
pipe_fp8, # FP8 pipeline
inputs, # Pipeline 输入参数
targets=["transformer"], # 目标组件
device_config=None, # 设备配置
sample_steps=None, # 采样的 steps(None=全部)
metrics=["rel_l2", "cosine", "sqnr"], # 计算的指标
analyze_kurtosis=True, # 是否分析峰度
kurtosis_mode="per_tensor", # 峰度模式: per_tensor | per_channel
enable_3d_plots=False, # 是否采集 3D 分布图(默认关闭)
enable_operator_analysis=False, # 是否采集 Top-K block 的算子级局部误差
max_3d_blocks=10, # 3D 分布采集的 blocks 数量上限
max_3d_tokens=64, # 激活图 token 维下采样上限
max_3d_channels=64, # 激活图 channel 维下采样上限
max_3d_weight_dim=64, # 权重图维度下采样上限
fp16_model_path=None, # FP16 权重目录(用于 3D 权重直读)
fp8_model_path=None, # FP8 权重目录(用于 3D 权重直读)
allow_hf_cache=True, # 允许从 HF cache 中读取权重
cache_context_mode="none", # KV cache 上下文模式: none | cond_only | auto_cfg
seed=42, # 随机种子
output_dir=None, # 输出目录
verbose=True, # 打印运行日志
)- Top Worst Blocks (by RelL2) 使用的是每个 block 在所有采样 step 上的
rel_l2.mean(不是单一 step)。 - 如果设置了
sample_steps,Top Worst 与 3D 分布图只使用这些 steps。 enable_3d_plots=True会额外执行一次 pass 采集 3D 分布数据(仅 Top-K blocks)。enable_operator_analysis=True会在同一额外 pass 中,对 Top-K blocks 采集算子级局部隔离误差。- Top-K 数量由
max_3d_blocks控制(3D 和算子级分析共用)。
- 算子发现范围:对 Top-K block 的下一层全部算子进行匹配与分析。
- 为避免误判,采用高置信匹配策略;低置信候选会被标记为 unmatched,不进入指标统计。
- 对识别为 FFN 的分支,会额外采集并展示
up_projection/down_projection。 - 误差口径为局部隔离误差:使用同一份 FP16 算子输入调用 FP8 对应算子,仅衡量该算子自身误差。
cache_context_mode用于控制 FP8 block 执行时的 cache 上下文:none:不传递上下文(默认)cond_only:固定 cond 上下文auto_cfg:按 step 内调用顺序自动 cond/uncond
- 建议先用
none或cond_only验证流程稳定性,再按模型行为切换auto_cfg。
列出 pipeline 中可用的组件。
targets_info = list_targets(pipe)
# 返回: [{name, type, has_blocks, block_count, ...}, ...]对比报告类。
report.summary() # 打印摘要
report.top_blocks(k=10) # 获取 Top-K 最差 blocks
report.get_high_kurtosis_blocks() # 获取高峰度 blocks
report.export("./output") # 导出报告| 指标 | 公式 | 说明 |
|---|---|---|
| rel_l2 | ||a-b||_2 / ||a||_2 |
相对 L2 误差(主指标) |
| cosine | 1 - cos(a, b) |
余弦距离(方向偏差) |
| sqnr | 10*log10(Σa² / Σ(a-b)²) |
信号量化噪声比 (dB) |
| mse | mean((a-b)²) |
均方误差 |
| max_abs | max(|a-b|) |
最大绝对误差 |
峰度(Kurtosis)用于识别量化敏感的 blocks:
- 权重峰度 > 3:权重分布有重尾,量化时容易产生较大误差
- 激活峰度变化大:FP8 可能无法精确表示动态范围
# 获取高峰度 blocks
high_k_blocks = report.get_high_kurtosis_blocks(threshold=3.0)
print(f"量化敏感的 blocks: {high_k_blocks}")quantize_compare/
├── qc/
│ ├── __init__.py # 包入口
│ ├── api.py # compare() 主函数
│ ├── adapter/ # 组件/Block 发现
│ ├── align/ # Block 对齐
│ ├── streaming/ # 流式执行
│ ├── metrics/ # 指标计算
│ ├── analysis/ # Timestep/峰度分析
│ └── report/ # 报告生成
├── examples/ # 示例代码
├── tests/ # 测试
├── requirements.txt
└── README.md
MIT