深入淺出騰訊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。
更多精彩文章，歡迎掃碼關注下方的公眾號 ~~ 歡迎關注和點贊，你的鼓勵將是我創作的動力
歡迎轉發至朋友圈，公眾號轉載請後臺留言申請授權～
