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

Overview

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

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

分层

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

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

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

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

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

深度神经网络里有个经典套路：一个激活函数层后面紧跟着一个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文件格式。

turbo.xxxlayer

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

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

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，计算速度当然更快。

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

总结

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

---

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

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