【NLP】Fastertransformer源碼解讀

最近拜讀了NVIDIA前陣子開源的fastertransformer，對CUDA編程不是很熟悉，但總算是啃下來一些，帶你們讀一下硬核源碼。

1. 簡介

英偉達公衆號推送的文章加上配圖其實已經把該要講的很清楚了，主要有如下幾方面：

1. 爲了減小kernel調用次數，將除了矩陣乘法的kernel都儘量合併
2. 針對大batch單獨進行了kernel優化
3. 支持選擇最優的矩陣乘法
4. 在使用FP16時使用half2類型，達到half兩倍的訪存帶寬和計算吞吐
5. 優化gelu、softmax、layernorm的實現以及選用rsqrt等

不瞭解底層的同窗可能不是很懂，沒事我剛看到的時候也不懂，也不敢問，強擼一下源碼就通透（fang qi）了

2. 硬核源碼解讀

首先簡略說一下第一點優化。Kernel在tensorflow裏的概念是operation的計算實現，在cuda裏是執行一個線程的函數，也是一次計算，只不過tensorflow的更加宏觀些。每次tensorflow執行一個operation，都要調用對應的OpKernel，試想一個經過TF實現的transformer，有將近60個operation，計算一次要執行60次上述過程，進行頻繁的GPU調度和顯存讀寫。所以fastertransformer儘量多地對kernel進行了合併。

2.1 總體結構

Fastertransformer目錄下主要有（如下簡稱FTF）：

1. fastertransformer：主要源碼
1. cuda：優化後的kernel以及對multi-head attention總體的封裝（沒過線性層）
2. tf_op：tensorflow operation和OpKernel的註冊（op理解爲聲明、Opkenerl是定義）
3. trt_plugin：tensorRT的實現（能夠支持multi streaming太讚了）
4. bertencodertransformer.h：transformer總體的封裝
2. sample：cpp、tensorflow、tensrflow_bert、tensorRT的調用FTF的示例
3. tools：根據參數選擇最優的矩陣乘法（GEMM=General Matrix Multiplication）

接下來我主要想講一下1.1的細節，1.2能夠參考這位大佬的文章，剩下的代碼可讀性很強，基本讀一兩遍就知道了。

除去矩陣乘法，做者把剩下的op合併成了4個（圖中藍色框）：

這四個op的cuda源碼分別在cuda_kernels.cu和open_attention.cu兩個文件中，接下來研究下每一個op。

2.2 add_QKVbias (open_attention.cu)

在FP32時，每一個block負責處理一個word(num_head*head_size)的add bias運算，先找到要處理QKV中的一個，再進行運算，由於要transpose，因此把結果存入[bsz, num_head, seq_len, head_size]的矩陣裏。

在FP16是每一個block同時處理QKV上的同一個word（多是由於fp16計算的更快一些），在實際的計算中把half都轉成了half2計算。add的話直接用封裝好的__hadd2運算。使用half2計算的緣由原文說的比較清楚：

針對半精度FP16，咱們對各個kernel也進行了相應優化。首先，在kernel的實現中，將輸入的half指針轉成half2類型，並使用了half2相關的數學函數。這樣不只僅能夠達到2倍於half的訪存帶寬和計算吞吐，還能夠極大地減小指令的發射數量。其次，在SoftMax以及Layer Normalization的操做中，爲防止求和溢出，將數據以half2的形式讀入後，會轉成float2類型，來作求和計算。
-- NVIDIA BERT推理解決方案Faster Transformer開源啦

2.3 softmax_kernel (open_attention.cu)

在計算softmax以前對block線程數進行了區間處理，由於block裏的線程數最好是wrap大小(32)的倍數，提升計算效率。

調用kernel以前，會根據batch_size * head_num選擇不一樣的softmax kernel，主要是爲了保證在大batch的狀況下的計算效率，這裏爲何使用120我也不是很清楚，但願懂的朋友助力一下

if(batch_size * head_num <= 120)
    {
      grid.x = batch_size * head_num * seq_len;
      softmax_kernel_v2<DataType_><<<grid, block, 0, stream>>>(qk_buf_, attr_mask, batch_size, head_num, seq_len, scaler); 
    }
    else
    {
      grid.x = batch_size * head_num;
      softmax_kernel<DataType_><<<grid, block, 0, stream>>>(qk_buf_, attr_mask, batch_size, head_num, seq_len, scaler); 
    }複製代碼

在算softmax時，分母有個求和操做，用到了經典的parallel reduce算法，能夠仔細看看參考，講的比較清楚。

這裏注意，使用最第一版源碼的同窗們須要照着實現一個blockReduceMax，以防止數值溢出，softmax嚴謹的實現應該是：

def softmax(x):
"""Compute the softmax of vector x."""
    exp_x = np.exp(x)
    softmax_x = exp_x / np.sum(exp_x)
    return softmax_x複製代碼

2.4 transpose (open_attention.cu)

這裏要transpose回[bsz, seq_len, num_head, head_size]的矩陣。由於c++裏面矩陣是行優先存儲，只要按順序乘過來就行了（最開始看的有點暈）。

2.5 add_bias_act & add_bias_input_layernorm (cuda_kernels.cu)

若是前面幾個函數啃下來了，這兩個就比較好懂，主要的優化點是：

1. x^3 -> x*x*x: c語言中x*x和pow(x,2)哪一個計算更快一點？
2. rsqrt: Why is SSE scalar sqrt(x) slower than rsqrt(x) * x?
3. 還有就是各類half2運算的使用

2.6 gemm (tools/gemm_test)

矩陣乘法根據fp16和fp32的不一樣在不一樣的cublas算法中選擇，選擇後記錄到http://gemm_config.in文件中：

問了下做者，其實fp32也可使用CUBLAS_GEMM_ALGO0_TENSOR_OP到CUBLAS_GEMM_ALGO15_TENSOR_OP的算法，只不過存在一些風險（使用後速度提高2倍）。

2.7 trt_plugin

做者額外封裝了一個tensorRT的層，tensorRT主要是經過engine，在給定的context和stream下進行異步計算，提供了multi stream inference的可能。關於TensorRT的異步編程推薦一個英偉達的PPT：

CUDA C/C++ Streams and Concurrency​ developer.download.nvidia.cn

---

這篇文章寫做週期比較長，主要是源碼比較硬核，邊看邊學cuda和c++，到如今也就懂了80%左右吧。不過fastertransformer是真的香，並且直接用tensorflow也很方便，各位須要inference的朋友們必定要用呀