更新时间:2024-11-21 GMT+08:00
分享

使用kv-cache-int8量化

kv-cache-int8是实验特性,在部分场景下性能可能会劣于非量化。当前支持per-tensor静态量化、per-tensor+per-head静态量化以及per-token,支持kv-cache-int8量化和FP16、BF16、AWQ、smoothquant的组合。

kv-cache-int8量化支持的模型请参见表3

per-tensor静态量化场景

在GPU机器上使用tensorRT 0.9.0版本工具进行模型量化,工具下载使用指导请参见https://github.com/NVIDIA/TensorRT-LLM/tree/v0.9.0

量化脚本convert_checkpoint.py存放在TensorRT-LLM/examples路径对应的模型文件夹下,例如:llama模型对应量化脚本的路径是examples/llama/convert_checkpoint.py。

执行convert_checkpoint.py脚本进行权重转换生成量化系数。

  1. 使用tensorRT量化工具进行模型量化。

    在GPU机器上使用tensorRT 0.9.0版本工具进行模型量化,工具下载使用指导请参见https://github.com/NVIDIA/TensorRT-LLM/tree/v0.9.0

    量化脚本convert_checkpoint.py存放在TensorRT-LLM/examples路径对应的模型文件夹下,例如:llama模型对应量化脚本的路径是examples/llama/convert_checkpoint.py。

    执行convert_checkpoint.py脚本进行权重转换生成量化系数,详细参数解释请参见https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/llama#int8-kv-cache

    python convert_checkpoint.py \
    --model_dir ./llama-models/llama-7b-hf   \
    --output_dir ./llama-models/llama-7b-hf/int8_kv_cache/ \
    --dtype float16  \
    --int8_kv_cache

    运行完成后,会在output_dir下生成量化后的权重。量化后的权重包括原始权重和kvcache的scale系数。

  2. 抽取kv-cache量化系数。

    该步骤的目的是将步骤1中生成的scale系数提取到单独文件中,供推理时使用。

    使用的抽取脚本由vllm社区提供:

    python3 examples/fp8/extract_scales.py \
    --quantized_model <QUANTIZED_MODEL_DIR> \
    --tp_size <TENSOR_PARALLEL_SIZE> \
    --output_dir <PATH_TO_OUTPUT_DIR>

    运行后在--output_dir下生成kv_cache_scales.json文件,里面是提取的per-tensor的scale值。内容示例如下:

    注意:

    1. 抽取完成后,可能提取不到model_type信息,需要手动将model_type修改为指定模型,如"llama"。
    2. 当前社区vllm只支持float8的kv_cache量化,抽取脚本中dtype类型是"float8_e4m3fn"。dtype类型不影响int8的scale系数的抽取和加载。

  3. 启动kv-cache-int8-per-tensor量化服务。

    在使用OpenAI接口或vLLM接口启动推理服务时添加如下参数:

    --kv-cache-dtype int8_pertensor  #只支持int8,表示kvint8 per-tensor量化
    --quantization-param-path kv_cache_scales.json  #输入2. 抽取kv-cache量化系数生成的json文件路径; 如果只测试推理功能和性能,不需要此json文件,此时scale系数默认为1,但是可能会造成精度下降。

per-tensor+per-head静态量化场景

如需使用该场景量化方法,请自行准备kv-cache量化系数,格式和per-tensor静态量化所需的2. 抽取kv-cache量化系数生成的json文件一致,只需把每一层的量化系数修改为列表,列表的长度为kv的头数,列表中每一个值代表每一个kv头使用的量化系数。内容示例如下:

在使用OpenAI接口或vLLM接口启动推理服务时添加如下参数,启动kv-cache-int8-per-tensor+per-head量化服务。

--kv-cache-dtype int8_pertensor_perhead  #只支持int8,表示kvint8 per-tensor+per-head量化
--quantization-param-path kv_cache_scales.json  #输入生成的json文件路径; 如果只测试推理功能和性能,不需要此json文件,此时scale系数默认为1,但是可能会造成精度下降。

per-token动态量化场景

如需使用该场景量化方法,推理前向会自动计算kv-cache量化系数,并进行kv的量化。

在使用OpenAI接口或vLLM接口启动推理服务时添加如下参数,启动kv-cache-int8-per-token量化服务。

--kv-cache-dtype int8_pertoken  #只支持int8,表示kvint8 per-token量化

相关文档