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。
更多精彩文章,欢迎扫码关注下方的公众号 ~~ 欢迎关注和点赞,你的鼓励将是我创作的动力
欢迎转发至朋友圈,公众号转载请后台留言申请授权~