推理加速

2023年12月8日修改
随着大模型参数规模的不断增长,在有限的算力资源下,提升模型的推理速度逐渐变为一个重要的研究方向。常用的推理加速框架包含lmdeploy、FasterTransformer和vLLM等。
一.lmdeploy推理部署
  lmdeploy由上海人工智能实验室开发,推理使用C++/CUDA,对外提供python/gRPC/http接口和WebUI界面,支持tensor parallel分布式推理、支持fp16/weightint4/kv cache int8量化。lmdeploy支持transformer结构(例如LLaMA、LLaMa2、InternLM、Vicuna等),目前支持fp16,int8和int4。
1.fp16推理
  安装预编译的Python包,即python3 -m pip install lmdeploy。把模型转成lmdeploy 推理格式,假设huggingface版LLaMa2模型已下载到/models/llama-2-7b-chat目录,结果会存到workspace文件夹,如下所示:
代码块
python3 -m lmdeploy.serve.turbomind.deploy llama2 /models/llama-2-7b-chat
  在命令行中测试聊天效果,如下所示:
代码块
python3 -m lmdeploy.turbomind.chat ./workspace
...
double enter to end input >>> who are you
...
Hello! I'm just an AI assistant ...
  也可以用gradio启动WebUI来聊天,如下所示:
代码块
python3 -m lmdeploy.serve.gradio.app ./workspace
  lmdeploy同样支持原始的facebook模型格式、支持70B模型分布式推理,用法查看lmdeploy官方文档
2.kv cache int8量化
  lmdeploy实现了kv cache int8量化,同样的显存可以服务更多并发用户。首先计算模型参数,结果是pth格式,保存到临时目录minmax,如下所示:
代码块
mkdir minmax
python3 -m lmdeploy.lite.apis.calibrate \
--model /models/llama-2-7b-chat \ # huggingface llama2 模型。也支持 llama/vicuna/internlm/baichuan 等
--calib_dataset 'c4' \ # 校准数据集,支持 c4, ptb, wikitext2, pileval
--calib_samples 128 \ # 校准集的样本数,如果显存不够,可以适当调小
--calib_seqlen 2048 \ # 单条的文本长度,如果显存不够,可以适当调小
--work_dir minmax \ # 保存 pth 格式量化统计参数和量化后权重的文件夹