Skip to content

JJ-Lin/quantize_compare

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

8 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Quantize Compare (QC)

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.txt

快速开始

1. 查看可用的 targets

from 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"]

2. 运行对比

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()

3. 导出报告

# 导出到目录
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.json

示例脚本(推荐)

Qwen-Image

python 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_analysis

SoulX-FlashTalk(含 KV cache 场景)

bash examples/run_soulx_flashtalk_qc.sh

如需开启算子级分析:

ENABLE_OPERATOR_ANALYSIS=1 bash examples/run_soulx_flashtalk_qc.sh

设备配置

模式 1: both_offload(默认,推荐大模型)

device_config = DeviceConfig.both_offload()

双方都使用 sequential_cpu_offload,适合显存有限的场景。

模式 2: fp16_resident(显存充足时)

device_config = DeviceConfig.fp16_resident()

FP16 模型常驻 GPU,只对 FP8 blocks 进行 offload,可以减少 FP16 block 的加载时间。

API 参考

compare()

主入口函数。

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 对应算子,仅衡量该算子自身误差。

KV cache 模型建议(如 SoulX-FlashTalk)

  • cache_context_mode 用于控制 FP8 block 执行时的 cache 上下文:
    • none:不传递上下文(默认)
    • cond_only:固定 cond 上下文
    • auto_cfg:按 step 内调用顺序自动 cond/uncond
  • 建议先用 nonecond_only 验证流程稳定性,再按模型行为切换 auto_cfg

list_targets()

列出 pipeline 中可用的组件。

targets_info = list_targets(pipe)
# 返回: [{name, type, has_blocks, block_count, ...}, ...]

CompareReport

对比报告类。

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

License

MIT

About

aigc模型量化对比工具,支持逐blockl可视化精度分析

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages