3倍加速CPU上的BERT模型部署

2020 年 7 月 13 日 ApacheMXNet

BERT模型 [1] 是由Google在2018年提出,由于其在各类NLP的任务上出色性能,一经推出便获得了很多的关注。但是BERT模型的缺点是计算量过于庞大,对于推理部署有着很大的挑战。最近我们看到越来越多的产品开始从一些轻量级的模型转向使用BERT,所以BERT模型的部署性能变得越来越重要。

最近,在学术界和工业界,大家都在加紧研发深度学习编译器用于优化模型部署的性能,比如TVM [2], MLIR [3], 和Glow [4]。并且一些近期的成果[2,5] 也证明了深度学习编译器可以有效地降低模型推理的时间。在AWS,我们在积极地对开源项目Apache TVM做开发和贡献,并且将TVM用于优化各种平台下的模型部署,其中也包括了BERT模型。同时,我们提供了一个非常简单易用的模型优化服务——Amazon SageMaker Neo。

在这篇文章中,我们想展示一下我们最新的BERT模型在CPU上部署优化的成果,并且手把手教你如何复现这一优化结果。我们可以在Amazon EC2 c5.9xlarge的服务器上将BERT模型的推理时间降低高达2.9倍,以及取得至多2.3倍更高的吞吐量。我们得到的fp32的BERT模型推理性能是目前已知在CPU上最好的。TVM之所以可以取得高性能主要源于三个点:

  • TVM将小算子融合在一起,可以大大提高缓存的使用效率

  • TVM既可以自己代码生成高性能的算子,又可以利用高性能的第三方库

  • TVM能够对计算图进行优化,同时也可以将一些计算速度较慢的数学运算替换成更快的近似计算

更重要的是,这一优化效果只需要十几行代码便可完成,完整复现代码文末可见,欢迎试用。另外,我们所有的优化代码也已全部开源。


BERT模型优化效果

我们使用了BERT base模型和DistilBERT模型来展示我们的优化性能。其中BERT base模型有12层transformer层,而DistilBERT只有6层,因此DistilBERT在推理时运行得更快。我们使用了GluonNLP (0.9.1)的模型库中提供的这两个模型,然后用TVM对其进行编译优化。所有的模型都使用fp32的精度。


下表比较了MXNet(mkl 1.6.0)和TVM优化过后的两个模型在c5.9xlarge上的运行时间。我们可以看到TVM可以降低2.1倍到2.9倍的运行时间。值得注意的是,DistilBERT模型在序列长度128的时候,在CPU上只需要9.5ms的推理时间。相比之下,之前ONNX runtime在类似的CPU上将BERT模型减少到只有3层才取得了9ms的推理时间。相对于TensorFlow和PyTorch实现的BERT模型,我们在CPU上也有相应的性能加速。

model batch seq length MXNet lat. (ms) TVM lat. (ms) speedup
BERT 1 64 26.1 12.6 2.1
BERT 1 128 45.8 19.2 2.4
BERT 1 256 99.4 35.3 2.8
DistilBERT 1 64 13.4 6.2 2.2
DistilBERT 1 128 23.2 9.5 2.5
DistilBERT 1 256 50.1 17.5 2.9

同时我们也在更多的EC2云服务器上(c5.2xlarge, c5.9xlarge, m5.2xlarge, m5.12xlarge)评测了BERT模型在batch 4和序列长度128下的吞吐量。下图显示TVM能够在各类型的服务器上将BERT和DistilBERT模型取得平均2倍更高的吞吐量。


如何复现这一结果?

接下来我们将一步步地教你如何来复现上面的优化结果。

首先,我们需要在Amazon EC2上使用AWS Deep Learning AMI, Ubuntu 18.04来创建一个CPU服务器。在我们ssh到这个服务器后,我们先要将conda环境切换成mxnet_p36,然后安装GluonNLP的python模块。

source activate mxnet_p36
pip install gluonnlp==0.9.1


然后,我们需要安装TVM。这里提供的TVM python模块是针对Deep Learning AMI编译的,有可能无法兼容其它的系统和运行环境。我们会在附录教大家如何手动编译及安装TVM。

pip install https://tvm-build-public.s3-us-west-2.amazonaws.com/dlami-cpu-mkl/tvm-0.7.dev1-cp36-cp36m-linux_x86_64.whl
pip install https://tvm-build-public.s3-us-west-2.amazonaws.com/dlami-cpu-mkl/topi-0.7.dev1-py3-none-any.whl


接下来,我们将使用TVM对BERT模型进行优化。在这里我们省去了如何用GluonNLP来训练和精调BERT模型,你可以参考GluonNLP网站上的教程。在我们载入BERT模型后,我们首先将BERT模型从MXNet转换到TVM Relay语言。Relay语言是TVM中用于表示模型的计算图的一种中间表达。

# Load model from GluonNLP
...
mx_model = nlp.model.BERTClassifier(...)

# Convert MXNet model to TVM
shape_dict = {
'data0': (batch, seq_length),
'data1': (batch, seq_length),
'data2': (batch,)
}
mod, params = relay.frontend.from_mxnet(mx_model, shape_dict)


在编译模型时,我们需要定义编译目标和优化级别。EC2上的c5和m5服务器的CPU都支持Intel的AVX-512指令,这类指令通过向量化来加速浮点运算。因此,我们在编译目标中加入了CPU架构(-mcpu=skylake-avx512)来使用这类指令。此外,我们也使用了Intel MKL的库来加速矩阵乘法等算子。第三,我们在编译时打开了一个优化选项叫“FastMath”。这个优化告诉TVM来用近似计算去替代一些较慢的数学运算,如erf。经过我们的验证,使用这类近似计算并不会影响最终的准确性。

target = "llvm -mcpu=skylake-avx512 -libs=cblas"
with relay.build_config(opt_level=3, required_pass=["FastMath"]):
graph, lib, cparams = relay.build(mod, target, params=params)


最后我们创建一个轻量级的TVM运行时来执行编译优化后的BERT模型,并检测它的输出是否和MXNet一致。

ctx = tvm.cpu()
rt = runtime.create(graph, lib, ctx)
rt.set_input(**cparams)
rt.set_input(data0=inputs, data1=token_types, data2=valid_length)
rt.run()
out = rt.get_output(0)
print(out.ansumpy())

# verify the correctness
tvm.testing.assert_allclose(out.asnumpy(), mx_out.asnumpy(), rtol=1e-3, atol=1e-3)


完整的复现代码可见于

https://gist.github.com/icemelon9/860d3d2c9566d6f69fa8112840dd95c1.


结论

综上,我们看到TVM可以显著地加速BERT模型在CPU上的推理性能。运行时间可以被降低至多2.9倍,而吞吐量可以被提高至2.3倍。值得注意的是,现在的这个BERT解决方案只适用于固定长度,我们还在继续开发TVM来支持动态的序列长度。


附录:从源代码编译TVM

TVM官方网站提供了一份详细的编译TVM指南。这里我们简单描述一下我们所使用的编译配置,以获得最佳的BERT推理性能。

首先,除了安装编译TVM所必须的依赖库以外,你还需要安装LLVM(>=6.0)和Intel MKL库(安装教程可见这个链接)。接下来,我们就可以编译TVM的C++库。

git clone --recursive https://github.com/apache/incubator-tvm.git tvm
cd tvm && mkdir build && cd build
cmake -DUSE_LLVM=/path/to/llvm-config -DUSE_BLAS=mkl -DUSE_OPENMP=intel ..
make -j4


然后,我们便可以安装TVM和Topi的python模块。

cd python; python setup.py install --user; cd ..
cd topi/python; python setup.py install --user; cd ../..


引用

[1] Jacob Devlin, et al. “BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding”. arXiv pre-print arXiv:1810.04805 (2018).
[2] Tianqi Chen, et al. “TVM: An Automated End-to-End Optimizing Compiler for Deep Learning.” 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18), 2018.
[3] Chris Lattner, et al. “MLIR: A Compiler Infrastructure for the End of Moore’s Law”. arXiv preprint arXiv:2002.11054, 2020.
[4] Nadav Rotem, et al. “Glow: Graph Lowering Compiler Techniques for Neural Networks”. CoRR, abs/1805.00907, 2018.
[5] Liu, Yizhi, et al. "Optimizing CNN Model Inference on CPUs."2019 USENIX Annual Technical Conference (USENIX ATC 19). 2019.
[6] Sanh, Victor, et al. "DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter." arXiv preprint arXiv:1910.01108 (2019).


登录查看更多
11

相关内容

专知会员服务
111+阅读 · 2020年8月22日
专知会员服务
77+阅读 · 2020年6月20日
【Amazon】使用预先训练的Transformer模型进行数据增强
专知会员服务
56+阅读 · 2020年3月6日
注意力机制模型最新综述
专知会员服务
259+阅读 · 2019年10月20日
深度神经网络模型压缩与加速综述
专知会员服务
125+阅读 · 2019年10月12日
社区分享|如何让模型在生产环境上推理得更快
TensorFlow 2.0 分布式训练
TensorFlow
8+阅读 · 2020年1月19日
英伟达Faster Transformer:作者带你揭秘BERT优化
机器之心
13+阅读 · 2019年9月18日
【GitHub】BERT模型从训练到部署全流程
专知
34+阅读 · 2019年6月28日
Colab 免费提供 Tesla T4 GPU,是时候薅羊毛了
机器之心
10+阅读 · 2019年4月25日
【强烈推荐】浅谈将Pytorch模型从CPU转换成GPU
机器学习研究会
7+阅读 · 2017年12月24日
TResNet: High Performance GPU-Dedicated Architecture
Arxiv
7+阅读 · 2020年3月30日
Arxiv
11+阅读 · 2019年6月19日
Star-Transformer
Arxiv
5+阅读 · 2019年2月28日
Meta-Learning with Latent Embedding Optimization
Arxiv
6+阅读 · 2018年7月16日
VIP会员
相关VIP内容
专知会员服务
111+阅读 · 2020年8月22日
专知会员服务
77+阅读 · 2020年6月20日
【Amazon】使用预先训练的Transformer模型进行数据增强
专知会员服务
56+阅读 · 2020年3月6日
注意力机制模型最新综述
专知会员服务
259+阅读 · 2019年10月20日
深度神经网络模型压缩与加速综述
专知会员服务
125+阅读 · 2019年10月12日
相关资讯
Top
微信扫码咨询专知VIP会员