模型介紹

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

使用方法如下：

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

優(yōu)化方法

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

Tile算子轉(zhuǎn)廣播

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

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

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

Permute算子消除

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

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

Slice+Concat算子融合

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

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

最終效果如下：

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

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

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

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

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

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

部署代碼

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

使用方法如下：

下面是實(shí)際的運(yùn)行結(jié)果展示：

結(jié)論

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