跳到主要内容

DeepSeek-R1 TPU

DeepSeek-R1 是由杭州深度求索公司开发, 该模型完全开源了所有训练技术和模型权重，性能对齐闭源的 OpenAI-o1, deepseek 通过 DeepSeek-R1 的输出，蒸馏了 6 个小模型给开源社区，包括 Qwen2.5 和 Llama3.1。 本文档将讲述如何使用 TPU-MLIR 将 DeepSeek-R1 蒸馏模型 DeepSeek-R1-Distill-Qwen-1.5B 和 DeepSeek-R1-Distill-Qwen-7B 大语言模型部署到 Airbox 上利用 TPU 进行硬件加速推理。 rkllm_2.webp

DeepSeek-R1 Qwen2.5 蒸馏模型部署​

  • 克隆仓库

    git clone https://github.com/zifeng-radxa/LLM-TPU.git

  • 打开 deepseek-r1 项目路径

    cd LLM-TPU/models/language_model/python_demo

  • 本案例提供 deepseek-r1-distill-qwen-1-5b 和 deepseek-r1-distill-qwen-7b INT4 量化模型与 cpython 预编译动态库下载

    用户可以参考 DeepSeek-R1 模型转换 自行转换不同量化方式的 DeepSeek-R1 蒸馏模型

    用户可以参考 DeepSeek-R1 cpython 动态库编译 自行编译 cpython 接口绑定文件

    使用 modelscope 下载与编译模型

    • pip3 安装 modelscope 
      pip3 install modelscope
    • 克隆 deepseek-r1-distill-qwen-1-5b_TPU 仓库 
      modelscope download --model radxa/deepseek-r1-distill-qwen-1-5b_TPU
    • 克隆 deepseek-r1-distill-qwen-7b_TPU 仓库 
      modelscope download --model radxa/deepseek-r1-distill-qwen-7b_TPU

  • 配置环境

    推荐创建虚拟环境，否则可能会影响其他应用的正常运行， 虚拟环境使用请参考这里

    python3 -m virtualenv .venv
    source .venv/bin/activate

  • 安装依赖包

    pip3 install --upgrade pip
    pip3 install transformers==4.45.1 Jinja2==3.1.5

  • 导入环境变量

    请使用 ldd 命令检查 chat.cpython-38-aarch64-linux-gnu.so 链接的 libbmlib.so 的路径是否为 LLM-TPU/support/lib_soc/lib*

    如 libbmlib.so 链接路径有误可运行下面的命令，请正确指定 LLM-TPU 路径

    export LD_LIBRARY_PATH=LLM-TPU/support/lib_soc:$LD_LIBRARY_PATH
    cp deepseek-r1-distill-qwen-1-5b_TPU/chat.cpython-38-aarch64-linux-gnu.so .
    ldd chat.cpython-38-aarch64-linux-gnu.so

  • 启动 DeepSeek-R1

    终端模式

    python3 pipeline.py --model_path ./deepseek-r1-distill-qwen-1-5b_TPU/qwen2.5-1.5b_int4_seq8192_1dev.bmodel --devid 0 --dir_path ./deepseek-r1-distill-qwen-1-5b_TPU  ```

    -m; --model_path： 指定模型路径

    -d; --device ： 指定 TPU 设备号，默认为 0

    -p; --dir_path： 指定包含 tokenizer 文件夹路径

    deepseek_r1_1.webp

DeepSeek-R1 模型转换​

用户可以参考本文自行转换不同量化类型的 Qwen2-7B-Instruct 模型到 bmodel

  • X86 工作站中准备环境

    请参考 TPU-MLIR 安装 配置 TPU-MLIR 环境。

    提示

    此例子使用的 tpu-mlir 版本为 tpu-mlir==1.14

    克隆仓库

    git clone https://github.com/zifeng-radxa/LLM-TPU.git

  • 通过 Huggingface 下载 DeepSeek-R1 开源蒸馏模型

    提示

    请确保已安装 git lfs

    • DeepSeek-R1-Distill-Qwen-1.5B 
      git clone https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
    • DeepSeek-R1-Distill-Qwen-7B 
      git clone https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B

  • 在工作目录 LLM-TPU/models/Qwen2_5 中创建虚拟环境

    虚拟环境使用请参考这里

    python3 -m virtualenv .venv
    source .venv/bin/activate
    pip3 install --upgrade pip
    pip install transformers_stream_generator einops tiktoken accelerate onnx torch==2.0.1 torchvision==0.15.2 transformers==4.45.2
    提示

    因为蒸馏模型的 base 模型是 qwen2.5, 所以使用 qwen2.5 的方法进行模型编译

  • 对齐模型环境

    将 LLM-TPU/models/Qwen2_5/compile/files/Qwen2-7B-Instruct/modeling_llama.py 复制到 transformers 库中, 注意此时 transformers 库应在 .venv 里

    • DeepSeek-R1-Distill-Qwen-1.5B

      cp ./compile/files/Qwen2.5-1.5B-Instruct/modeling_qwen2.py .venv/lib/python3.10/site-packages/transformers/models/qwen2
      • DeepSeek-R1-Distill-Qwen-7B
      cp ./compile/files/Qwen2.5-7B-Instruct/modeling_qwen2.py .venv/lib/python3.10/site-packages/transformers/models/qwen2

  • 生成 onnx 文件 这里以 8192 长度为例子

    cd compile
    python3 export_onnx.py --model_path your_torch_model --seq_length 8192 --device cpu

    --model_path： 下载的 DeepSeek-R1-Distill-Qwen 文件夹路径

    --seq_length： 固定导出的 sequence length, 根据需要可选 512, 1024，2048, 8192 等长度

  • 生成 bmodel 文件

    生成 bmodel 之前需要退出虚拟环境

    deactivate

    编译模型

    • DeepSeek-R1-Distill-Qwen-1.5B
    ./compile.sh --mode int4 --name qwen2.5-1.5b --addr_mode io_alone --seq_length 8192
    • DeepSeek-R1-Distill-Qwen-7B
    ./compile.sh --mode int4 --name qwen2.5-7b --addr_mode io_alone --seq_length 2048

    --mode：量化模式，可选 int4, int8

    --seq_length： 序列长度，需和生成 onnx 文件时指定相同 seq_length

    --name： 模型名称，可选 qwen2.5-1.5b 和 qwen2-7b

    --addr_mode： 记录网络的地址分配模式，此处选择 io_alone 模式

    提示

    生成 bmodel 耗时大概 1 小时以上，建议 64G 内存以及 100G 以上硬盘空间，不然很可能 OOM 或者 no space left

DeepSeek-R1 cpython 动态库编译​

Airbox 中编译 cpython 动态库,

提示

预编译文件已经包含在下载包中，如通过 modelscope 下载无需编译

pip3 install pybind11[global]
cmake -B build && cmake --build build
cp build/*cpython* .

推理性能​

测试模型量化方式模型长度first token latency(s)token per second(tokens/s)
deepseek-r1-distill-qwen-1.5bINT481925.15930.448
deepseek-r1-distill-qwen-7bINT420482.84311.008