深入浅出腾讯BERT推理模型--TurboTransformers

Overview

TurboTransformers是腾讯最近开源的BERT推理模型,它的特点就是一个字,快。本人用BERT(huggingface/transformers)在V100上做了测试,测试结果和官宣的基本一致:TurboTransformers的推理速度要比Pytorch快上1~4倍。

图片来源:https://github.com/Tencent/TurboTransformers

它之所以快,是因为它是专用于BERT的轻量级推理模型。

分层

不管是计算机的硬件、软件,还是现在的深度学习,它们都遵循着一个很重要的设计思想--分层:

  • 用简单的代码(或电路)来实现一个基本功能组件。
  • 用几个基本组件组合成一个功能更强的复杂组件。
  • 从简单到复杂,像搭积木一样,一层层地搭建出拥有很强功能的组件。

开发者只需要基于PyTorch的几个基本组件就能搭建出BERT模型,而且这些组件本身对他们来说都是透明的。正因如此,PyTorch才越来越受到研究者青睐。

BERT模型的基本组件

分层设计的优点很多,例如,可以简化问题、降低创新门槛、加速开发等,但它的缺点也很明显:

  • 流程固定化
  • 存在中间层延迟

深度神经网络里有个经典套路:一个激活函数层后面紧跟着一个dropout层。PyTorch需要lanuch两个GPU kernel程序来完成这两步计算。

F.dropout(F.relu(x))

实际上,这两项计算都是element-wise的,是可以合并成一个kernel的。但目前来说,不管是PyTorch,还是其他的通用训练框架,它们都很少有提供这种融合计算的API。

至于中间层延迟,最经典的要属“hello world”程序。虽然只有几行代码,但实际上要经过的中间层数根本数不过来。

你可以阅读深入浅出PyTorch(算子篇)来了解下矩阵相乘这个最基本的计算在PyTorch里要经过多少个中间层。

分层展开

要想将程序的低延迟最大化,就需要把分层的代码完全展开,并重构代码。典型例子就是嵌入式系统,为了实现某种需求,它可以打破应用程序、程序库、操作系统甚至是硬件设备的界限,打造一个软硬件一体化产品。

这种分层展开的设计模式当然也有它的局限性:专用。由于高度定制化,它通常只能用于完成某个特定功能。低延迟和专用化是呈绝对的正相关的。

TurboTransformers就是采用这种设计:只实现BERT模型前向传播所需要的算子,并融合那些可以合并的算子。

turbo.Tensor

首先,它用CUDA开发了一个轻量级的tensor计算库,所谓的轻量级,指的是不用考虑反向传播、稀疏矩阵等操作,只实现BERT前向传播所必需的operator。

虽然tensor库是用C++写的,但考虑到python在AI开发中的地位,它用pybind11将C++ API暴露给前端的python Tensor类。

# turbo_transformers/python/pybind.cpp
 72   py::class_<core::Tensor>(m, "Tensor")                      
 73       .def_static("from_dlpack",
 74                   [](py::capsule capsule) -> std::unique_ptr<core::Tensor> {
 75                     auto tensor = (DLManagedTensor *)(capsule);
 76                     PyCapsule_SetName(capsule.ptr(), "used_tensor");
 77                     return absl::make_unique<core::Tensor>(tensor);
 78                   })
 79       .def("to_dlpack",
 80            [](core::Tensor &tensor) -> py::capsule {
 81              auto *dlpack = tensor.ToDLPack();                    
 82              return py::capsule(dlpack, "dltensor", DLPack_Capsule_Destructor);
 83            })
 84       .def("n_dim", &core::Tensor::n_dim)
 85       .def("shape", &core::Tensor::shape)

从预训练模型(PyTorch)那迁移参数时,turbo.Tensor不能直接对接torch.Tensor,需要先将PyTorch的参数转成dlpack格式, 再通过from_dlpack()将这些数据导入生成TurboTransformers tensor。除了dlpack之外,还支持*.npz文件格式。

图片来源:https://github.com/Tencent/TurboTransformers

turbo.xxxlayer

TurboTransformers用CUDA重构了Embedding、self-attention、intermediate、output、LayerNorm和pooler等layer。turbo.layer不仅代码结构简洁,overhead少,还合并了一部分算子。

图片来源:https://www.oschina.net/news/115146/tencent-wechat-opensource-turbotransformers

这里以intermediate layer为例,来分析这些算子的特点。

turbo_transformers/layers/bert_intermediate.cpp

intermediate layer的实现比较简单:一个Linear layer后面紧跟着一个gelu activation layer。

PyTorch的intermediate layer的会lanuch 3个kernel来完成这部分计算:

  • #1: y = input.matmul(weight)
  • #2: y = y + bias
  • #3: y = gelu(y)

由于#2和#3都是element-wise kernel,turbo把它们进行了融合--AddBiasAct(),相同的计算操作,只需要lanuch 2个kernel,计算速度当然更快。

turbo_transformers/layers/kernels/mat_mul.cpp

和PyTorch一样,turbo的MatMul算子也是调用cuBLAS来进行矩阵运算,而且turbo还启用了Tensor Core来加速计算(CUBLAS_TENSOR_OP_MATH)。

总结

到此,本文基本上讲清了TurboTransformers的速度优势来源,由于篇幅所限,不能分析所有的算子。BERT的核心模块是self-attention,如果想了解更多,可以阅读深入浅出Transformer


更多精彩文章,欢迎扫码关注下方的公众号 ~~ 欢迎关注和点赞,你的鼓励将是我创作的动力

欢迎转发至朋友圈,公众号转载请后台留言申请授权~

AI修炼手册:一个有深度的公众号

posted @ 2020-06-27 09:20  A君来了  阅读(1821)  评论(0编辑  收藏  举报