FasterTransformer BERT

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

模型結(jié)構(gòu)

標(biāo)準(zhǔn)的 BERT 和 高效的 FasterTransformer

FasterTransformer 編碼器支持以下配置。

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

Sequence length (S): 序列長度 <= 4096。對于 INT8 模型，當(dāng) S > 384 時(shí) 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).

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

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

接下來，基于Effective Transformer的思想，我們在 FasterTransformer v2.1 中通過去除無用的 padding 來進(jìn)一步優(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中，我們重構(gòu)了代碼，將 mask building 和 padding 移動(dòng)到 Bert 的 forward 函數(shù)中，并在 Ampere GPU 上基于稀疏特性來加速GEMM。

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

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

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

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

對于 Effective FasterTransformer，主要思想是去除句子的填充以防止計(jì)算無用的標(biāo)記。當(dāng)一個(gè) Batch 的平均序列長度遠(yuǎn)小于最大序列長度時(shí)，此方法可以節(jié)省大量時(shí)間。

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

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

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

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

因?yàn)槲覀兛梢詫⑦@些重建/移除融合到其他 kernel 中，額外的開銷也是可以忽略的。

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

源代碼在這里。該 kernel 同時(shí)支持 Effective FasterTransformer 和標(biāo)準(zhǔn) BERT 模型。

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

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

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

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

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

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

Constructor of BERT

Input of BERT

Output of BERT

上面聲明了 Bert 模型的輸入?yún)?shù)，以及輸入和輸出Tensor的shape。

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

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

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

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

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

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

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

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

優(yōu)化點(diǎn)解讀

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

我么先把 BERT 的多頭注意力機(jī)制的實(shí)現(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

這里的意思就是計(jì)算 Q，K，V 的時(shí)候有兩種方式，一種是用3個(gè)單獨(dú)的gemm算子對應(yīng)FasterTransformer v1.0版本的這段代碼：

另外一種就是通過一個(gè) Batch GEMM算子同時(shí)完成對 Q, K, V 的計(jì)算，

add_QKV_bias 優(yōu)化

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

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

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

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

高效的softmax kernel

這里我沒有怎么看，因?yàn)閛neflow已經(jīng)有一個(gè)比FasterTransformer更好的softmax kernel實(shí)現(xiàn)了。

transpose kernel

這個(gè) kernel 是對應(yīng)上面 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]。

因?yàn)锳ttetion 層不會(huì)改變輸入的形狀，因?yàn)?Attention 的計(jì)算過程是：q * k 轉(zhuǎn)置(.transpose(2, 3))，除以 d_k ** 0.5，輸出維度是 [b, head_num , seq_length, seq_length] 即單詞和單詞直接的相似性 ，然后對最后一個(gè)維度進(jìn)行 softmax 操作得到 [b, head_num, seq_length, seq_length] , 最后和 v（shape 也是 [batch_size, head_num, seq_length, size_per_head]） 做一個(gè)矩陣乘法，結(jié)果的 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) 進(jìn)一步將 shape 重新排列為 [batch_size, seq_length, head_num * size_per_head] 。

對于 FP32 模式，啟動(dòng) batch_size * head_num * seq_length 個(gè) Block , 然后每個(gè) Block 啟動(dòng) size_per_head 個(gè)線程處理一個(gè)序列（一個(gè)序列對應(yīng) size_per_head 個(gè)元素）。如下：

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實(shí)現(xiàn)也比較簡單，根據(jù)blockIdx.x計(jì)算下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)化方式，每個(gè) block 處理 4 個(gè)sequence。具體來說，就是現(xiàn)在啟動(dòng) batch_size * head_num * seq_len / 4 個(gè) Block， 每個(gè) Block 使用 2 * size_per_head 個(gè)線程處理 4 個(gè)序列。為什么 2 * size_per_head 個(gè)線程可以處理 4 個(gè)序列（一個(gè)序列對應(yīng) size_per_head 個(gè)元素），原因是因?yàn)槭褂昧?half2 來做數(shù)據(jù)讀取。half 類型的 kernel 實(shí)現(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

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

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

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

add_bias_input_layernorm

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

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

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

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

add_bias_act

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

番外

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

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

總結(jié)

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

審核編輯：劉清