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

简介: 今天,阿里巴巴技术专家丁亦川梳理了神马搜索 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 工具,查看性能热点:

image.png

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

  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 设备即可。

image.png

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

■ 4.1.2 计算

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

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

image.png

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

image.png

5.3 分析

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

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

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

六、后续工作

  1. 挖掘新型号显卡(T4) 的性能,尝试半精度、int8 等。预测时对精度要求较低,而 T4 对低精度计算的支持较好,预期性能还可以大幅提升;
  1. 尝试知识蒸馏等小模型替换大模型。这方面还需要与各业务同学多多合作,发挥模型的最大性能。

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

引用资料

[1]https://arxiv.org/abs/1810.04805

[2]https://github.com/tensorflow/tensorflow/blob/v1.8.0/tensorflow/core/graph/graph_partition.cc#L938

[3]https://github.com/tensorflow/tensorflow/blob/v1.8.0/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h#L38

[4 ] https://github.com/zhihu/cuBERT

相关实践学习
部署Stable Diffusion玩转AI绘画(GPU云服务器)
本实验通过在ECS上从零开始部署Stable Diffusion来进行AI绘画创作,开启AIGC盲盒。
目录
相关文章
|
4月前
|
机器学习/深度学习 人工智能 算法
没想到!AlphaZero式树搜索也能用来增强大语言模型推理与训练
【7月更文挑战第26天】Xidong Feng等研究人员提出了一项创新方法,通过采用AlphaZero式的树搜索算法来增强大语言模型(LLMs)的推理与训练能力。这项技术,称为TS-LLM(Tree-Search for LLMs),将LLMs的解码过程视为搜索问题,并运用AlphaZero的树搜索来指导这一过程。TS-LLM不仅提升了模型的通用性和适应性,还在多个任务中实现了显著的性能提升。此外,它能在训练阶段指导LLMs学习更优的解码策略。尽管如此,TS-LLM依赖于高质量的预训练LLM,并面临较高的计算成本挑战。[论文](https://arxiv.org/abs/2309.17179)
74 5
|
4月前
|
机器学习/深度学习 自然语言处理 并行计算
淘宝搜索中的深度语义模型:从理论到实践
淘宝搜索系统通过引入深度语义模型,极大地提升了搜索质量和用户体验。这些模型不仅能够准确理解用户的需求,还能够智能地匹配和推荐商品,为用户提供了一个更加便捷、个性化的购物环境。随着技术的不断发展和完善,淘宝搜索将会变得更加智能和高效。
|
5月前
|
人工智能 自然语言处理 算法
大模型+蒙特卡洛树搜索,一招让LLaMa-3 8B奥数水平直逼GPT-4
【6月更文挑战第25天】 - 复旦大学和上海AI Lab的研究者提出这一算法,用于增强大型语言模型在复杂数学推理任务中的能力,解决现有模型推理准确性问题。 - **MCTSr**流程包括初始化、选择、自细化、自评估、反向传播和UCT更新,通过多轮迭代提升答案质量。 - 实验显示,该算法在**GSM8K**、**GSM Hard**、**MATH**和**Olympiad-level**数据集上表现出色,尤其在多次迭代后。 - 尽管计算成本高且不适用于所有问题类型,但研究揭示了强化LLMs推理能力的新途径,对未来的AI应用具有指导意义。
187 8
|
5月前
|
机器学习/深度学习 自然语言处理 算法
用神经架构搜索给LLM瘦身,模型变小,准确度有时反而更高
【6月更文挑战第20天】研究人员运用神经架构搜索(NAS)压缩LLM,如LLaMA2-7B,找到小而精准的子网,降低内存与计算成本,保持甚至提升性能。实验显示在多个任务上,模型大小减半,速度加快,精度不变或提升。NAS虽需大量计算资源,但结合量化技术,能有效优化大型语言模型。[论文链接](https://arxiv.org/pdf/2405.18377)**
55 3
|
机器学习/深度学习 算法 文件存储
ICLR2021 | 从50篇投稿,看模型搜索最新进展
他们提出一种称为FTSO的方法,通过分别搜索网络拓扑和运营商,将NAS的搜索时间从几天缩短至0.68秒,同时在ImageNet上实现76.42%的测试精度,在CIFAR10上实现97.77%的测试精度。
ICLR2021 | 从50篇投稿,看模型搜索最新进展
|
机器学习/深度学习 编解码 TensorFlow
MnasNet架构解析与复现-神经架构搜索
为移动设备设计卷积神经网络 (CNN) 具有挑战性,因为移动模型需要小而快,但仍要准确。尽管在所有维度上都致力于设计和改进移动 CNN,但当需要考虑如此多的架构可能性时,很难手动平衡这些权衡。在本文中,我们提出了一种**自动移动神经架构搜索 (MNAS) 方法**,该方法明确地将模型延迟纳入主要目标,以便搜索可以识别出在准确性和延迟之间取得良好折衷的模型。与之前的工作不同,延迟是通过另一个通常不准确的代理(例如 FLOPS)来考虑的,我们的方法通过在手机上执行模型来直接测量现实世界的推理延迟。为了进一步在灵活性和搜索空间大小之间取得适当的平衡,我们**提出了一种新颖的分解层次搜索空间,它鼓励整
541 0
MnasNet架构解析与复现-神经架构搜索
【算法提高——第二讲】搜索(3)
【算法提高——第二讲】搜索(3)
【算法提高——第二讲】搜索(3)
【算法提高——第二讲】搜索(1)
【算法提高——第二讲】搜索(1)
【算法提高——第二讲】搜索(1)
【算法提高——第二讲】搜索(2)
【算法提高——第二讲】搜索(2)
【算法提高——第二讲】搜索(2)
|
机器学习/深度学习 自然语言处理 搜索推荐
深度语义模型以及在淘宝搜索中的应用
传统的搜索文本相关性模型,如BM25通常计算Query与Doc文本term匹配程度。由于Query与Doc之间的语义gap, 可能存在很多语义相关,但文本并不匹配的情况。为了解决语义匹配问题,出现很多LSA,LDA等语义模型。
11399 0
下一篇
无影云桌面