基于BERT，神马搜索在线预测性能如何提升？

原创丁亦川阿里云开发者

小叽导读：俗话说，天下没有免费的午餐，BERT 拥有出众效果的代价就是同样“出众”的资源消耗。今天，阿里巴巴技术专家丁亦川梳理了神马搜索 BERT 在线预测过程中，遇到的性能问题，探索过程以及当前初步的成果，分享给大家。

《阿里巴巴机器智能技术精选集》已正式推出，收录多篇 AAAI2018 优秀论文，并向业界开放下载。本书集中展示了阿里最新研究成果，包罗NLP、搜索推荐、社交网络、基础算法、深度学习等多个方向。点击下方按钮，立即阅读~

一、背景介绍

BERT 的全称是 Bidirectional Encoder Representation from Transformers，即双向 Transformer 的 Encoder，由 google 在 2018 年提出[1]。模型的主要创新点是，在海量语料上同时用 Masked LM 和 Next Sentence Prediction 两种任务做预训练，利用 transformer 中 self-attentnion 的特征抽取能力，全面获取词语和句子级别的 representation。

根据论文中数据， BERT 在多个 NLP 任务中都明显超过同时期的其他方法，它的里程碑意义在于：终于有一个 deep 的模型可以显著提高 NLP 任务的准确率，而且这个模型可以从无标记数据集中预训练得到。

二、BERT 性能之痛

在我们搜索场景中，对模型的高吞吐、低延迟有格外强烈的需求。以典型的搜索结果排序为例，一次检索会经过召回、粗排、精排等阶段，即使在最后的精排阶段，候选集也至少有 10 – 20 个文档，即 BERT 预测请求的平均 batch size 为 20，所以每个模型预测请求都有相当大的计算开销；同时搜索对延迟的要求也很严格，需要模型预测的延迟在 20ms 以内，这对模型服务的性能提出了挑战。

BERT 计算开销巨大，只能用 GPU 才能满足延迟需求，而需要 BERT 的业务场景远不止此，如果都靠堆机器来满足性能要求，硬件成本会远超预算，即使最终把模型从 12 层裁剪到 3 层，性能也远不及预期，因此我们对 BERT 模型的预测做了专项优化。

三、BERT 性能分析

首先，通过 tensorflow 的 Profile 工具，查看性能热点：

可以看到消耗主要由两部分组成：

1. 内存到显存的参数传输；
2. Transformer 计算；

更抽象地看，这两个问题本质就是显存和计算的管理问题，优化也是针对这两方面来实施。

我们先尝试了基于 tensorflow 的优化，这个方法受限于 tensorflow 沉重的框架，收益较小，不过分析的过程让我们更清晰地看到了问题所在；然后我们基于开源代码重写了 BERT 预测逻辑，获得了近一倍的性能提升，下面会分别介绍。由于 BERT 算法的主要开销在于 transformer 结构，所以下文叙述时就不特意区分 BERT 和 transformer 了。

四、BERT 性能优化

4.1 基于原生 tensorflow

■ 4.1.1 显存

结合性能热点和 tensorboard，可以发现传输的数据就是模型的各个参数，因为 tensorflow 加载模型用于预测时，会对计算图做切割，启发式[2]地将各个子图放到可用的计算设备上（典型的就是 CPU 和 GPU）。

要减少这部分开销也很简单，通过 tensorflow 的 device 机制将各个 OP 指定到 GPU 设备即可。

不过显存问题并没有完全解决，就像 CPU 程序运行过程中会产生内存碎片，显存也有一样的问题，tensorflow 实现了一套显存管理器[3]，不过运行过程中（尤其是请求的 batch 变化较大时），还是会看到延迟有抖动。

■ 4.1.2 计算

tensorflow 框架在实现算法时经常需要大量 op，依赖这种灵活的机制，tensorflow 能构造各种网络结构，而缺点也很明显：不但会引入节点间数据交互的开销，在 GPU 上运行时更会有大量的 cuda kernel launch，不利于 GPU 性能发挥。

通过 tensorboard 看到，BERT 的 transformer encoder 用了 568 个节点。

而且 Transformer 算法包含大量的矩阵计算，这些计算开销无法取巧绕过，tensorflow 本身对 GPU 上的量化计算支持也很差，除非定制融合 op，否则没有好的解决办法。

■ 4.1.3 收益

如上所述，我们通过设置 device 的方法将参数都放到显存中，经测试延迟下降了 1/3，但是 transformer 计算量没有变化，所以模型的吞吐没有显著提升；基于 tensorflow 的优化开发代价较小，但是效果也有限。

4.2 重写 BERT 预测逻辑

tensorflow 框架帮我们完成了显存的管理、计算节点的调度，减少使用者的负担，降低入门的门槛，但这个沉重的框架也是性能优化最大的束缚。对于 transformer 这种复用性很强的经典结构，我们就考虑轻装上阵，重新实现。

完全重写的开发工作量较大，幸运的是，知乎的开源项目[4]已经帮我们铺好了道路，这边再次感谢知乎同学的工作。Transformer 结构基本可以直接复用，我们需要做的主要是：

1. 解析 tensorflow 模型中的参数，适配到 cubert 中；
2. 实现 Transformer 后面的网络，常见的就是一个 MLP；
3. 对 cubert 略作修改，根据具体模型结构（矩阵 size）选择合适的 cuda/cublas 接口；

相比 tensorflow 模型两大性能问题：显存、计算，重新实现的 BERT 则几乎没有这些负担。

■ 4.2.1 显存

如 tensorflow 的优化方式，我们也可以在模型加载时就把需要的参数都直接加载到显存中，而且通过设置 max batch size，我们可以把计算过程中存储中间结果的显存全部申请好，实现了预测过程中显存分配的零消耗；

■ 4.2.2 计算

重新实现的 BERT 牺牲了 tensorflow 计算图的灵活性，好处是大幅减少了 kernel launch 的消耗（约原生 tensorflow 模型的 20%）。而且计算过程完全可控，很方便就能实现半精度计算，在支持 tensor core 的显卡（V100/T4）上也能充分利用硬件加速的红利；

■ 4.2.3 收益

吞吐提升一倍，同等压力下的延迟下降 50%，数据详见下一节。

五、性能对比数据

5.1 环境/BERT 模型参数

内核：3.10.0
gcc：4.9.2

transformer 层数：3
隐层大小：768
序列长度：64
请求 batch size：20

5.2 性能数据