深入淺出騰訊BERT推理模型--TurboTransformers

A君來了發表於2020-06-27
模型ORM

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修煉手冊：一個有深度的公眾號

相關文章