W4A16量化

大模型推理中，模型权重数据类型（weight），推理计算时的数据类型（activation）和kvcache一般使用半精度浮点FP16或BF16。量化指将高比特的浮点转换为更低比特的数据类型的过程。例如int4、int8等。

模型量化分为weight-only量化，weight-activation量化和kvcache量化。

量化的一般步骤是：1、对浮点类型的权重镜像量化并保存量化完的权重；2、使用量化完的权重进行推理部署。

1. 使用该量化工具，需要切换conda环境，运行以下命令。

   conda create --name llmcompressor --clone PyTorch-2.5.1  
   conda activate llmcompressor

2. 安装llm-compressor。

   pip install llmcompressor==0.6.0
   pip install transformers==4.51.3
   pip install zstandard

3. 创建量化脚本，完整的llm_compressor_W4A16.py脚本请参见llm_compressor_W4A16.py脚本。
4. 配置校准数据集。W4A16量化脚本中默认使用"mit-han-lab/pile-val-backup"数据集作为校准数据集，如需指定数据集，请修改"llm_compressor_W4A16.py"脚本中如下字段。

   DATASET_ID = "mit-han-lab/pile-val-backup"
   DATASET_SPLIT = "validation"
   
   NUM_CALIBRATION_SAMPLES = 256
   MAX_SEQUENCE_LENGTH = 512

   如果当前环境无法访问HuggingFace网站获取数据集，可以先将数据集通过浏览器下载到本地后上传至服务器，并将DATASET_ID指定到服务器路径，数据集下载地址：https://huggingface.co/datasets/mit-han-lab/pile-val-backup/tree/main

   下载后将val.jsonl.zst文件上传至服务器自定义目录，例如<local-dir>，并解压：

   cd <local-dir>
   zstd -d val.jsonl.zst

   解压后的文件名为val.jsonl，解压成功后需要删除原始文件，避免影响后续量化时读取数据：

   rm -f val.jsonl.zst

   DATASET_ID指定到val.jsonl文件所在的文件夹目录名称即可：

   DATASET_ID = "<local-dir>"

5. 配置量化算法。W4A16量化脚本中默认提供AWQ非对称量化算法，如下所示：

   # Configure the quantization algorithm to run.
   recipe = [
       AWQModifier(ignore=["lm_head"], scheme="W4A16_ASYM", targets=["Linear"]),
   ]

   如果使用对称量化算法，可以配置scheme="W4A16"。

6. 运行"llm_compressor_W4A16.py"文件进行模型量化，模型权重路径需要根据实际情况修改，量化时间和模型大小有关，预计30分钟~3小时。

   # 根据机器卡的空闲情况指定量化使用的卡号，不设置则默认使用0号卡
   export ASCEND_RT_VISIBLE_DEVICES=0
   python llm_compressor_W4A16.py --model-path /home/ma-user/Qwen3-32B/ --quant-path /home/ma-user/Qwen3-32B-quant/ 

   参数说明：

   • --model-path：原始模型权重路径
   • --quant-path：转换后保存权重路径
7. 参考启动在线推理服务，在启动服务时添加如下量化参数：

    -q compressed-tensors 或者--quantization compressed-tensors

   W4A16模型当前不支持eager或acl-graph模式，服务启动时需要开启ascend_turbo图模式，配置如下：

   --additional-config: {"ascend_turbo_graph_config": {"enabled": true}} 

   W4A16模型不支持配置Qwen系列性能优化环境变量，服务启动时不支持设置如下环境变量：

   unset ENABLE_QWEN_HYPERDRIVE_OPT
   unset ENABLE_QWEN_MICROBATCH
   unset ENABLE_PHASE_AWARE_QKVO_QUANT
   unset DISABLE_QWEN_DP_PROJ

   针对Qwen2.5-72B-instruct模型和Qwen2-72B-instruct模型，在量化完成后启动服务之前，需要先修改量化后的W4A16模型配置文件config.json，将intermediate_size的值从29568改为29696，原因是Qwen2.5-72B-instruct模型的FFN blocks为29568，W4A16量化group size为128，29568/128=231，该数字无法被大于1的TP整除，因此在TP大于1的场景下无法正常运行服务。此处将intermediate_size改为29696，29696/128=232，该数可以被TP=4或者TP=8整除，因此可以进行多卡推理（参考社区说明）。修改配置后，服务启动时会自动将权重padding到指定shape。修改方式如下：

   vim config.json
   # 修改如下配置项
   "intermediate_size": 29696,

import argparse
import os
from functools import partial
import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu

from datasets import load_dataset
from llmcompressor import oneshot
from llmcompressor.modifiers.awq import AWQModifier
from transformers import AutoModelForCausalLM, AutoTokenizer


os.environ['PYTORCH_NPU_ALLOC_CONF'] = 'expandable_segments:False'

# Select calibration dataset.
DATASET_ID = "mit-han-lab/pile-val-backup"
DATASET_SPLIT = "validation"

# Select number of samples. 256 samples is a good place to start.
# Increasing the number of samples can improve accuracy.
NUM_CALIBRATION_SAMPLES = 256
MAX_SEQUENCE_LENGTH = 512


def preprocess(example, tokenizer):
    return {
        "text": tokenizer.apply_chat_template(
            [{"role": "user", "content": example["text"]}],
            tokenize=False,
        )
    }


def quantize(model_id, quantized_path):
    # Select model and load it
    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto")
    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)

    ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]")
    ds = ds.shuffle(seed=42)
    preprocess_func = partial(preprocess, tokenizer=tokenizer)
    ds = ds.map(preprocess_func)

    # Configure the quantization algorithm to run.
    recipe = [
        AWQModifier(ignore=["lm_head"], scheme="W4A16_ASYM", targets=["Linear"]),
    ]

    oneshot(
        model=model,
        dataset=ds,
        recipe=recipe,
        max_seq_length=MAX_SEQUENCE_LENGTH,
        num_calibration_samples=NUM_CALIBRATION_SAMPLES,
    )

    # Save to disk compressed.
    model.save_pretrained(quantized_path, save_compressed=True)
    tokenizer.save_pretrained(quantized_path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", type=str, required=True, help="Path to the input model")
    parser.add_argument("--quant-path", type=str, required=True, help="Path to save the compressed model")
    args = parser.parse_args()
    quantize(args.model_path, args.quant_path)