重塑翻譯與識別技術(shù)：開源語音識別模型Whisper的編譯優(yōu)化與部署

模型介紹

Whisper模型是一個由OpenAI團(tuán)隊開發(fā)的通用語音識別模型。它的訓(xùn)練基于大量不同的音頻數(shù)據(jù)集，是一個多任務(wù)模型，可以執(zhí)行語音識別、語言翻譯、語言識別。下面是模型的整體架構(gòu)：

使用方法如下：

通過修改TPU-MLIR編譯器代碼，可以對Whisper模型性能進(jìn)行深度優(yōu)化，使得模型在SOPHON BM1684X處理器上運行時間減少到原來的一半，本篇文章將帶領(lǐng)大家對Whisper模型進(jìn)行編譯與優(yōu)化，并完成實際應(yīng)用的部署。

優(yōu)化方法

本次模型優(yōu)化很具有典型性，不僅適用當(dāng)前模型，對其他模型也有幫助，下面對相關(guān)的優(yōu)化方法進(jìn)行介紹

Tile算子轉(zhuǎn)廣播

模型轉(zhuǎn)換完成后，觀察Profile時序圖和對應(yīng)的MLIR文件，可以看到MatMulOp前面有兩個Tile操作，通過復(fù)制數(shù)據(jù)來使數(shù)據(jù)shape對齊，并且耗時占用十分明顯（如紅框所示），實際上完全可以利用算子的廣播功能實現(xiàn)，解決方法是使MatMulOp在hdim_is_batch的情況下支持n維度的廣播 （形如：5x1x8x64@1x1500x8x64 ）。

KVCache動態(tài)轉(zhuǎn)靜態(tài)

由于模型結(jié)構(gòu)是基于Transformer模型結(jié)構(gòu)構(gòu)建的，KVCache的使用是比較常見的，對于使用KVCache的模型，我們的做法是，將KVCache作為模型的輸入和輸出。但在推理過程中KVCache會不斷增長，呈現(xiàn)一定的動態(tài)性，這就使得有一些數(shù)據(jù)需要PCIe重復(fù)搬運；內(nèi)存管理比較復(fù)雜。通過padding將動態(tài)模型轉(zhuǎn)換為輸入和輸出為固定大小的靜態(tài)模型，并且在內(nèi)部利用Slice和Concat算子自動完成數(shù)據(jù)拼接，減少外部內(nèi)存處理復(fù)雜性，而且通過優(yōu)化runtime過程避免多余的PCIe搬運。

Permute算子消除

完成上面兩步優(yōu)化后，通過觀察Profile時序圖和分析final.MLIR文件發(fā)現(xiàn)，這一部分的網(wǎng)絡(luò)被分為了兩個LayerGroup，之間還存在許多Global Layer，這些都會導(dǎo)致額外的數(shù)據(jù)搬運。追溯這一現(xiàn)象的原因，很大可能是由于Slice、Concat和Permute三種操作的存在阻斷了LayerGroup的劃分。而且Permute、Concat和Slice的GDMA操作帶寬利用率低，浪費了較多時間。下面從這三個算子入手進(jìn)行優(yōu)化。

在Transformer類的模型里，由于數(shù)據(jù)需要維度翻轉(zhuǎn)整理，造成模型里有很多的Permute操作，但是通過算子的實現(xiàn)，一些運行前后的Permute是可以相互抵消的。如下圖所示，MatMul算子輸入和輸出的Permut是可以消除的，主要步驟是將輸入的Permute算子移動到輸出，MatMul利用TPU指令的特性，實現(xiàn)轉(zhuǎn)置的矩陣乘法，并在輸出處與原來的Permute抵消掉。最終結(jié)果下圖右側(cè)所示。

Slice+Concat算子融合

Slice和Concat本質(zhì)上是將已經(jīng)計算好但放置位置錯誤的結(jié)果進(jìn)行截取或搬運。如果我們能提前知道結(jié)果應(yīng)該放到哪里，就可以完全去掉這兩種操作。下面是一個典型的Slice+Concat的Pattern。通過分析右圖，可以看到為了將數(shù)據(jù)排到前面，Slice將Load后的數(shù)據(jù)進(jìn)行搬運，之后Concat將MatMul的結(jié)果搬運到Slice后數(shù)據(jù)后面。其實這兩次搬運如果提前知道了放置位置，是可以去掉的。

下圖示意了簡單的優(yōu)化方法，在Load操作中引入一個Offset參數(shù)，表示數(shù)據(jù)在Global中的偏移，相當(dāng)于在Load時直接做Slice, 減少了重復(fù)數(shù)據(jù)搬運，另外直接將Concat的輸出地址分配給Load，將數(shù)據(jù)直接寫在Concat的輸出地址，省去Concat的GDMA搬運時間。

最終效果如下：

可以對比下優(yōu)化前后的final.MLIR文件

比較上面兩圖，可以看到Concat和Slice, Permute大部分都去掉了。從下面的Profile也可以看出明顯的提升：

優(yōu)化結(jié)果

經(jīng)過上述優(yōu)化，模型的運行時間由原來的23.143ms變?yōu)?3.589ms。為方便定量分析，下面提供了優(yōu)化前后的算子性能統(tǒng)計結(jié)果。

優(yōu)化前性能統(tǒng)計

優(yōu)化后性能統(tǒng)計

部署代碼

部署代碼目前已經(jīng)開源(https://github.com/JKay0327/whisper-TPU_pyd)。本W(wǎng)hisper應(yīng)用整體是由多個環(huán)節(jié)串聯(lián)起來的，包括前處理、Encoder、Language Detect Decoder、以及主循環(huán)中的Decoder迭代過程。上面優(yōu)化的主要是針對主循環(huán)中的Decoder模型進(jìn)行的。具體運行過程如下圖所示。

使用方法如下：

下面是實際的運行結(jié)果展示：

結(jié)論

本文是對在Whisper模型應(yīng)用過程中的總結(jié)，說明了在模型優(yōu)化過程中采用的各種思路和方法，最終將模型的性能翻倍。Whisper模型是一個很有價值的應(yīng)用，可以實現(xiàn)各種語音任務(wù)，期待大家借助算能產(chǎn)品進(jìn)行更多功能的開發(fā)。