W8A8量化

什么是W8A8量化

W8A8量化是一种将模型权重和激活值都量化为8位数据的技术。该技术把高位浮点数转为8位，通常是将权重和激活值从16位或32位浮点数转换为8位整数（int8）格式。量化后，模型权重体积会减少，同时使用int8格式数据进行矩阵乘法（MatMul）运算时，可减少计算量，进而提升推理性能。

W8A8量化方案能降低模型显存以及需要部署的卡数。也能同时降低首token时延和增量推理时延。

约束限制

• 支持W8A8量化的模型列表请参见支持的模型列表。
• 激活量化支持动态per-token，支持对称量化。
• 权重量化支持per-channel，支持对称量化。

获取量化后的模型权重

有两种方式获取量化模型权重：

方式一：可以在Huggingface开源社区获取量化后的模型权重，具体参考从开源社区下载发布的llm-compressor量化模型。

方式二：获取FP16/BF16的模型权重之后，通过llm-compressor工具进行量化，具体参考使用llm-compressor量化工具进行量化。

使用llm-compressor工具量化模型

本章节介绍如何在NPU的机器上使用开源量化工具llm-compressor量化模型权重，然后在NPU的机器上实现推理量化。了解更多开源量化工具信息，请参考llm-compressor。

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. W8A8量化默认使用"HuggingFaceH4/ultrachat_200k"数据集作为校准数据集，如需指定数据集请修改"llm_compressor_W8A8.py"脚本中如下字段：

   DATASET_ID = "HuggingFaceH4/ultrachat_200k"
   DATASET_SPLIT = "train_sft"
   
   NUM_CALIBRATION_SAMPLES = 512
   MAX_SEQUENCE_LENGTH = 2048

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

   下载的数据集需要保留原始的data目录，下载后目录结构示例如下，其中<local-data-dir>是存放数据集的自定义目录：

   <local-data-dir>
      -- ultrachat_200k
          -- data
             -- train_sft-xxx.parquet
             -- train_gen-xxx.parquet

   指定本地数据集目录，需要指定到data的上一层级，不包含data：

   DATASET_ID="<local-data-dir>/ultrachat_200k"

   完整的llm_compressor_W8A8.py脚本请参见llm_compressor_W8A8.py脚本。

4. 运行"llm_compressor_W8A8.py"文件进行模型量化，量化时间和模型大小有关，预计30分钟~3小时。

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

   参数说明：

   • --model-path：原始模型权重路径
   • --quant-path：转换后保存权重路径

   量化执行过程中如果打印如下WARNING日志，可以不用关注，不影响量化功能正常使用：

   get_GPU_usage_nv | WARNING - Pynml library error:
    NVML Shared Library Not Found

5. 启动在线推理服务，在启动服务时添加如下命令。

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

llm_compressor_W8A8.py脚本

llm_compressor_W8A8.py脚本完整代码如下

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 llmcompressor.modifiers.quantization import GPTQModifier
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
from llmcompressor.utils import dispatch_for_generation
from transformers import AutoModelForCausalLM, AutoTokenizer


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

# Select calibration dataset.
DATASET_ID = "HuggingFaceH4/ultrachat_200k"
DATASET_SPLIT = "train_sft"

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


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


# Tokenize inputs.
def tokenize(sample, tokenizer):
    return tokenizer(
        sample["text"],
        padding=False,
        max_length=MAX_SEQUENCE_LENGTH,
        truncation=True,
        add_special_tokens=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)

    # Load dataset and preprocess.
    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)
    tokenize_func = partial(tokenize, tokenizer=tokenizer)
    ds = ds.map(tokenize_func, remove_columns=ds.column_names)

    # Configure algorithms. In this case, we:
    #   * apply SmoothQuant to make the activations easier to quantize
    #   * quantize the weights to int8 with GPTQ (static per channel)
    #   * quantize the activations to int8 (dynamic per token)
    recipe = [
        SmoothQuantModifier(smoothing_strength=0.8),
        GPTQModifier(targets="Linear", scheme="W8A8", ignore=["lm_head"]),
    ]

    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)