「來源： ｜GiantPandaCV ID：BBuf233」
引言
Transformer，Bert模型在NLP領域取得巨大成功，得到廣泛應用。而Transformer系列模型大小通常很大，在應用層提供相應服務是一個巨大的挑戰。
位元組提出的lightseq是一款高效能訓練，推理庫，包含了各種GPU最佳化技術，並且能夠很好相容tf/torch的模型實現。相比Tensorflow原生實現能達到14倍加速，比英偉達出品的FasterTransformer能夠加速1。4倍。
論文：https：//arxiv。org/abs/2010。13887
倉庫地址：https：//github。com/bytedance/lightseq
介紹
類似LightSeq的高效能加速庫也有很多，下面的三個主要特性是我們比別的加速庫表現好的原因：
我們將Tensorflow/Pytorch實現中的一些細粒度Kernel，進一步融合實現成一個粗粒度的Kernel，從而避免大量核函式啟動和GPU memory IO帶來的時間成本
我們設計了一種hierarchical（層級） auto regressive search來替代auto regressive search，進一步加速
我們提出了一種動態視訊記憶體複用策略，在NLP處理中，我們經常會遇到變長資料，給記憶體分配帶來了困難。LightSeq預先定義了每個kernel最大可使用視訊記憶體，並給不存在依賴關係的kernel進行共享，能夠減少8倍記憶體分配。
方法
Transformer主要分為兩部分，特徵計算層和輸出層。特徵計算層就是自注意力機制+FFN這些，而輸出層則會隨著任務不同而有些許改變。在NLU上是分類，而在NLG上是搜尋（用beam search）。
我們做了一下三個最佳化來分別解決前面提到的問題
kernel fusion
簡單來說就是將多個kernel，融合進一個大kernel。
以LayerNormalization（tf版本）為例子，在XLA最佳化過後，仍然需要3次kernel launch和2個臨時變數儲存（mean和variance）。
我們可以藉助cuda來自定義一個完整的LayerNorm Kernel。（程式碼地址：https：//github。com/bytedance/lightseq/blob/master/lightseq/training/csrc/kernels/normalize_kernels。cu#L35-L77）
template
__global__ voidker_layer_norm（T *ln_res， T *vars， T *means， const T *inp，
const T *scale， const T *bias， int hidden_size）{
// step 0。 compute local sum
float l_sum = 0；
float l_square_sum = 0；
const float4 *inp_f4 = （const float4 *）inp + blockIdx。x * hidden_size；
for （uint idx = threadIdx。x； idx < hidden_size； idx += blockDim。x） {
float4 val = inp_f4［idx］；
l_sum += val。x + val。y + val。z + val。w；
l_square_sum +=
val。x * val。x + val。y * val。y + val。z * val。z + val。w * val。w；
}
首先以float4的形式讀入資料（即一次性讀4個float），然後分別計算local sum和square sum。這裡取四個資料是以 val。x， val。y， val。z， val。w來取。計算square sum是用於後面var的計算。
接著是做reduce sum操作了
// step 1。 compute reduce sum
float mean_dim = float（hidden_size） * 4。f；
float reduce_val［2］ = {l_sum， l_square_sum}；
blockReduce
__shared__ float s_mean， s_var；
if （threadIdx。x == 0） {
s_mean = reduce_val［0］ / mean_dim；
if （means ！= nullptr） {
means［blockIdx。x］ = s_mean；
}
s_var = reduce_val［1］ / mean_dim - s_mean * s_mean + LN_EPSILON；
vars［blockIdx。x］ = s_var；
s_var = rsqrtf（s_var）；
}
__syncthreads（）；
分別對sum，和square_sum做一次reduce操作。然後在0號執行緒上，計算得到mean和var。
這裡的var用的是公式
得到，然後進行同步。注意這裡的s_mean和s_var是一個shared變數，可以被同一個block內的其他執行緒得到。
最後就是減均值，除方差，拉伸變換這部分操作
// step 2。 layer norm result
float4 *output_f4 = （float4 *）ln_res + blockIdx。x * hidden_size；
for （uint idx = threadIdx。x； idx < hidden_size； idx += blockDim。x） {
float4 vscale = __ldg（（const float4 *）scale + idx）；
float4 vbias = __ldg（（const float4 *）bias + idx）；
float4 val = inp_f4［idx］；
val。x = （val。x - s_mean） * s_var * vscale。x + vbias。x；
val。y = （val。y - s_mean） * s_var * vscale。y + vbias。y；
val。z = （val。z - s_mean） * s_var * vscale。z + vbias。z；
val。w = （val。w - s_mean） * s_var * vscale。w + vbias。w；
output_f4［idx］ = val；
}
kernel fusion這種最佳化也是在別的加速庫很常見，nvidia出品的Megatron也有做layernorm，mask+softmax，triangle（取上三角）+mask+softmax的融合，具體可參考：https：//github。com/NVIDIA/Megatron-LM/tree/main/megatron/fused_kernels
Hierarchical Auto Regressive Search
auto regressive search方法如beam seach， diverse beam search等方法計算過於複雜。通常我們只需要置信度較高的幾個label/tokens，而不是所有token都需要參與最終輸出的計算。
我們簡單回顧下beam search做法：
使用softmax計算，並將結果寫入gpu memory
讀取結果，並從中選取top-K的beams和tokens
而lightseq受啟發於推薦系統的檢索+重排，採用了兩階段策略：
隨機將logits組分到k個組
計算每個group（）的最大值，記為
在中，計算最小值，可以看作是logits粗略的top-k值
選擇值大於R的logits，並寫入到gpu memory
為什麼這種做法是粗略的top-k呢？假設我們有［1， 2， 3， 4， 5， 6］，我想要取top2，那這裡應該是6。按照lightseq做法，我們先隨機分兩組，如果是這種情況［1， 5， 6］， ［2， 3， 4］。那麼每組最大值分別是6， 4。然後取最小值，得到top2是4。所以只能看作是一種粗略的topk
檢索+重排動態視訊記憶體複用
LightSeq預先定義好動態shape的最大長度，在一開始先分配好最大視訊記憶體，此外GPU視訊記憶體將共享給不存在依賴關係的中間結果。
我理解這種做法是用於變長資料中，因為以往遇到變長資料，我們都會統一padding到一個固定長度，在最後計算loss也是加上一個padding mask。使用這種做法就能節省視訊記憶體。
另外位元組也出品了一個Effective Transformer，也是解決變長資料問題。透過對mask求一個字首和，在計算attention前後進行相應的刪除/恢復padding。
具體可參考作者的知乎文章（https：//www。zhihu。com/search？type=content&q=Effective%20Transformer）相關程式碼在（https：//github。com/bytedance/effective_transformer）
結果
profile資料
經過一系列最佳化後，在lightseq上，GEMM通用矩陣乘能夠佔大部分計算，計算效率更高。
與其他Transformer庫比較