TensorRT加速BERT

Date: 2019/10/10 Categories: 工作工程 Tags: TensorRT BERT DeepLearningInference

之前有调研过一些框架在BERT前向方面的对比, 如具体结果所示 TensorRT在GPU上有很大优势, 本文记录使用TensorRT加速BERT的一些关键点.

Introduction

基本概念

TensorRT是Nvidia开发的深度学习推断框架
BERT: Google发布的2018年最火的NLP模型
TensorRT-Bert: TensorRT目前还没有开源, 但一些周边的工具/demo已经开源了, 其中包括了BERT的实现, 之前网上有一些NV的PR文, 如“2毫秒完成AI推理”等都是用的它
FasterTransformer: 也是来自NV的BERT加速工具, 由国内团队开发. 与上面的TensorRT-Bert的区别是可以作为plugin嵌入到tensorflow中, 更方便复用现有代码, 测试速度略差于TensorRT-Bert.

可以简单的说TensorRT-Bert是一个TensorRT plugin, 而FasterTrnasformer是一个Tensorflow plugin.

下面的篇幅将详细介绍TensorRT-BERT的使用

版本和依赖

目前TensorRT最新的版本是TensorRT-6.1, 而TensorRT-BERT也只支持这个最新版本.

TensorRT还没有开源, 需要注册一个nvidia账号下载. 下载页面

用tlinux2.2的话, 需要选择centos版本的包, 其中包含了库文件, 头文件和例子代码等, 可选三个版本

TensorRT-6.0.1.5.CentOS-7.6.x86_64-gnu.cuda-10.1.cudnn7.6.tar.gz
TensorRT-6.0.1.5.CentOS-7.6.x86_64-gnu.cuda-10.0.cudnn7.6.tar.gz
TensorRT-6.0.1.5.CentOS-7.6.x86_64-gnu.cuda-9.0.cudnn7.6.tar.gz

分别对应cuda10.¹⁄₁₀.0和9.0版本,而cuda版本又分别有各自最低的nvidia内核驱动要求, 选择时候需要参考CUDA Compatibility, 摘录如下:

CUDA Toolkit	Linux x86_64 Driver Version
CUDA 10.1 (10.1.105)	>= 418.39
CUDA 10.0 (10.0.130)	>= 410.48
CUDA 9.0 (9.0.76)	>= 384.81

主机上的nvidia 驱动版本可以用命令nvidia-smi --query-gpu=driver_version --format=csv,noheader查看, 考虑最新的tensorflow-1.14.0的pypi版本使用的cuda10.0, 后续我们都以cuda10.0为例, 如果驱动版本过低(如384.81)需要联系运维升级nvidia内核驱动.

Installtion

如前所述, TensorRT-BERT实际是TensorRT的plugin, 具体来说就是两个so文件, 因此在使用前需要编译之.

配置编译环境

编译环境直接用TensorRT配好的centos环境即可, 编译出来的包和tlinux是兼容的, 使用 Dockerfile build一个开发镜像, 构建时需要指定cuda版本为10.0:docker build -a CUDA_VERSION=10.0 -t trt-bert-dev . 这个开发镜像里已经包含了cuda-10.0的文件

还需要把TensorRT文件放入编译环境, 使用docker bind mount即可, 最终启动编译环境shell的命令类似:

CUDA_VERSION=${CUDA_VERSION:-10.0}
git clone https://github.com/NVIDIA/TensorRT.git
docker run -it --rm -w /workspace/TensorRT/demo/BERT/ \
    -v $PWD/TensorRT/:/workspace/TensorRT \
    -v TensorRT-6.0.1.5-cuda${CUDA_VERSION}:/tensorrt trt-bert-dev bash

一些代码修改

貌似默认的cmake配置编译不过?, 查看几处即可

TensorRT/demo/BERT/CMakeLists.txt:
1. 修改CMAKE_CUDA_FLAGS包括自己gpu需要的版本
2. 修改include_directories和link_directories下的路径, 目前还是硬编码的10.1
3. 修改CMAKE_CXX-FLAGS, 加入-std=c++11, 因为nv使用的高版本gcc, 默认启动c++11, 而我们用的gcc-4.8.5默认还是c++03

这里CMAKE_CUDA_FLAGS如何修改与你的GPU型号有关, 一些文章详细讲了如何设置可以参考:

编译TensorRT-Bert

进入编译容器shell后,执行

mkdir build
cd build
cmake ..
make -j

在TensorRT/demo/BERT/build/下得到libbert_plugins.so和libbert_plugins.so两个文件

转换模型

TensorRT的模型文件叫plan, 可以从很多深度学习训练框架中转换得到, TensorRT-Bert已经提供了一个python的转换脚本bert_builder.py 用于转换ckpt模型文件, 其他格式需要自己去修改了.

一个例子调用

python3 bert_builder.py  -m model/model.ckpt-5474 -c model -s 384 -b 8 -o model.engine

这个命令使用batch_size为8, 序列长度为384的配置转换模型权重model.ckpt-5474, 得到tensorrt engine文件model.engine

TensorRT有Python API, 参考文档 User Guide

在python下运行推断

类似tensorflow中的session.run()方式, 在python中执行

一个例子, 其实就是bert_inference.py

TensorRT-Inference-Server

简称trtis, 要支持BERT的inference也需要使用最新版本, nvidia有提供docker镜像, 但编译是基于cuda10.1的, 所以为了兼容我们机器上的cuda10.0及驱动, 需要编译一个cuda10版本, 直接在tlinux2.2开发机上编译, 得到trtserver可执行文件后拷贝出来即可运行

为了避免处理tf/torch/onnx等依赖, 选择在cmake中把这些都屏蔽

cmake ../build -DCMAKE_BUILD_TYPE=Release  \
    -DTRTIS_ENABLE_METRICS=ON \
    -DTRTIS_ENABLE_GCS=OFF \
    -DTRTIS_ENABLE_S3=OFF \
    -DTRTIS_ENABLE_CUSTOM=ON \
    -DTRTIS_ENABLE_TENSORFLOW=OFF \
    -DTRTIS_ENABLE_TENSORRT=ON \
    -DTRTIS_ENABLE_CAFFE2=OFF \
    -DTRTIS_ENABLE_ONNXRUNTIME=OFF \
    -DTRTIS_ENABLE_PYTORCH=OFF

CMakeLists.txt写的不太好, 需要comment掉几个target才可以编译:

trtis-client
trtis-test-utils
aws-sdk-cpp
google-cloud-cpp
prometheus-cpp

另外需要确保有/usr/local/cuda且把tensorrt安装到ld.so.conf中, 版本和之前编译TensorRT-BERT一致即可

要运行trtis, 需要配置一下model文件, 主要就是之前得到的engine文件, 目录结构如下:

models/
`-- bert32
    |-- 1
    |   `-- model.plan
    `-- config.pbtxt

2 directories, 2 files

其中config.pbtxt文件定义了输入输出, max_batch_size等参数

name: "bert32"
platform: "tensorrt_plan"
max_batch_size: 1
input [
  {
    name: "input_ids"
    data_type: TYPE_INT32
    dims: [ 32 ]
  },
  {
    name: "input_mask"
    data_type: TYPE_INT32
    dims: [ 32 ]
  },
  {
    name: "segment_ids"
    data_type: TYPE_INT32
    dims: [ 32 ]
  }
]
output [
  {
    name: "cls_dense"
    data_type: TYPE_FP32
    dims: [ 32, 2, 1, 1 ]
  }
]

启动命令是

$ LD_PRELOAD=libbert_plugins.so trtserver  --model-repositor models/
I1010 09:52:05.779515 25314 server.cc:110] Initializing TensorRT Inference Server
I1010 09:52:05.841512 25314 server_status.cc:83] New status tracking for model 'bert32'
I1010 09:52:05.841606 25314 model_repository_manager.cc:668] loading: bert32:1
I1010 09:52:08.402921 25314 plan_backend.cc:239] Creating instance bert32_0_gpu0 on GPU 0 (6.1) using model.plan
W1010 09:52:22.904966 25314 logging.cc:46] TensorRT was linked against cuDNN 7.6.3 but loaded cuDNN 7.6.1
W1010 09:52:23.952128 25314 logging.cc:46] TensorRT was linked against cuDNN 7.6.3 but loaded cuDNN 7.6.1
I1010 09:52:23.971254 25314 plan_backend.cc:397] Created instance bert32_0_gpu0 on GPU 0 (6.1) with stream priority 0 and optimization profile 
I1010 09:52:24.041133 25314 model_repository_manager.cc:810] successfully loaded 'bert32' version 1
I1010 09:52:24.041217 25314 main.cc:417] Starting endpoints, 'inference:0' listening on
I1010 09:52:24.042710 25314 grpc_server.cc:1730] Started GRPCService at 0.0.0.0:8001
I1010 09:52:24.042728 25314 http_server.cc:1125] Starting HTTPService at 0.0.0.0:8000

根据Warning about Dynamic Batching and thread count #95但利用trtis做批处理/离线运算时记得把http thread count设大，否则dynamic batching会因为收不到足够的请求失效

参考了

Custom Operations: trtis启动时如何加载plugin
Model Repository: 模型文件如何摆放
Model Configuration: config.pbtxt文件的编写

andyfei的notebook

have fun