量化之从 LLM 中提取 PTQ 校准数据

什么是 PTQ

PTQ 是 Post-Training Quantization 的缩写，中文叫训练后量化。简单说，就是把已经训练好的模型参数从高精度（比如 32 位浮点数）转换成低精度（比如 8 位整数），让模型跑得更快、占内存更少。但量化需要一个前提：得知道每一层输入数据的数值范围大概是多少，才能合理地确定缩放比例和截断边界。这些用来确定范围的数据，就叫校准数据。

校准样本是一小批用来“试跑”模型的输入样本。它本身不是用户平时使用模型时输入的内容，而是挑选出来用来观察模型内部各层数值分布的数据。这里用的是 Pile validation set，这是 HuggingFace 上 mit-han-lab/pile-val-backup 这个数据集，属于英文文本语料。工具会从里面取 128 条样本，让模型跑一遍，记录每一层的输入 activation。所以校准样本可以理解为“为了量化而喂给模型的一小批输入”。

校准样本是输入文本，而真正用来确定数值范围的校准数据，是模型跑完这些样本后产生的 activation。

amct_pytorch.extract_ptq_data 是一个工具（参考cann/amct仓），用来从大语言模型里提取 PTQ 所需的校准数据。它会用校准样本跑一遍模型推理，把每一层的输入 activation 提取出来保存成 .pkl 文件。后续量化时，就用这些文件里的数据来计算 scale 和 clip 值。

activation 是模型内部某一层产生的输出张量。当它传给下一层时，就成为下一层的输入。extract_ptq_data 要捕获的是被量化层的输入，也就是上一层产生的 activation。

.pkl 是 Python pickle 格式的后缀。pickle 是 Python 内置的一种序列化方式，可以把 Python 对象直接存成二进制文件，之后再用 pickle 读回来，对象的结构和内容都能还原。这里保存的是每一层的输入 activation，本质上是张量数据，用 pickle 存比较方便，不需要额外定义复杂的文件格式。

完整命令示例

先设置环境变量：

export PATH=/usr/local/python3.12.13/bin:$PATH
source /usr/local/Ascend/cann-9.0.0/set_env.sh
export HF_ENDPOINT=https://hf-mirror.com
export HF_HUB_DISABLE_XET=1

然后运行提取命令：

MODEL_DIR=/workspace/models/Qwen3-0.6B

python3 -m amct_pytorch.extract_ptq_data \
    --model "$MODEL_DIR" \
    --model_name qwen3 \
    --seq_len 4096 \
    --granularity block \
    --device npu:0 \
    --data_dir ptq_data/qwen3/attn-linear \
    --quant_target attn-linear 2>&1

参数说明

–model 指定本地已下载好的模型目录。

–model_name 填 qwen3，用来选择正确的预处理逻辑。

–seq_len 是序列长度，填 4096，需要匹配模型的最大上下文长度。

–granularity 是提取粒度，填 block 表示按层为单位保存。

–device 填 npu:0，表示使用 NPU 0 推理。

–data_dir 是输出目录，支持相对路径，相对于当前工作目录。

–quant_target 是要提取的目标模块，可选 attn-linear 或 mlp。attn-linear 提取 Attention 中线性层的输入，mlp 提取 MLP 层的输入。也可以提取其他子模块，视具体模型和量化需求而定。

attn-linear 指的是 Attention 模块里的线性层。Transformer 的 Attention 部分通常包含几个线性变换，比如把输入映射成 query、key、value 的投影层，以及最后的输出投影层。这些层都是全连接层，也就是线性层。attn-linear 这个目标就是提取这些线性层的输入 activation。

mlp 指的是 Transformer 里的前馈网络部分，通常叫 MLP 或 FFN。它一般由两个线性层组成，中间夹一个激活函数。mlp 这个目标就是提取这部分网络的输入 activation。

这两个目标分开提取，是因为量化时不同模块的数值分布可能差别很大，需要各自独立的校准数据。

运行过程

阶段一：数据集下载

首次运行时，工具会自动下载 Pile validation set，这是 HuggingFace 上的 mit-han-lab/pile-val-backup 数据集。缓存路径在 ~/.cache/huggingface/hub/datasets–mit-han-lab–pile-val-backup/。最终大小约 450MB，解压后约 1.3GB。下载完成后会变成 3 个 Arrow 分片文件。

Arrow 是 Apache Arrow 定义的一种列式内存格式，也常用于磁盘上的数据存储。HuggingFace 的 datasets 库默认把数据集存成 Arrow 文件，按分片切分，方便并行读取。这里就是 pile-val-backup-validation-00000-of-00003.arrow 这样的三个文件。

阶段二：加载校准样本

数据集下载完毕后，日志会出现 Loaded 128 calibration samples for extract_ptq_data. 这表示从 Pile 数据集中选取了 128 条样本作为校准数据。

阶段三：模型加载与 NPU 推理

模型被加载到 CPU 内存，Qwen3-0.6B 约占用 7.6GB RSS。CPU 接近跑满，约 10 核，980% 左右。NPU AICore 使用率会在实际推理时变为非零。模型按 block 逐块送 NPU 推理，每块的输入 activation 被保存。

阶段四：输出结果

最终输出目录结构大致如下：

ptq_data/qwen3/attn-linear/
├── position_ids.pkl
├── position_embeddings.pkl
├── block_0_attn_in.pkl
├── block_1_attn_in.pkl
├── block_2_attn_in.pkl
│   ...
└── block_27_attn_in.pkl

每个 block 对应模型的一个 transformer layer。Qwen3-0.6B 有 28 层，所以是 block_0 到 block_27。每个 block 文件大小约 1GB。

关于同时提取两个目标

如果需要同时提取 attn-linear 和 mlp，需要分两次运行，每次指定不同的 --quant_target 和 --data_dir。第二次运行会跳过数据集下载阶段。(逻辑上应该支持跑一次而提取二次)