LLM模型一站式量化
==========

#### 简介

本工具提供大语言模型（Large Language Model，以下简称LLM）的4bit低位量化能力，采用标准的三段式量化流程：权重量化、激活量化和量化参数提取。三段式量化流程说明如下表所示。

**表1** 大语言模型4bit低位量化三阶段流程

| 阶段 | 输入 | 输出 |
| :-- | :-- | :-- |
| 阶段一（权重量化） | 
\- 开发者提供

\- JSON数据集（text字段作为prompt）

\- HuggingFace模型路径（开发者浮点模型）

\- 量化配置及执行文件

\- config.yaml

\- run.sh

 | 

\- dopt\_config.json：开发者需完成该文件的量化策略配置后，重新执行此阶段

\- trained\_quant\_weight.pth：权重量化参数

**说明：** 开发者需根据阶段一生成的dopt\_config.json文件手动进行量化策略配置后，再次执行该阶段。

 |
| 阶段二（激活量化） | 

\- 开发者提供

\- JSON数据集（text字段作为prompt）

\- HuggingFace模型路径（开发者浮点模型）

\- 量化配置及执行文件

\- config.yaml

\- run.sh

 | trained.pth：用于阶段三加载或量化仿真验证 |
| 阶段三（量化参数提取） | 

\- 开发者提供

\- JSON数据集（text字段作为prompt）

\- HuggingFace模型路径（开发者浮点模型）

\- 量化配置及执行文件

\- config.yaml

\- run.sh

 | 

\- fake\_quant\_weight.pth：用于ONNX模型导出

\- quant\_params\_file：权重及激活量化参数

 |

#### 量化前准备工作

*   HuggingFace浮点模型
    
*   JSON格式数据集，使用“text”字段作为prompt关键字
    
*   量化配置文件[config.yaml](#configyaml示例)
    
*   [run.sh](#runsh示例)执行脚本
    

![](https://contentcenter-vali-drcn.dbankcdn.cn/pvt_2/DeveloperAlliance_scene_100_1/f9/v3/fXC7fWGRRf2sYrd3QWuQYQ/note_3.0-zh-cn.png?HW-CC-KV=V1&HW-CC-Date=20260612T073654Z&HW-CC-Expire=86400&HW-CC-Sign=9A1F3C302BAC04A3911DBF09BD93D5D45A33C9CFEB910D53B0B827987E1DEA37)

*   LLM在量化过程中使用AutoModelForCausalLM以及AutoTokenizer加载，所以开发者的HuggingFace浮点模型需要满足该加载方式约束。
    
*   run.sh执行脚本环境须与开发者推理或者训练环境保持一致，否则模型加载或推理失败。
    

#### \[h2\]config.yaml示例

config.yaml用于模型量化配置，开发者可根据实际需要进行配置。以下示例可供参考，主要参数说明请参见[config.yaml配置参数说明](#configyaml配置参数说明)。

kd:
  enable: False
  loss: mse
  micro\_batch\_size: 2
  gradient\_accumulation\_steps: 4
  weight\_decay: 0.0
  warmup\_steps: 10
  num\_epochs: 3
  learning\_rate: !!float 1e-4
  eval\_step: 1
  logging\_step: 50
  lr\_scheduler\_type: cosine
  trainable\_keys:
    - quant\_alpha
    - norm
  no\_split\_module\_classes:
    - Qwen2DecoderLayer
    - GlmDecoderLayer
    - LlamaDecoderLayer
    - HunYuanDecoderLayer
dataset:
  train\_files: dataset.json
  train\_samples: 1024
  ptq\_samples: 1024
extra\_training\_config:
  fp16: True
cutoff\_len: 128
num\_samples: 256
quant\_param\_2: False
embedding\_separate: True
lm\_head\_size:

#### \[h2\]config.yaml配置参数说明

以下为config.yaml文件的关键配置参数，具体说明如下表所示。

| 参数 | 参数 | 默认取值 | 说明 |
| :-- | :-- | :-- | :-- |
| kd：量化蒸馏相关 | enable | False | 
\- True：使用蒸馏量化策略

\- False：使用无训练量化策略，速度快，效果较稳定

 |
| kd：量化蒸馏相关 | loss | mse | 量化蒸馏损失函数，仅支持mse |
| kd：量化蒸馏相关 | micro\_batch\_size | 2 | 单卡batch |
| kd：量化蒸馏相关 | gradient\_accumulation\_steps | 4 | 梯度累计步数 |
| kd：量化蒸馏相关 | weight\_decay | 0.0 | 权重衰减系数 |
| kd：量化蒸馏相关 | warmup\_steps | 10 | 预热步数 |
| kd：量化蒸馏相关 | num\_epochs | 3 | 训练迭代次数 |
| kd：量化蒸馏相关 | learning\_rate | 1e-4 | 学习率 |
| kd：量化蒸馏相关 | eval\_step | 1 | 每1个训练步进行一次评估 |
| kd：量化蒸馏相关 | logging\_step | 50 | log打印步数 |
| kd：量化蒸馏相关 | lr\_scheduler\_type | cosine | 学习率调整策略。 |
| kd：量化蒸馏相关 | trainable\_keys | \- | 

配置可训练参数。

\- quant\_alpha：量化层的可训练参数

\- norm：layer\_norm层的可训练参数

 |
| kd：量化蒸馏相关 | no\_split\_module\_classes | \- | 

多卡切分时，选择切分粒度。

\- Qwen2DecoderLayer

\- GlmDecoderLayer

\- LlamaDecoderLayer

\- HunYuanDecoderLayer

 |
| dataset | train\_files | \- | JSON格式训练数据文件，dataset.json |
| dataset | train\_samples | 1024 | kd使能为True时， 训练集样本数，缺省默认全量数据集 |
| dataset | ptq\_samples | 1024 | PTQ优化样本数，缺省默认全量数据集 |
| dataset | extra\_training\_config | fp16 | 训练数据类型，仅支持fp16 |
| dataset | cutoff\_len | 128 | 激活量化样本序列长度。 |
| dataset | num\_samples | 256 | 激活量化校准样本数 |
| dataset | quant\_param\_2 | False | 

\- True：Kirin9020

\- False：KirinX90

 |
| dataset | embedding\_separate | True | 

\- True：单独保存为bin文件

\- False: 导出embedding的量化参数到量化文件，合并形态。

 |
| dataset | lm\_head\_size: | \- | 可指定lmhead长度，硬件对齐。 |

#### 执行三段式量化

按以下步骤执行run.sh文件，stagex为传入参数，具体可参考[run.sh示例](#runsh示例)。

1.  权重量化。
    
    1.  选配量化策略。
        
        输出：[dopt\_config.json](#dopt_configjson量化配置文件说明)，权重量化参数trained\_quant\_weight.pth【为阶段二的输入文件】。
        
        bash run.sh stage1
        
        量化策略可选配置为：
        
        Quant\_act\_weight\_eco   decode层策略
        Quant\_lm\_head          lm\_head层策略
        Quant\_Embed\_MinMax     embedding层策略
        
    2.  dopt\_config.json文件手动修改量化策略配置后，再次执行该阶段，进行权重量化。
        
         bash run.sh stage1
        
2.  激活量化。
    
    bash run.sh stage2
    
    输出：trained.pth文件，可以用作标杆和量化仿真。
    
3.  提取量化参数。
    
    bash run.sh stage3
    
    输出：fake\_quant\_weight.pth+quant\_params\_file，其中fake\_quant\_weight.pth在导出onnx模型时需要替换模型的权重为该文件中的权重。
    

#### \[h2\]run.sh示例

#!/bin/bash
set -e
qlibs='path/to/dopt\_pytorch\_py3'
export WANDB\_DISABLED=true
export HF\_DATASETS\_OFFLINE=0
export PYTHONPATH=${qlibs}:$PYTHONPATH
# 设置为cuda或npu模式，二选一
# cuda模式
export DEVICE=cuda
export CUDA\_VISIBLE\_DEVICES=1
# npu模式
# export DEVICE=npu
# export ASCEND\_RT\_VISIBLE\_DEVICES=2,3
ROOT=.
testcase='output\_dir'
RUN\_FILE=${qlibs}/dopt/dopt\_lm/opt\_main.py
output\_dir=${ROOT}/${testcase}/train\_output
mkdir -p ${output\_dir}
cp ${ROOT}/config.yaml $output\_dir
model\_path='path/to/model'
dopt\_config=./${testcase}/dopt\_config.json
quant\_stage=$1
block\_size=128
python -u \\
    ${RUN\_FILE} --model-path $model\_path \\
    --dopt-config $dopt\_config \\
    --optimize-config ${ROOT}/config.yaml \\
    --quant-stage $quant\_stage \\
    --block-size $block\_size \\
    --output-dir ${output\_dir} 2>&1 | tee ${output\_dir}/logs.log

#### \[h2\]dopt\_config.json量化配置文件说明

工具支持开发者插件式自定义LLM量化规格：

*   逐层权重量化位宽(bit)和分组大小(group\_size)。
    
*   逐层激活量化位宽(bit)。
    
*   量化策略：decoder层使用"Quant\_act\_weight\_eco"，lm\_head层使用"Quant\_lm\_head"，embedding层使用"Quant\_Embed\_MinMax"。
    

{
     "layer\_strategy": {
         "model.layers.0.self\_attn.q\_proj": {
             "type": "<class 'torch.nn.modules.linear.Linear'>",
             "quant\_strategy": "Quant\_act\_weight\_eco",
             "weight":{
                 "bit": 4,
                 "group\_size": 64
             },
             "input":{
                 "bit": 16
             }
         }
 }
}

#### 量化工具输出件

trained.pth            ### 量化仿真可使用该文件
fake\_quant\_weight.pth  ### 导出onnx文件使用该权重
trained\_quant\_weight.pth ### 阶段一的输出，阶段二的输入

#### 量化仿真

量化完成后，开发者可进行量化仿真推理，通过对比量化模型与原始浮点模型的输出结果，来评估量化模型精度是否满足要求。量化仿真推理工程可参考[qwen2模型量化仿真推理demo](#qwen2模型量化仿真推理demo)。

![](https://contentcenter-vali-drcn.dbankcdn.cn/pvt_2/DeveloperAlliance_scene_100_1/92/v3/iK5iY4tOTt-2dxPLLLxs1g/zh-cn_image_0000002594829434.jpg?HW-CC-KV=V1&HW-CC-Date=20260612T073654Z&HW-CC-Expire=86400&HW-CC-Sign=9AD96F5D75C89558AD9606E29E36BBD7080FB84CD67D7CC10D14A40EACD8716B)

#### \[h2\]插件方法

浮点模型完成加载和初始化后，使用以下接口将浮点模型转换为量化模型，模型推理逻辑不变，可参考[qwen2模型量化仿真推理demo](#qwen2模型量化仿真推理demo)。

import sys
sys.path.append("path/to/dopt")
def get\_quanted\_model(base\_model, dopt\_config, quanted\_ckpt):
     from dopt.dopt\_lm.do\_opt import(
         optimize\_model,
         set\_quant\_state,
         set\_calibrate\_state,
         set\_run\_mode,
     )
     model = optimize\_model(base\_model, dopt\_config)
     model.load\_state\_dict(torch.load(quanted\_ckpt, map\_location=torch.device('cpu')), strict=True)
     set\_quant\_state(model, weight\_state=True, input\_state=True)
     set\_calibrate\_state(model, False)
     model.eval()
     return model

*   base\_model：开发者浮点模型（使用transformers库AutoModelForCausalLM类进行加载）。
    
*   dopt\_config：量化配置文件。
    
*   quanted\_ckpt：量化后pth文件。
    

#### \[h2\]qwen2模型量化仿真推理demo

import os
import sys
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
sys.path.append('path/to/dopt')
os.environ\["CUDA\_VISIBLE\_DEVICES"\] = "0"
def get\_quanted\_model(base\_model, dopt\_config, quanted\_ckpt):
    """
    加载量化模型核心函数
    :param base\_model: 原始HF模型对象
    :param dopt\_config: 量化配置文件路径（dopt\_config.json）
    :param quanted\_ckpt: 量化权重路径（trained.pth）
    :return: 量化后的模型
    """
    from dopt.dopt\_lm.do\_opt import(
        optimize\_model,
        set\_quant\_state,
        set\_calibrate\_state,
        set\_run\_mode,
    )
    # 模型量化优化（根据配置文件应用4bit量化策略）
    model = optimize\_model(base\_model, dopt\_config)
    # 加载量化权重（强制CPU加载避免显存冲突）
    model.load\_state\_dict(torch.load(quanted\_ckpt, map\_location=torch.device('cpu')), strict=True)
    # 设置量化状态
    set\_quant\_state(model, weight\_state=True, input\_state=True)  # 启用权重和激活量化
    set\_calibrate\_state(model, False)  # 关闭校准模式
    model.eval()
    return model
def generate(prompt = "Give me a short introduction to large language model."):
    """
    量化模型推理函数
    :param prompt: 输入文本（默认示例prompt）
    :return: 模型生成的响应
    """
    # 构建Qwen2专用对话模板
    messages = \[
        {"role": "system", "content": "You are Qwen, created by Alibaba Cloud..."},
        {"role": "user", "content": prompt}
    \]
    # 应用模板并tokenize
    text = tokenizer.apply\_chat\_template(
        messages,
        tokenize=False,
        add\_generation\_prompt=True
    )
    model\_inputs = tokenizer(\[text\], return\_tensors="pt").to(model.device)
    # 执行生成
    generated\_ids = model.generate(
        \*\*model\_inputs,
        max\_new\_tokens=512,  # 控制最大生成长度
    )
    # 后处理：去除输入部分
    generated\_ids = \[
        output\_ids\[len(input\_ids):\]
        for input\_ids, output\_ids in zip(model\_inputs.input\_ids, generated\_ids)
    \]
    return tokenizer.batch\_decode(generated\_ids, skip\_special\_tokens=True)\[0\]
if \_\_name\_\_ == '\_\_main\_\_':
    # === 1. 加载原始模型 ===
    model\_name = "path/to/model"  # 需替换为实际模型路径
    model = AutoModelForCausalLM.from\_pretrained(
        model\_name,
        torch\_dtype=torch.float16,  # 半精度加载
        device\_map="auto"           # 自动分配GPU
    )
    tokenizer = AutoTokenizer.from\_pretrained(model\_name)
    # === 2. 加载量化模型 ===
    quant\_res\_root = 'dsr1\_qwen7b\_ptq'  # 量化结果目录
    dopt\_config = f"./{quant\_res\_root}/dopt\_config.json"    # 阶段一生成的配置文件
    quanted\_ckpt = f"./{quant\_res\_root}/train\_output/trained.pth"  # 阶段二生成的权重
    model = get\_quanted\_model(
        model,
        dopt\_config,  # 需确保已正确配置量化参数
        quanted\_ckpt   # 需与当前模型架构匹配
    )
    # === 3. 执行推理测试 ===
    prompt = "who are you?"
    response = generate(prompt)
    print("量化模型推理结果：", response)
    """
    预期输出示例：
    "Hi, I am Qwen, ..."
    """