本文經(jīng)AI新媒體量子位（公眾號(hào)ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請聯(lián)系出處。

 4月9日，英偉達(dá)x量子位分享了一期nlp線上課程，來自NVIDIA的GPU計(jì)算專家、FasterTransformer 2.0開發(fā)者之一的薛博陽老師，與數(shù)百位開發(fā)者共同探討了：

• FasterTransformer 2.0 新增功能介紹
• 如何針對Decoder和Decoding進(jìn)行優(yōu)化
• 如何使用Decoder和Decoding
• Decoder和Decoding能夠帶來什么樣的加速效果

應(yīng)讀者要求，我們將分享內(nèi)容整理出來，與大家一起學(xué)習(xí)。文末附有本次直播回放、PPT鏈接，大家也可直接觀看。

以下為本次分享的內(nèi)容整理：

大家好，今天為大家介紹的是FasterTransformer 2.0的原理與應(yīng)用。

什么是FasterTransformer？

首先，參加本期直播的朋友對Transformer架構(gòu)應(yīng)該都有了一定了解。這個(gè)架構(gòu)是在“Attention is All You Need”這篇論文中提出的。在BERT Encoder中使用了大量Transformer，效果很好。因此，Transformer已成為 NLP 領(lǐng)域中非常熱門的深度學(xué)習(xí)網(wǎng)絡(luò)架構(gòu)。

但是，Transformer計(jì)算量通常是非常大的。因此，Transformer的時(shí)延往往難以滿足實(shí)際應(yīng)用的需求。

Transformer架構(gòu)可以應(yīng)用于Encoder或Decoder上。在Encoder中，Transformer包含1個(gè)multi-head attention和1個(gè)feed forward網(wǎng)絡(luò)，在Decoder中，包含2個(gè)multi-head attention和1個(gè)feed forward網(wǎng)絡(luò)。

其中，純Encoder架構(gòu)在目前的很多應(yīng)用中都有很好的表現(xiàn)，比如Q&A系統(tǒng)、廣告推薦系統(tǒng)等，因此，針對Encoder的優(yōu)化是非常有必要的。

而在一些場景中，如翻譯場景，我們需要Encoder和Decoder架構(gòu)。在這種架構(gòu)下，Decoder消耗的時(shí)間占比是非常高的，可能達(dá)到90%以上，是推理的主要瓶頸。因此，針對Decoder的優(yōu)化也是一項(xiàng)重要的工作，能帶來明顯的加速效果。

實(shí)際應(yīng)用中，F(xiàn)asterTransformer 1.0版本針對BERT中的Encoder為目標(biāo)做了很多優(yōu)化和加速。在2.0版本中，則主要新增了針對Decoder的優(yōu)化，其優(yōu)越的性能將助力于翻譯、對話機(jī)器人、文字補(bǔ)全修正等多種生成式的場景。

上表比較了Encoder和Decoder計(jì)算量的不同。當(dāng)我們需要編解碼一個(gè)句子的時(shí)候，Encoder可以同時(shí)編碼很多個(gè)字，甚至可以直接編碼一個(gè)句子。

但是Decoder是一個(gè)解碼的過程，每次只能解碼一個(gè)字，因此，解碼一個(gè)句子時(shí)我們需要多次Decoder的forward，對GPU更不友善。

上圖列出了FasterTransformer中針對BERT優(yōu)化的模塊。在編碼方面，以BERT為基準(zhǔn)，提供了一個(gè)單層的、等價(jià)于BERT Transformer 的模塊，供使用者做調(diào)用。當(dāng)我們需要多層的Transformer時(shí)，只需調(diào)用多次Encoder即可。

解碼方面更為復(fù)雜，為了兼顧靈活性與效率，我們提供兩個(gè)不同大小和效果的模塊：

Decoder（黃色區(qū)塊） 由單層的 Transformer layer 組成，它包含兩個(gè)attention和一個(gè)feed forward 網(wǎng)絡(luò)；而Decoding（藍(lán)色區(qū)塊）除了包含多層的 Transformer layer 之外，還包括了其他函數(shù)，例如 embedding_lookup、beam search、position Encoding 等等。

我們用一個(gè)簡單的虛擬碼展示Decoder和Decoding的區(qū)別。

在Decoding中通常有兩個(gè)終止條件，一是是否達(dá)到預(yù)先設(shè)定的最大的sequence length，第二個(gè)條件是所有的句子是否都已經(jīng)翻譯完畢，未終止時(shí)會(huì)不斷循環(huán)。

以句子長度為128的句子翻譯場景為例，若其 Decoder 是由6層的 Transformer layer 組成的，總共需要調(diào)用 128x6=768 次的Decoder；如果是使用 Decoding 的話，則只需要調(diào)用一次Decoding，因此Decoding的推理效率更高。

小結(jié)

首先，F(xiàn)asterTransformer提供了高度優(yōu)化過的Transformer layer：在Encoder方面是基于BERT實(shí)現(xiàn)的；在Decoder方面基于OpenNMT-TensorFlow開源的庫做為標(biāo)準(zhǔn)；Decoding包含了翻譯的整個(gè)流程，也是基于OpenNMT-TensorFlow。

其次，F(xiàn)asterTransformer 2.0的底層由CUDA和cuBLAS實(shí)現(xiàn)，支持FP16 和 FP32 兩種計(jì)算模式，目前提供C++ API和TF OP。

現(xiàn)在，F(xiàn)asterTransformer 2.0已經(jīng)開源，大家可以在DeepLearningExamples/FasterTransformer/v2 at master · NVIDIA/DeepLearningExamples · GitHub獲取全部源代碼。

如何進(jìn)行優(yōu)化？

先以Encoder為例。

參數(shù)：no XLA，batch size 1，12 heads，size per head 64，F(xiàn)P 32

圖中藍(lán)色方塊表示GPU在實(shí)際運(yùn)行，空白的表示GPU在閑置，因此GPU在很多時(shí)間是閑置狀態(tài)。造成GPU閑置的原因是kernels太小，GPU要不斷閑置以等待CPU啟動(dòng)kernel的時(shí)間，這也稱為kernel launch bound問題。

如何解決這個(gè)問題？

我們嘗試開啟TF的XLA，其他參數(shù)不變。圖中我們看到，從原本計(jì)算1層Transformer layer需要50個(gè)kernel縮減到24個(gè)左右。大部分kernel變得比較寬，雖有加速，但是閑置的時(shí)間還是比較多。

因此，我們提出FasterTransformer針對Encoder進(jìn)行優(yōu)化。

首先，我們把矩陣計(jì)算部分挑選出來，用NVIDIA高度優(yōu)化的庫cuBLAS 來計(jì)算，此外的部分，我們把能融合的kernel都盡可能融合起來。

最終的結(jié)果如上圖右邊，經(jīng)過整體的優(yōu)化后，我們只需要8個(gè)矩陣計(jì)算加6個(gè)kernel就可以完成單層Transformer layer計(jì)算，也就是說所需kernel從24個(gè)減少到14個(gè)。

我們可以看到，優(yōu)化后每一個(gè) kernel 都相對比較大，時(shí)間占比小的kernel也減少了。但還是有很多空白的片段。

我們直接調(diào)用C++ API，如圖，GPU閑置的時(shí)間幾乎沒有了。因此，小batch size情況下，我們推薦使用C++ API以獲得更快的速度。當(dāng)batch size比較大時(shí)，GPU閑置時(shí)間會(huì)比較少。

接下來我們看下Decoder。

參數(shù)：no XLA，batch size 1，8 heads，size per head 64，F(xiàn)P32

經(jīng)過統(tǒng)計(jì)，TF需要使用70個(gè)左右kernel來計(jì)算1層Transformer layer。直觀來看，非常小、時(shí)間占比非常短的kernel更多。因此，batch size比較小的情況下，優(yōu)化效果會(huì)更明顯。

Decoder的優(yōu)化同上述Encoder，特別之處是，Decoder里面的矩陣計(jì)算量非常少，因此我們把整個(gè)multi-head attention以一個(gè)kernel來完成。經(jīng)過優(yōu)化之后，原本需要70個(gè)kernel才能完成的計(jì)算，只需要使用16個(gè)kernel就能夠完成。

在更大的Decoding模塊中，另一個(gè)時(shí)間占比較多的kernel是beam search，這里我們針對top k做出優(yōu)化。在GPU中可以同時(shí)執(zhí)行多個(gè)block和多個(gè)warp，并行運(yùn)算，大大節(jié)省時(shí)間。

如何使用FasterTransformer？

大家可以在DeepLearningExamples/FasterTransformer/v2 at master · NVIDIA/DeepLearningExamples · GitHub根目錄下找到對應(yīng)資料：

針對 Decoder 和 Decoding，F(xiàn)asterTransformer 分別提供了 C++ 和 TensorFlow OP 這兩種接口。

C++接口

首先創(chuàng)建一個(gè)Eecoder，超參數(shù)如圖:

其次，設(shè)定訓(xùn)練好的模型權(quán)重；

設(shè)置好后，直接調(diào)用forward即可。

TF OP接口

首先，我們需要先載入OP。這里以Decoder為例，會(huì)自動(dòng)創(chuàng)建TF需要使用的庫，調(diào)用接口時(shí)先導(dǎo)入.so文件（圖中已標(biāo)紅）：

然后調(diào)用Decoder，放入input、權(quán)重、超參數(shù)，然后針對out put 做Session run。

這里需要注意的是，參數(shù)里有一個(gè)虛擬的輸入 (pseudo input)。這個(gè)輸入是為了避免 TensorFlow 的 decoder 和 FasterTransformer Decoder 發(fā)生并行，因?yàn)槲覀儼l(fā)現(xiàn)并行執(zhí)行時(shí)，Decoder中的memory可能會(huì)被污染。實(shí)際應(yīng)用的時(shí)候可以將這個(gè)輸入拿掉。

優(yōu)化效果

最后我們來看下優(yōu)化的效果如何。首先測試環(huán)境設(shè)置：

使用的GPU是NVIDIA的Tesla T4和V100。

Encoder模塊在Tesla V100的結(jié)果

超參數(shù)設(shè)置：12 layers，32 sequence length，12 heads，64 size per head（BERT base），under FP 16

結(jié)果如上圖，batch size從100逐步增加到500的過程中，F(xiàn)asterTransformer對比TF開啟XLA，大概可以提供1.4倍的加速。

Decoder和Decoding模塊在Tesla T4的結(jié)果

超參數(shù)設(shè)置：Batch size 1，beam width 4，8 heads，64 size per head，6 layers，vocabulary size 30000，F(xiàn)P 32

結(jié)果如上圖，不同的sequence length下，相比于TF，F(xiàn)asterTransformer Decoder可以帶來3.4倍左右的加速效果，Decoding可以帶來7-8倍的加速，效率更高。

超參數(shù)設(shè)置：Batch size 256，sequence length 32，beam width 4，8 heads，64 size per head，6 layers，vocabulary size 30000

結(jié)果如上圖，把batch size固定在較高值時(shí)，不同的FP下，F(xiàn)asterTransformer Decoder和Decoding也帶來一定的加速效果。