310芯片运行LLM

https://gitee.com/yinghuo302/ascend-llm

简介

本项目基于昇腾310芯片部署大语言模型，目前已经成功运行meta-llama/Llama-2-7b-hf和TinyLlama/TinyLlama-1.1B-Chat-v1.0。

本实践项目由南京大学计算机科学与技术系杜骋同学主导，由朱光辉老师进行指导，由昇腾CANN生态使能团队提供技术支持，并在昇腾开发者大会2024进行了展示。

效果预览

invalid image(图片无法加载)

关键技术

静态图方案

在Transformer模型中，基于模型的自回归推理特性，业界普遍采用kvcache缓存的方式增加模型的推理性能。kvcache会缓存上一次推理得到的kv矩阵用于本次推理，大大减少了推理计算量。

由于缓存的kv矩阵要和当前输入字符计算出的kv矩阵进行拼接，因此每次推理完整的kv矩阵长度一直在增加，致使模型shape不固定，会走动态推理流程，存在大量算子编译时间，推理性能大大下降。

本方案基于原先动态图方案，将kv矩阵固定到一个最大长度，结合attention_mask屏蔽输入序列部分位置的特性实现了静态图的方案。在kvcache达到上限时通过KV缓存驱逐（StreamingLLM和Heavy-Hitter Oracle）让模型可以反复推理。
量化方案

大模型权重过大，在端侧设备由于内存限制通常难以运行，因此通常将大模型权重从fp16量化到int8甚至int4降低内存消耗.

本项目采用平滑激活（SmoothQuant），动态混合精度分解（类似LLM.int8），静态混合精度分解量化方案，通过对权重和激活值均采用int8量化，显著节省了内存并提升了推理速度。

运行方式

环境准备

昇腾软硬件解决方案(驱动+固件+CANN)前往昇腾社区，按照说明下载安装。
或者下载香橙派0318镜像，烧录到sd卡，启动环境，参考香橙派AIpro快速上手指南。
第三方依赖模型导出和推理相关文件夹下requirements.txt，使用pip 进行安装。
```
pip install -r requirements.txt
```

本项目测试环境：香橙派AI pro，CANN 7.0/7.2，python 3.9。

算子适配

protoc安装
```
 根据昇腾文档选择合适的protoc版本，protoc版本和CANN版本强相关。CANN7.0/7.2使用的protoc 1.13.0

 ```
```
安装protoc==1.13.0，找一空闲目录下载
wget https://obs-9be7.obs.cn-east-2.myhuaweicloud.com/wanzutao/tiny-llama/protobuf-all-3.13.0.tar.gz --no-check-certificate
tar -zxvf protobuf-all-3.13.0.tar.gz
cd protobuf-3.13.0
apt-get update
apt-get install autoconf automake libtool
./autogen.sh
./configure
make -j4
make install
sudo ldconfig
protoc --version # 查看版本号
```

算子编译部署

# 将./custom_op/matmul_integer_plugin.cc 拷贝到指定路径
cd tiny_llama
export ASCEND_PATH=/usr/local/Ascend/ascend-toolkit/latest
cp custom_op/matmul_integer_plugin.cc $ASCEND_PATH/tools/msopgen/template/custom_operator_sample/DSL/Onnx/framework/onnx_plugin/
cd $ASCEND_PATH/tools/msopgen/template/custom_operator_sample/DSL/Onnx

打开build.sh，找到下面四个环境变量，解开注释并修改如下：

export ASCEND_TENSOR_COMPILER_INCLUDE=/usr/local/Ascend/ascend-toolkit/latest/include
export TOOLCHAIN_DIR=/usr
export AICPU_KERNEL_TARGET=cust_aicpu_kernels
export AICPU_SOC_VERSION=Ascend310B4

编译运行

./build.sh 
cd build_out/
./custom_opp_ubuntu_aarch64.run
# 生成文件到customize到默认目录 $ASCEND_PATH/opp/vendors/，删除冗余文件
cd $ASCEND_PATH/opp/vendors/customize
rm -rf op_impl/ op_proto/

模型量化与导出

导出的模型可以从阿里云盘中下载

导出onnx：将transformer库中的modeling_llama替换为export_llama文件下的modeling_llama。通过一下命令将模型导出为onnx（相对路径均为相对export_llama.py文件）

 ```bash
 python export_llama.py \
         --model <model_name_or_path> \
         --output <output_onnx_file_path> \
         --act-path <act_scales_file_path>
         --quant <quant_config_file_path>
 ```
 模型量化具体见[readme](./export_llama/readme.md)。对于TinyLlama-1.1B建议采用per-token的absmax量化（即w8x8.py）或者平滑激活（即smooth.py）；对于Llama-2-7b-hf，建议采用静态混合精度分解（即sd.py）或者平滑激活+静态混合精度分解(即smsd.py)。已经测试的方案为TinyLlama-1.1B per-token的absmax量化，Llama-2-7b-hf 静态混合精度分解。

ATC模型转换

 ``` bash
 atc --framework=5 --model="xxx.onnx"  --output="xxx" --input_format=ND --input_shape="input_ids:batch,seq_len;attention_mask:batch,seq_len+kv_len;position_ids:batch,seq_len;past_key_values:n_layer,2,batch,n_head,kv_len,head_dim" --log=debug --soc_version=Ascend310B1 --precision_mode=must_keep_origin_dtype
 ```
 上述的n_layer, n_head, head_dim变量由模型决定。对于Llama-2-7b，n_layer=32, n_head=32, head_dim=128；对于TinyLlama-1.1B，n_layer=22, n_head=4, head_dim=64

 对于batch, seq_len, kv_len, 请根据需要填入，建议设置batch=1, seq_len=1, kv_len=1024。如对于TinyLlama-1.1B

 ```bash
 atc --framework=5 --model="./tiny-llama.onnx"  --output="tiny-llama" --input_format=ND --input_shape="input_ids:1,1;attention_mask:1,1025;position_ids:1,1;past_key_values:22,2,1,4,1024,64" --log=debug --soc_version=Ascend310B1 --precision_mode=must_keep_origin_dtype
 ```

 对于Llama-2-7b，ATC转换占用内存较大，建议采用其他设备转换，如采用香橙派进行模型转换可以`export MAX_COMPILE_CORE_NUMBER=1`和`export TE_PARALLEL_COMPILER=1`，并开swap分区（推理时请关闭swap，会影响性能）。

模型推理运行

目前提供两种运行模式：

cli模式：在终端运行，每一次输入一行，一次性返回所有的推理结果。
web模式：前端代码在github或者gitee，打包出dist文件夹，放在inference文件夹下即可。

cd inference
python main.py \
        --model <path_to_onnx_or_om_model> \
        --hf-dir <path_to_huggingface_model_dir>  \ # 需要tokenizer和模型配置文件，权重不需要 
        --engine <acl/onnx> \
        --sampling <greedy/top_p/top_k> --sampling_value <>  --temperature <> \ #采样相关配置
        --cli # 添加--cli表示在终端运行