自定义 INT4 Block 量化：从 llm-compressor 到 vLLM 完整讲解

代码地址：llm-compressor添加自定义量化策略+vllm推理资源-CSDN下载

本文档对应仓库目录 custom_quant_scheme_develop/，面向已具备 PyTorch / Hugging Face 基础、首次接触 llm-compressor 与 vLLM compressed-tensors 的同学。示例已在真实环境跑通：128×128 block、对称 INT4、仅权重量化（W4A16）。

1. 你要解决什么问题？

大模型权重量化后，通常要经历两段流水线：

阶段	工具	目标
训练后量化 / 压缩	llm-compressor	读 FP 模型 → 按自定义算法量化 → 写出带 `quantization_config` 的 HF 目录
推理	vLLM	读量化目录 → 为每层选 Scheme → 加载权重 → `generate`

本示例的算法特点：

INT4 对称量化 ，每个 128×128 的权重块共用一个 scale（absmax / 7）。
仅量化 Linear 权重 ，不量化激活（无校准数据，走 datafree pipeline）。
磁盘格式为 compressed-tensors 的 pack-quantized（8 个 int4 打进 1 个 int32）。
vLLM 原版不支持 strategy=block + int4 的 scheme；我们通过 Monkey-patch + 自定义 Scheme 在加载时 一次性反量化 成 bf16，再用普通 F.linear 推理。

2. 目录结构与职责

复制代码

 custom_quant_scheme_develop/
 ├── environment.yaml              # Conda 环境参考（llmcompressor + vllm 0.19 + compressed-tensors）
 ├── for_llm_compressor/           # 阶段一：量化与保存
 │   ├── my_custom_scheme.py       # 核心：Modifier + block 量化算子
 │   └── run_quantization.py       # 入口：oneshot + save_pretrained
 └── for_vllm/                     # 阶段二：推理
     ├── patch.py                  # 给 vLLM 打补丁，识别 BLOCK scheme
     ├── CompressedTensorsW4A16Block.py  # 自定义 vLLM Scheme（加载后 dequant）
     └── llm_inference.py          # 入口：先 patch，再 LLM.generate

运行顺序（不可颠倒）：

在 for_llm_compressor 下执行 run_quantization.py → 得到 *-INT4-BLOCK 模型目录。
在 for_vllm 下执行 llm_inference.py，且 patch_vllm() 必须在 LLM(...) 之前。

3. 整体数据流（建议先看懂这张图）

阶段一 · run_quantization.py（内存里完成量化，再落盘）

复制代码

阶段二 · llm_inference.py（读磁盘目录，在 vLLM 里还原成可算的权重）

复制代码

两排如何衔接： 阶段一输出的 E（如 Qwen3-0.6b-INT4-BLOCK/） 是目录路径；阶段二里 LLM(model=qmodel_dir) 读的就是这个目录。patch_vllm 不读磁盘，只是在进 LLM() 前改 vLLM 的 scheme 分发逻辑。

与代码的对应关系（避免误解）：

图示步骤	实际发生的事
② oneshot	`oneshot(..., pipeline="datafree")`：无校准数据，只触发 `CALIBRATION_EPOCH_START/END`；`pre_process` 还会给模型挂上压缩版 `save_pretrained`
③ Modifier	发生在 ② 内部的 `CALIBRATION_EPOCH_END`：`compute_block_scales` → `fake_quantize` → `pack_to_int32` → 写 `weight_packed` / `weight_scale` / `quantization_scheme`，并删掉原 `weight`（不是 save 时才 pack）
④ save	`model.save_pretrained(..., save_compressed=True)`：把内存里的量化模块写入 safetensors ，并用各层 `quantization_scheme` 汇总成 `config.json` 里的 `quantization_config`
① patch	`patch.py` 替换 `CompressedTensorsConfig._get_scheme_from_parts`，让 `strategy=block` 不再 `NotImplementedError`
② LLM 读 config	从 `config.json` 解析 `quantization_config`，为每个 Linear 准备 `CompressedTensorsLinearMethod`
③ 建层 + 加载	`get_scheme` → 返回 `CompressedTensorsW4A16Block`；`create_weights` 注册空槽位；checkpoint loader 把磁盘里的 `weight_packed`、`weight_scale` 拷入（QKV 时走 `BlockScaleParameter` 修正 shard）
④ 加载后反量化	`process_weights_after_loading`：`unpack_from_int32` + 按 block 展开 scale → bf16 `layer.weight` ，再删掉 packed 参数；不是 forward 里每次反量化
⑤ generate	`scheme.apply_weights` → `F.linear(x, layer.weight, bias)`

4. 阶段一：llm-compressor（`for_llm_compressor/`）

4.1 为什么用自定义 Modifier？

llm-compressor 内置的 QuantizationModifier + Preset（如 W4A16）适合框架已经实现 的量化策略。本示例的 128×128 block INT4 不在 stock 路径里，需要自己完成：

分块、求 scale、fake-quant
用 pack_to_int32 写出 weight_packed
为每层挂上 QuantizationScheme（strategy=BLOCK），以便 save_pretrained 写出正确的 config.json

因此本仓库只讲 Modifier 路线 ：算法与磁盘布局都在 my_custom_scheme.py 里写清楚，并与 vLLM 侧字段命名对齐。

4.2 `run_quantization.py` 在做什么？

复制代码

 recipe = MyCustomInt4WeightOnlyModifier(
     targets=["Linear"],
     ignore=["lm_head"],
     block_size=(128, 128),
 )
 oneshot(model=model, recipe=recipe, pipeline="datafree")
 model.save_pretrained(SAVE_DIR, save_compressed=True)

要点：

pipeline="datafree" ：权重量化不需要校准数据；框架只会发 CALIBRATION_EPOCH_START/END 事件（见下文）。
ignore=["lm_head"]：输出头常保持高精度，与官方示例一致。
save_compressed=True ：走 modify_save_pretrained 包装的保存逻辑，与 quantization_config 一致（模型在 oneshot 的 pre_process 里已被 patch）。

运行前请修改 MODEL_ID 为本机可访问的基座模型路径。

4.3 Modifier 生命周期（`MyCustomInt4WeightOnlyModifier`）

内置 Modifier 的惯例是 在 on_event 里响应校准事件 ，而不是只依赖基类的 BATCH_START/END：

钩子	本示例行为
`on_initialize`	`match_named_modules` 检查至少有一层 `Linear`；清理旧 buffer；记录 `block_size`
`on_event` + `CALIBRATION_EPOCH_START`	`on_start`（标记 started）
`on_event` + `CALIBRATION_EPOCH_END`	`_quantize_targets`：真正量化所有层
`on_finalize`	清空 `_quantized` 集合

易错点： 若只监听 BATCH_END 而不监听 CALIBRATION_EPOCH_END，在 datafree 下 一层都不会被量化。

4.4 量化算法（`compute_block_scales` / `fake_quantize_block_int4`）

权重布局为 PyTorch nn.Linear 约定：weight.shape == [out_features, in_features] == [N, K]。

对每个 128×128 块：

若 N 或 K 不能整除 128，先在右下角 zero-pad （只影响块边界，反量化时按 weight_shape 裁回逻辑形状）。
scale = 块内 absmax / 7（对称 INT4 最大正值为 7）。
q = round(w / scale) 再 clamp 到 [-8, 7]；反量化 dq = q * scale。

weight_scale 的形状为 [ceil(N/128), ceil(K/128)] ，dtype 存为 bf16。

4.5 单层写回（`_quantize_linear`）------与 vLLM 对齐的关键

每层完成后：

字段	含义
`weight_packed`	`pack_to_int32(int4_codes, num_bits=4)`，形状 `[N_pad, K_pad/8]`，必须用官方 API
`weight_scale`	2D block scale，bf16
`weight_shape`	`[n, k]` 逻辑形状（int64），供解压
`quantization_scheme`	`strategy=BLOCK`，`block_structure=[128,128]`
`quantization_status`	`COMPRESSED`

并删除原来的 weight Parameter（避免与 packed 格式冲突）：

复制代码

 update_offload_parameter(module, "weight_packed", packed.contiguous())
 if "weight" in module._parameters:
     module._parameters.pop("weight")

注意：

量化完成后不要在同一进程里对原模型做 model.generate()（已无合法 weight）。
__all__ 里写的是 MyCustomWeightOnlyModifier，类名实际是 MyCustomInt4WeightOnlyModifier，导入时以类名为准。

4.6 保存到磁盘后应看到什么？

打开 SAVE_DIR：

config.json ：quantization_config 中 weights.strategy 为 "block"，block_structure 为 [128, 128]，format 为 pack-quantized。
*.safetensors ：线性层有 weight_packed 、weight_scale 、weight_shape （不再有 FP 的 weight）。

5. 阶段二：vLLM（`for_vllm/`）

5.1 为什么需要 `patch.py`？

vLLM 0.19 的 CompressedTensorsConfig._get_scheme_from_parts 对 INT4 + pack-quantized 主要匹配 GROUP/CHANNEL（WNA16 / Marlin） ，不包含 BLOCK。否则会：

复制代码

 NotImplementedError: No compressed-tensors compatible scheme was found.

patch_vllm() 在 类方法 上替换 _get_scheme_from_parts：若检测到 BLOCK + int4 + pack-quantized + 无激活量化，则返回我们的 CompressedTensorsW4A16Block，否则调用原始逻辑。

使用方式（必须同一进程、且在加载模型前）：

复制代码

 from patch import patch_vllm
 patch_vllm()
 
 from vllm import LLM
 llm = LLM(model=qmodel_dir, ...)

不能先在一个终端 vllm serve，再在另一个终端单独 import patch。

5.2 `CompressedTensorsW4A16Block` 三部分

（1）`create_weights`：声明 vLLM 参数槽位

weight_packed：PackedvLLMParameter，packed_factor=8，K 维除以 8。
weight_scale：自定义 BlockScaleParameter ，形状 (out/128, in/128)。
weight_shape：长度 2 的 int64 向量。

分区约束：out % 128 == 0 且 in % 128 == 0（对 TP 分片后的 per-partition 尺寸）。

（2）`BlockScaleParameter`：QKV / 列并行加载

vLLM 合并 QKV 时，对 scale 的 shard_offset / shard_size 按行块缩放。因 block scale 的第一维是 out/128 而非 out，需要在 load_merged_column_weight / load_qkv_weight 里 除以 block_h，否则 scale 会加载错位。这是本示例里最容易踩的坑之一。

（3）`process_weights_after_loading`：一次性反量化

思路：不在每次 forward 里 unpack，加载完成后做一次：

用 weight_packed.shape 反推逻辑形状：[N, K_packed*8]（不要依赖会被 QKV 覆盖的 weight_shape buffer）。
unpack_from_int32(..., num_bits=4, shape=orig_shape, packed_dim=1)。
将 2D scale 按 block 扩回 与权重同 shape，做 w * scale。
设 layer.weight 为 bf16 Parameter，删除 weight_packed / weight_scale。

之后 apply_weights 就是标准 F.linear。代价是显存接近 FP16 权重、推理无 INT4 加速；优点是实现简单、数值易与 compressor 侧对齐验证。

5.3 `llm_inference.py`

patch_vllm()
LLM(model=qmodel_dir, ...) --- 将 qmodel_dir 改为阶段一输出目录名
SamplingParams + generate

调试建议：设置 VLLM_LOGGING_LEVEL=DEBUG，搜索日志 Using scheme: CompressedTensorsW4A16Block。

6. 端到端操作清单（给学生照着做）

环境

复制代码

 # 参考 environment.yaml 创建环境，或确保已安装：
 # llmcompressor, compressed-tensors, transformers, vllm==0.19.x, torch

步骤 A：量化

复制代码

 cd custom_quant_scheme_develop/for_llm_compressor
 # 编辑 run_quantization.py 中的 MODEL_ID
 python run_quantization.py

检查：控制台有 INT4 [128,128] block quantization 进度条；输出目录存在且 config.json 含 block 量化配置。

步骤 B：推理

复制代码

 cd custom_quant_scheme_develop/for_vllm
 # 编辑 llm_inference.py 中的 qmodel_dir
 python llm_inference.py

检查：启动时打印 [vllm-block-quant] BLOCK quantization support enabled.；能正常 generate 且无 NotImplementedError。

7. 常见问题（FAQ）

现象	可能原因
oneshot 后层数不变 / 无 scale	未监听 `CALIBRATION_EPOCH_END` 或 `targets` 未匹配
save 后没有 `weight_packed`	未设 `quantization_scheme` / `save_compressed=False`
vLLM `No compressed-tensors compatible scheme`	未 `patch_vllm()`，或 config 中 strategy/format 与 patch 条件不一致
推理乱码、scale 对不上	pack 未用 `pack_to_int32`；或 QKV scale 未用 `BlockScaleParameter` 修正 shard
`strategy` 比较失败	config 里多为字符串 `"block"`，patch 里用 `QuantizationStrategy.BLOCK`；若失效可改为 `in ("block", QuantizationStrategy.BLOCK)`
量化后 HF `generate` 报错	已删除 `weight`，应在 vLLM 或另写 decompress 脚本验证

8. 扩展方向（作业 / 进阶）

不在 vLLM 里全量 dequant：实现 fused INT4 block GEMM（工作量大，需 CUDA kernel 或第三方库）。
改 block 为 64×64 ：同时改 Modifier 的 block_size 与 scheme 的 block_structure，并确认各层 N,K 可被整除或接受 padding。
数值回归 ：单层导出 weight_packed / scale，用 compressed_tensors 的 PackedQuantizationCompressor.decompress 与 Modifier 侧 fake_quantize 对比 max error。
合并为包 ：pip install -e . 后统一入口 serve_block.py（内部先 patch_vllm() 再调 vLLM CLI）。

9. 核心文件速查

文件	一句话
`my_custom_scheme.py`	block INT4 算法 + Modifier 生命周期 + 写 `weight_packed`
`run_quantization.py`	oneshot + 保存量化模型
`patch.py`	让 vLLM 认识 BLOCK int4 scheme
`CompressedTensorsW4A16Block.py`	加载权重、QKV scale、一次性反量化、linear 推理
`llm_inference.py`	patch 后 `LLM.generate` 演示

10. 推荐阅读顺序

本文 §2--§3（目录与数据流）
my_custom_scheme.py：先看 on_event，再看 _quantize_linear
run_quantization.py
patch.py → CompressedTensorsW4A16Block.py → llm_inference.py
官方文档：Adding a New Modifier

文档版本与示例代码 custom_quant_scheme_develop 同步；若你修改了 block_size 或保存目录命名，请同步更新 run_quantization.py / llm_inference.py 中的路径。

后续升级版本可能改为 Preset + QuantizationModifier 简化 compressor 入口；当前仅覆盖 Modifier 实现，不包含该路径。

自定义 INT4 Block 量化：从 llm-compressor 到 vLLM 完整讲解

1. 你要解决什么问题？

2. 目录结构与职责

3. 整体数据流（建议先看懂这张图）

3. 整体数据流（建议先看懂这张图）

4. 阶段一：llm-compressor（for_llm_compressor/）

4.1 为什么用自定义 Modifier？

4.2 run_quantization.py 在做什么？

4.3 Modifier 生命周期（MyCustomInt4WeightOnlyModifier）

4.4 量化算法（compute_block_scales / fake_quantize_block_int4）

4.5 单层写回（_quantize_linear）------与 vLLM 对齐的关键

4.6 保存到磁盘后应看到什么？

5. 阶段二：vLLM（for_vllm/）

5.1 为什么需要 patch.py？

5.2 CompressedTensorsW4A16Block 三部分

（1）create_weights：声明 vLLM 参数槽位

（2）BlockScaleParameter：QKV / 列并行加载

（3）process_weights_after_loading：一次性反量化

5.3 llm_inference.py

6. 端到端操作清单（给学生照着做）

环境

步骤 A：量化

步骤 B：推理

7. 常见问题（FAQ）

8. 扩展方向（作业 / 进阶）

9. 核心文件速查

10. 推荐阅读顺序

4. 阶段一：llm-compressor（`for_llm_compressor/`）

4.2 `run_quantization.py` 在做什么？

4.3 Modifier 生命周期（`MyCustomInt4WeightOnlyModifier`）

4.4 量化算法（`compute_block_scales` / `fake_quantize_block_int4`）

4.5 单层写回（`_quantize_linear`）------与 vLLM 对齐的关键

5. 阶段二：vLLM（`for_vllm/`）

5.1 为什么需要 `patch.py`？

5.2 `CompressedTensorsW4A16Block` 三部分

（1）`create_weights`：声明 vLLM 参数槽位

（2）`BlockScaleParameter`：QKV / 列并行加载

（3）`process_weights_after_loading`：一次性反量化

5.3 `llm_inference.py`