總結FasterTransformer Encoder(BERT)的cuda相關優(yōu)化技巧

FasterTransformer BERT

FasterTransformer BERT 包含優(yōu)化的 BERT 模型、高效的 FasterTransformer 和 INT8 量化推理。

模型結構

標準的 BERT 和 高效的 FasterTransformer

FasterTransformer 編碼器支持以下配置。

Batch size (B1): 批量大小 <= 4096

Sequence length (S): 序列長度 <= 4096。對于 INT8 模型，當 S > 384 時 S 需要是 32 的倍數(shù)。

Size per head (N): 小于 128 的偶數(shù)。

Head number (H): 在 FP32 下滿足 H * N <= 1024 或在 FP16 下滿足 H * N <= 2048 的任何數(shù)字。

Data type: FP32, FP16, BF16, INT8 and FP8 (Experimental).

如果內存足夠，任意層數(shù)（N1）

在 FasterTransformer v1.0 中，我們提供了高度優(yōu)化的 BERT 等效編碼器模型。

接下來，基于Effective Transformer的思想，我們在 FasterTransformer v2.1 中通過去除無用的 padding 來進一步優(yōu)化BERT推理，并提供 Effective FasterTransformer。

在 FasterTransformer v3.0 中，我們提供了 INT8 量化推理以獲得更好的性能。

在 FasterTransformer v3.1 中，我們優(yōu)化了 INT8 Kernel 以提高 INT8 推理的性能，并將 TensorRT 的多頭注意力插件集成到 FasterTransformer 中。

在 FasterTransformer v4.0 中，我們添加了多頭注意力 Kernel 支持 V100 的 FP16 模式和 T4, A100 的 INT8 模式。

下圖演示了除 INT8 外的這些優(yōu)化的流程圖。在FasterTransformer v5.0中，我們重構了代碼，將 mask building 和 padding 移動到 Bert 的 forward 函數(shù)中，并在 Ampere GPU 上基于稀疏特性來加速GEMM。

在 FasterTransformer v5.1 中，我們支持對 Bert FP16 進行進行多節(jié)點多 GPU 推理。

[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-v85Qr1r0-1674747365445)(null)]

BERT 模型是 google 在2018年提出的。FasterTransformer 的encoder 相當于 BERT 模型，但是做了很多優(yōu)化。

圖 1 最左邊的流程顯示了 FasterTransformer 中的優(yōu)化。經過優(yōu)化后，F(xiàn)asterTransformer 僅使用 8 或 6 個 gemms（藍色塊）和 6 個自定義 CUDA kernel（綠色塊）來實現(xiàn)一個 transformer 塊。

對于 Effective FasterTransformer，主要思想是去除句子的填充以防止計算無用的標記。當一個 Batch 的平均序列長度遠小于最大序列長度時，此方法可以節(jié)省大量時間。

圖 2 顯示了我們使用的想法和偏移量（橙色）。要實現(xiàn) Effective FasterTransformer，我們需要考慮兩個問題。

首先，我們需要去除 BERT 之前的 padding，離開 BERT 之后重建 padding 以保持結果的形狀。

這很簡單，帶來的開銷基本可以忽略。第二個問題是多頭注意力的計算。

一個天真的解決方案是在多頭注意力之前重建填充并在多頭注意力之后移除填充，如圖 1 的第二個流程圖所示。

因為我們可以將這些重建/移除融合到其他 kernel 中，額外的開銷也是可以忽略的。

為了進一步提高多頭注意力的性能，我們集成了 TensorRT 的多頭注意力，將整個注意力計算融合到一個 kernel 中。

源代碼在這里。該 kernel 同時支持 Effective FasterTransformer 和標準 BERT 模型。

圖 1 中的第三個和第四個流程圖顯示了工作流程。

有了這樣的 kernel ，我們就不用擔心多頭注意力的填充問題了。

這個 kernel 需要另一個偏移量，如圖 2 所示。

第一個偏移量 [0, 0, 1, 3, 3, 3]比較好理解，直接和[0, 1, 2, 3, 4, 5]迭代就可以得到原始的位置了。第二個偏移量是從0位置開始，記錄連續(xù)的原始token個數(shù)，比如我們將[0, 2, 3, 6]做差分，得到[2, 1, 3]也對應了原始的數(shù)據(jù)中每行做的padding的tokn數(shù)目。

此外，我們發(fā)現(xiàn) padding 會影響某些任務的準確性，盡管它們應該是無用的。因此，我們建議刪除下游任務最終輸出中的填充。

編碼器的參數(shù)、輸入和輸出：

Constructor of BERT

Input of BERT

Output of BERT

上面聲明了 Bert 模型的輸入參數(shù)，以及輸入和輸出Tensor的shape。

此外，注意到 TensorRT 的多頭注意力Kernel雖然功能很強大但是也有一些限制。首先，這個kernel需要 Turing 或者更新架構的 GPU，并且每個頭的大小必須是64。當條件不滿足時，我們使用FasterTransformer的原始多頭注意力實現(xiàn)。其次，它需要一個額外的序列長度偏移量，如Figure2所示，更多的細節(jié)在這里 。

當輸入有 padding 時，序列長度偏移的形狀為 。假設這里有3個序列，長度分別為 , , ，然后 padding 之后的序列長度為 。那么序列長度偏移時 。即，序列長度偏移記錄了每個句子的序列長度。當我們有 padding 時，我們將 padding 視為一些獨立的句子。

在 FasterTransformer v4.0 中，我們實現(xiàn)了兩條 INT8 推理的流水線，如圖 3 所示。對于 int8_mode == 1 (int8v1)，我們不量化殘差連接，使用 int32 作為 int8 gemms 的輸出，并對權重采用逐通道的量化方式。

對于 int8_mode == 2 (int8v2)，我們量化殘差連接，使用 int8 作為 int8 gemms 的輸出，并對權重采用逐張量的量化。一般來說，int8_mode == 1 的精度更高，而 int8_mode == 2 的性能更好。

對于 INT8 推理，需要量化模型。我們提供了 TensorFlow 量化工具和示例代碼，同時還提供了帶有 TensorRT 量化工具的 PyTorch 示例代碼。

請先參考bert-quantization/bert-tf-quantization和examples/pytorch/bert/bert-quantization-sparsity中的README。

在 FasterTransformer v5.0 中，我們支持稀疏 gemm 以利用 Ampere GPU 的稀疏特性。我們還提供了一個關于 PyTorch 的示例。

在 FasterTransformer v5.1 中，我們支持 BERT 模型的多 GPU 多節(jié)點推理。

優(yōu)化點解讀

優(yōu)化主要是針對 Figure 1 也就是 BERT 的編碼器模塊的各個組件來講(我這里忽略了 Figure1 的和 padding 相關的組建的講解，感興趣的讀者可以自己看看 FasterTransformer)。

我么先把 BERT 的多頭注意力機制的實現(xiàn)貼一下，方便下面的講解：

importtorch.nnasnn

classAttention(nn.Module):
"""
Compute'ScaledDotProductAttention
"""

defforward(self,query,key,value,mask=None,dropout=None):
scores=torch.matmul(query,key.transpose(-2,-1))
/math.sqrt(query.size(-1))

ifmaskisnotNone:
scores=scores.masked_fill(mask==0,-1e9)

p_attn=F.softmax(scores,dim=-1)

ifdropoutisnotNone:
p_attn=dropout(p_attn)

returntorch.matmul(p_attn,value),


classMultiHeadedAttention(nn.Module):
"""
Takeinmodelsizeandnumberofheads.
"""

def__init__(self,h,d_model,dropout=0.1):
super().__init__()
assertd_model%h==0

#Weassumed_valwaysequalsd_k
self.d_k=d_model//h
self.h=h

self.linear_layers=nn.ModuleList([nn.Linear(d_model,d_model)for_inrange(3)])
self.output_linear=nn.Linear(d_model,d_model)
self.attention=Attention()

self.dropout=nn.Dropout(p=dropout)

defforward(self,query,key,value,mask=None):
batch_size=query.size(0)

#1)Doallthelinearprojectionsinbatchfromd_model=>hxd_k
query,key,value=[l(x).view(batch_size,-1,self.h,self.d_k).transpose(1,2)
forl,xinzip(self.linear_layers,(query,key,value))]

#2)Applyattentiononalltheprojectedvectorsinbatch.
x,attn=self.attention(query,key,value,mask=mask,dropout=self.dropout)

#3)"Concat"usingaviewandapplyafinallinear.
x=x.transpose(1,2).contiguous().view(batch_size,-1,self.h*self.d_k)

returnself.output_linear(x)

Compute Q, K, V by three GEMMs or one Batch GEMM

這里的意思就是計算 Q，K，V 的時候有兩種方式，一種是用3個單獨的gemm算子對應FasterTransformer v1.0版本的這段代碼：

另外一種就是通過一個 Batch GEMM算子同時完成對 Q, K, V 的計算，

add_QKV_bias 優(yōu)化

這個是針對上面forward函數(shù)中 (1) 這部分存在的分別對 Q, K, V進行bias_add以及transpose的優(yōu)化，將其融合成一個cuda kernel。這里啟動 add_QKV_bias 的參數(shù)來看。

對于FP32，F(xiàn)asterTransformer是啟動 batch_size * seq_len * 3 個 Block， 每個 Block 里面啟動 head_num * size_per_head 個線程只處理一個token（對應 head_num * size_per_head 次計算）的 bias_add 計算。

我們注意到這里還將輸入的shape進行了改變，也就是將原始的[batch_size, seq_length, head_num * size_per_head] -> [batch_size, seq_length, head_num, size_per_head]（對應 .view(batch_size, -1, self.h, self.d_k)）->[batch_size, head_num, seq_length, size_per_head]（對應.transpose(1, 2)），這個過程對應了 https://github.com/NVIDIA/FasterTransformer/blob/release/v1.0_tag/fastertransformer/cuda/open_attention.cu#L149 這里的索引代碼。

而對于FP16模式，F(xiàn)asterTransformer是啟動 batch_size * seq_len 個 Block，,每個 Block 里面啟動 head_num * size_per_head 個線程同時處理QKV的同一個token（對應head_num * size_per_head次計算），在實際計算時會把half pack成half2進行計算：https://github.com/NVIDIA/FasterTransformer/blob/release/v1.0_tag/fastertransformer/cuda/open_attention.cu#L172 ，并使用了half2相關的數(shù)學函數(shù)。這樣不僅僅可以達到2倍于half的訪存帶寬和計算吞吐，還可以極大地減少指令的發(fā)射數(shù)量。

高效的softmax kernel

這里我沒有怎么看，因為oneflow已經有一個比FasterTransformer更好的softmax kernel實現(xiàn)了。

transpose kernel

這個 kernel 是對應上面 BERT 的 Encoder 部分的：

x=x.transpose(1,2).contiguous().view(batch_size,-1,self.h*self.d_k)

這里的 x 的 shape 仍然和之前的 q 的 shape 一致， 為[batch_size, head_num, seq_length, size_per_head]。

因為Attetion 層不會改變輸入的形狀，因為 Attention 的計算過程是：q * k 轉置(.transpose(2, 3))，除以 d_k ** 0.5，輸出維度是 [b, head_num , seq_length, seq_length] 即單詞和單詞直接的相似性 ，然后對最后一個維度進行 softmax 操作得到 [b, head_num, seq_length, seq_length] , 最后和 v（shape 也是 [batch_size, head_num, seq_length, size_per_head]） 做一個矩陣乘法，結果的 shape 和輸入的 shape 形狀都是：[batch_size, head_num, seq_length, size_per_head] 。因此這里的 x.transpose(1, 2) 就是把 shape 為 [batch_size, head_num, seq_length, size_per_head] 的 x 重新排列為 [batch_size, head_num, size_per_head, seq_length]。然后 x.contiguous().view(batch_size, -1, self.h * self.d_k) 進一步將 shape 重新排列為 [batch_size, seq_length, head_num * size_per_head] 。

對于 FP32 模式，啟動 batch_size * head_num * seq_length 個 Block , 然后每個 Block 啟動 size_per_head 個線程處理一個序列（一個序列對應 size_per_head 個元素）。如下：

constintseq_per_block=1;
grid.x=batch_size*head_num*seq_len/seq_per_block;
block.x=seq_per_block*size_per_head;
transpose<<>>(transpose_dst_,dst,
batch_size,seq_len,head_num,size_per_head);

而 transpose 的kernel實現(xiàn)也比較簡單，根據(jù)blockIdx.x計算下batch_id和seq_id以及head_id（輸入 x 的 shape 為 [batch_size, head_num, seq_length, size_per_head]）：

template
__global__
voidtranspose(T*src,T*dst,constintbatch_size,constintseq_len,constinthead_num,constintsize_per_head)
{
intbatch_id=blockIdx.x/(head_num*seq_len);
intseq_id=blockIdx.x%seq_len;
inthead_id=(blockIdx.x%(head_num*seq_len))/seq_len;
dst[batch_id*(head_num*seq_len*size_per_head)+seq_id*head_num*size_per_head
+head_id*size_per_head+threadIdx.x]=src[blockIdx.x*size_per_head+threadIdx.x];
}

對于 half 來說，采用和 add_QKV_bias 一樣的優(yōu)化方式，每個 block 處理 4 個sequence。具體來說，就是現(xiàn)在啟動 batch_size * head_num * seq_len / 4 個 Block， 每個 Block 使用 2 * size_per_head 個線程處理 4 個序列。為什么 2 * size_per_head 個線程可以處理 4 個序列（一個序列對應 size_per_head 個元素），原因是因為使用了 half2 來做數(shù)據(jù)讀取。half 類型的 kernel 實現(xiàn)如下：

__inline____device__
inttarget_index(intid1,intid2,intid3,intid4,intdim_1,intdim_2,intdim_3,intdim_4)
{
returnid1*(dim_2*dim_3*dim_4)+id3*(dim_2*dim_4)+id2*dim_4+id4;
}

template<>
__global__
voidtranspose(__half*src,__half*dst,
constintbatch_size,constintseq_len,constinthead_num,constintsize_per_head)
{
inttid=blockIdx.x*blockDim.x+threadIdx.x;

intbatch_id=tid/(head_num*seq_len*size_per_head);
inthead_id=(tid%(head_num*seq_len*size_per_head))/(seq_len*size_per_head);
intseq_id=(tid%(seq_len*size_per_head))/size_per_head;
intid=tid%size_per_head;

inttarget_id=target_index(batch_id,head_id,seq_id,id,batch_size,head_num,seq_len,size_per_head);
half2*src_ptr=(half2*)src;
half2*dst_ptr=(half2*)dst;

dst_ptr[target_id]=src_ptr[tid];
}

trt_add_QKV_bias 和 TensorRT fused multi-head attention kernel

實際上從 Figure1 也可以看出我們上面講到的 batch GEMM，softmax, GEMM，transpose 等操作都可以被合成一個超大的 cuda kernel，進一步進行優(yōu)化，也就是這里的 TensorRT fused multi-head attention kernel。這個是將 TensorRT 的這個插件作為第三方倉庫引入到 FasterTransformer 進行加速的，具體的代碼我沒有研究過，這里就不展開了。

現(xiàn)在 MultiHeadAttention 部分涉及到的優(yōu)化其實就講完了，我們接著看一下FasterTransformer 對 BERT Encoder 的其它部分的優(yōu)化。我們這里貼一下 Transformer 的結構圖：

在 MultiHeadAttention 的后面接了一個 Add & Norm，這里的 Add 其實就是殘差，Norm 就是 LayerNorm。所以 Encoder 部分的兩個 Add & Norm 可以總結為：

add_bias_input_layernorm

這里的 LayerNorm(X + MultiHeadAttention(X)) 就對應了 FasterTransformer 里面的 add_bias_input_layernorm 這個優(yōu)化

實際上 layernorm 在 oneflow 也有比 Faster Transformer 更好的 kernel，

對于 softmax 和 layernorm 我還沒看 FasterTransformer 的源碼，后續(xù)研究了之后再分享。

總的來說就是 add_bias_input_layernorm 這個優(yōu)化把殘差連接和LayerNorm fuse到一起了，性能更好并且降低了kernel launch的開銷。

add_bias_act

在上圖的 Feed Forward 的實現(xiàn)中，還有一個 bias_add 和 gelu 激活函數(shù)挨著的 pattern ，所以 FasterTransformer 實現(xiàn)了這個 add_bias_act kernel 將兩個操作融合起來，常規(guī)操作。

番外

除了上述優(yōu)化技巧之外，我們可以為 BERT 模型在我們的GPU上試跑 GEMM，并且保存 GEMM 性能最高的超參數(shù)配置，這個對于 cublas 和 cutlass 實現(xiàn)的卷積應該都是成立的。

另外我觀察到在 Faster Transformer 的cuda kernel實現(xiàn)中，大量應用了__ldg這個指令，查了一下資料說這個是只讀緩存指令，在讀地址比較分散的情況下，這個只讀緩存比L1的表現(xiàn)要好，對一些帶寬受限的kernel有性能提升。后續(xù)有空繼續(xù)研究下...

總結

我這邊從文檔上初步看的東西也就這么多，后續(xù)可能會繼續(xù)研究學習下Faster Transformer的softmax/layernorm實現(xiàn)，或者解讀一下其它Transformer架構的優(yōu)化技巧。

審核編輯：劉清