Awesome

[DEPRECATED]

This project is deprecated. Use LLM-TPU instead. Thank you.

ChatGLM2-TPU

本项目实现BM1684X部署语言大模型ChatGLM2-6B。通过TPU-MLIR编译器将模型转换成bmodel，并采用c++代码将其部署到BM1684X的PCIE环境，或者SoC环境。

下文中默认是PCIE环境；如果是SoC环境，按提示操作即可。

在知乎上写了关于ChatGLM2-6B的解读，方便大家理解源码：

ChatGLM2-6B流程解析与TPU-MLIR部署

开发环境

下载docker，启动容器，如下：

docker pull sophgo/tpuc_dev:latest

# myname1234 is just an example, you can set your own name
docker run --privileged --name myname1234 -v $PWD:/workspace -it sophgo/tpuc_dev:latest

后文假定环境都在docker的/workspace目录。

如果是要在SoC环境运行，则需要安装如下库：

apt-get install gcc-aarch64-linux-gnu g++-aarch64-linux-gnu

下载ChatGLM2-6B，比较大，会花较长时间

git lfs install
git clone git@hf.co:THUDM/chatglm2-6b

并对该工程做三点修改：

将config.json文件中seq_length配置为512；
将modeling_chatglm.py文件中的如下代码：

if attention_mask is not None:
    attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))

修改为：

if attention_mask is not None:
    attention_scores = attention_scores + (attention_mask * -10000.0)

这样修改可以提升效率，使用masked_fill效率低下；另一方面masked_fill转ONNX存在些bug。

将modeling_chatglm.py文件中的如下代码：

pytorch_major_version = int(torch.__version__.split('.')[0])
if pytorch_major_version >= 2:

修改为：

pytorch_major_version = int(torch.__version__.split('.')[0])
if False:

下载TPU-MLIR代码并编译，(也可以直接下载编译好的release包解压)

git clone git@github.com:sophgo/tpu-mlir.git
cd tpu-mlir
source ./envsetup.sh
./build.sh

下载sentencepiece，并编译得到sentencepiece.a

git clone git@github.com:google/sentencepiece.git
cd sentencepiece
mkdir build
cd build
cmake ..
make -j

如果要编译SoC环境，则需要在CMakeLists.txt加入如下代码：

set(CMAKE_C_COMPILER aarch64-linux-gnu-gcc)
set(CMAKE_ASM_COMPILER aarch64-linux-gnu-gcc)
set(CMAKE_CXX_COMPILER aarch64-linux-gnu-g++)

下载libsophon库并安装

在算能官网https://developer.sophgo.com/site/index/material/all/all.html可以找到SDK最新版本，如下：

wget https://sophon-file.sophon.cn/sophon-prod-s3/drive/23/06/15/16/Release_230501-public.zip

解压sdk后安装libsophon，如下：

apt install sophon-libsophon-dev_0.4.8_amd64.deb

注意如果是SoC环境则安装arm64版本sophon-libsophon-dev_0.4.8_arm64.deb

下载本项目ChatGLM2-TPU，如下：

git clone git@github.com:sophgo/ChatGLM2-TPU.git

编译模型

指定ChatGLM2-6B的python路径

export PYTHONPATH=/workspace/chatglm2-6b:$PYTHONPATH

导出所有onnx模型，如果过程中提示缺少某些组件，直接pip install 组件即可

cd chatglm2-tpu/compile
python3 export_onnx.py

此时有大量onnx模型被导出到tmp目录。

对onnx模型进行编译

目前TPU-MLIR支持对ChatGLM2进行F16, INT8和INT4量化，且支持多芯分布式推理，默认情况下会进行F16量化和单芯推理，最终生成chatglm2-6b.bmodel文件

./compile.sh

若想进行INT8或INT4量化，则执行以下命令，最终生成chatglm2-6b_int8.bmodel或chatglm2-6b_int4.bmodel文件

./compile.sh --mode int8 # or int4

若想进行2芯推理，则执行以下命令，最终生成chatglm2-6b_f16_2dev.bmodel文件

./compile.sh --num_device 2

编译程序(C++版本)

cd chatglm2-tpu/demo
mkdir build
cd build
cmake ..
make -j

如果是SoC环境，则在CMakeLists.txt中加入以下代码，并将SoC版本的libsentencepiece.a替换过来：

set(CMAKE_C_COMPILER aarch64-linux-gnu-gcc)
set(CMAKE_ASM_COMPILER aarch64-linux-gnu-gcc)
set(CMAKE_CXX_COMPILER aarch64-linux-gnu-g++)

编译生成chatglm2可执行程序，将chatglm2、chatglm2-6b?.bmodel和tokenizer.model拷贝到运行环境就可以执行了。 (tokenizer.model来自ChatGLM2-6B)。

运行chatglm2，默认单芯运行chatglm2-6b.bmodel:

./chatglm2

如果是要运行INT8或INT4模型，则命令如下：

./chatglm2 --model chatglm2-6b_int8.bmodel # same with int4

如果是2芯分布式推理，使用如下命令(比如指定在2号和3号芯片上运行, 用bm-smi查询芯片id号)：

./chatglm2 --model chatglm2-6b_f16_2dev.bmodel --dev_id 2,3

编译程序(Python版本)

cd chatglm2-tpu/python_demo
mkdir build
cd build
cmake ..
make -j

编译成功会生成ChatGLM2.cpython-37m-x86_64-linux-gnu.so，之后将chatglm2-6b.bmodel放到python_demo目录下。另外这里也直接给出了so文件，可以直接省略上面的编译这一步 (但是必须为python3.7版本)。

若想采用INT8或INT4量化，则需要在编译前将ChatGLM2.cpp中的CHATGLM_MODEL更改为对应的bmodel名称。

python run.py

即可成功运行python的demo。

如果是SoC环境，则在CMakeLists.txt中加入以下代码，并将SoC版本的libsentencepiece.a替换过来：

set(CMAKE_C_COMPILER aarch64-linux-gnu-gcc)
set(CMAKE_ASM_COMPILER aarch64-linux-gnu-gcc)
set(CMAKE_CXX_COMPILER aarch64-linux-gnu-g++)

编译程序(Python Web版本)

cd chatglm2-tpu/web_demo
mkdir build
cd build
cmake ..
make -j

编译成功会生成libtpuchat.so*，在chat.py中指定bmodel_path token_path device_id。

python web_demo.py

即可成功运行web的demo。

如果是SoC环境，则在CMakeLists.txt中加入以下代码，并将SoC版本的libsentencepiece.a替换过来：

set(CMAKE_C_COMPILER aarch64-linux-gnu-gcc)
set(CMAKE_ASM_COMPILER aarch64-linux-gnu-gcc)
set(CMAKE_CXX_COMPILER aarch64-linux-gnu-g++)

运行效果

以下为单芯片下INT8量化模式的运行效果：