文档首页/ AI开发平台ModelArts/ 最佳实践/ LLM大语言模型训练推理/ 主流开源大模型基于DevServer适配PyTorch NPU推理指导（6.3.908）/ 推理模型量化/ 使用kv-cache-int8量化
更新时间：2024-09-14 GMT+08:00
查看PDF
分享

使用kv-cache-int8量化

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

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

Step1使用tensorRT量化工具进行模型量化，必须在GPU环境

在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系数。

Step2 抽取kv-cache量化系数

该步骤的目的是将Step1使用tensorRT量化工具进行模型量化中生成的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系数的抽取和加载。

Step3 启动kv-cache-int8量化服务

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

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

相关文档