谷歌搜索百度推广网Google搜索360好搜狗神马百度搜索引擎€

简介： 今天，阿里巴巴技术专家丁亦川梳理了神马搜索 BERT 在线预测过程中，遇到的性能问题，探索过程以及当前初步的成果，分享给大家。

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

----

一、背景介绍

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 性能数据

5.3 分析

压测过程中发现，有时 qps 上升，延迟反而下降，因为每个请求都要 GPU 多核并行计算，于是存在竞争的风险；

理想情况下，请求均匀到达，相互之间没有竞争，能保持低延迟打满 GPU；

如果在某个时间点同时到达多个请求，就会相互抢夺 GPU 资源，导致所有请求的延迟都明显上涨。

六、后续工作

1. 挖掘新型号显卡（T4） 的性能，尝试半精度、int8 等。预测时对精度要求较低，而 T4 对低精度计算的支持较好，预期性能还可以大幅提升；

2. 尝试知识蒸馏等小模型替换大模型。这方面还需要与各业务同学多多合作，发挥模型的最大性能。

求贤若渴：随着 BERT 算法出现，可以看到，在工程/算法的结合点我们还大有可为。我们团队的工作涉及搜索引擎、推荐引擎、模型预测、向量检索等，可以深度参与搜索、信息流推荐业务，期待各位算法或架构大牛加盟我们团队，联系邮箱：yichuan.dingyc@alibaba-inc.com

更多文章核心内容：点击下方“了解更多”

原文出处：阿里云大学开发者社区