【实战指南】如何在ARM架构下高效运行PaddleNLP的Ernie-3模型

209 阅读2分钟

本篇文章首发于【算法工程笔记】公众号,更多内容,欢迎关注。

最近项目中使用到了 paddlenlp 中的 Ernie-3 模型,但是比较麻烦的事情是硬件是 arm 架构的 910B4 显卡,经过一番折腾最终还是使用 paddlenlp sever 跑起来了,并且也有了与英伟达 GPU 相匹配的推理速度(单条文本分类10毫秒级别)。

这里简单记录如下。

TL;DR

  • 文章主要分享如何使用paddlenlpErnie-3模型在arm架构的910B4显卡上运行。
  • 操作步骤包括安装特定版本的驱动和Docker镜像。
  • 修改paddlenlp服务文件以支持npu设备。
  • 设置环境变量FLAGS_npu_jit_compile=0以优化服务速度。

昇腾 NPU 支持模型列表

  • BERT
  • ERNIE-3.0
  • UIE
  • UTC
  • RoBERTa

操作步骤

  • 参考 昇腾 NPU,此处需要注意的是昇腾驱动版本,官方建议是 23.0.3 ,建议驱动不要低于此

  • Dockerfile,根据 CPU 架构版本选择不同的基础镜像,下面以arm架构举例

# Use the specified base image
# arm 架构 CPU
FROM registry.baidubce.com/device/paddle-npu:cann80T13-ubuntu20-aarch64-gcc84-py39

# Set the working directory
WORKDIR /work

# Install paddlepaddle and paddlenlp
RUN python3.9 -m pip install paddlepaddle==3.0.0.dev20240520 -i https://www.paddlepaddle.org.cn/packages/nightly/cpu/ && \
    python3.9 -m pip install paddle_custom_npu==3.0.0.dev20240719 -i https://www.paddlepaddle.org.cn/packages/nightly/npu/ && \
    pip3 install paddlenlp -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple

# Command to run the application
CMD ["python3", "-c", "'import paddle; paddle.utils.run_check()'"]
  • 构建镜像
docker build -t your-paddle-image-name .
  • 验证是否可以运行,注意此处可能会有某些库文件加载的问题,如果有的话,需要将启动命令前加上 source ~/.bashrc
docker run -it --name paddle-npu-dev -v $(pwd):/work \
    --privileged --network=host --shm-size=128G -w=/work \
    -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
    -v /usr/local/dcmi:/usr/local/dcmi \
    -e ASCEND_RT_VISIBLE_DEVICES="0,1,2,3,4,5,6,7" \
    your-paddle-image-name /bin/bash
  • 修改 paddlenlp simple server 的部分文件内容
# PADDLENLP_PATH 为 paddlenlp 包的绝对路径
sed -i 's/gpu/npu/' PADDLENLP_PATH/server/model_manager.py
sed -i 's/gpu/npu/' PADDLENLP_PATH/server/predictor.py 
FLAGS_npu_jit_compile=0 paddlenlp server server:app --host 0.0.0.0 --port 8189 --workers 1