器集群產(chǎn)品圖..png)
器產(chǎn)品圖.png)
器產(chǎn)品圖.png)
器產(chǎn)品圖.png)
器產(chǎn)品圖1-8080.png)
品主圖.png)
算豐技術(shù)揭秘｜重塑翻譯與識別技術(shù)：開(kāi)源語(yǔ)音識別模型Whisper的編譯優(yōu)化與部署
模型介紹
Whisper模型是一個(gè)由OpenAI團隊開(kāi)發(fā)的通用語(yǔ)音識別模型。它的訓練基于大量不同的音頻數據集，是一個(gè)多任務(wù)模型，可以執行語(yǔ)音識別、語(yǔ)言翻譯、語(yǔ)言識別。下面是模型的整體架構：
使用方法如下：
通過(guò)修改TPU-MLIR編譯器代碼，可以對Whisper模型性能進(jìn)行深度優(yōu)化，使得模型在SOPHON BM1684X處理器上運行時(shí)間減少到原來(lái)的一半，本篇文章將帶領(lǐng)大家對Whisper模型進(jìn)行編譯與優(yōu)化，并完成實(shí)際應用的部署。
優(yōu)化方法
本次模型優(yōu)化很具有典型性，不僅適用當前模型，對其他模型也有幫助，下面對相關(guān)的優(yōu)化方法進(jìn)行介紹
Tile算子轉廣播
模型轉換完成后，觀(guān)察Profile時(shí)序圖和對應的MLIR文件，可以看到MatMulOp前面有兩個(gè)Tile操作，通過(guò)復制數據來(lái)使數據shape對齊，并且耗時(shí)占用十分明顯（如紅框所示），實(shí)際上完全可以利用算子的廣播功能實(shí)現，解決方法是使MatMulOp在hdim_is_batch的情況下支持n維度的廣播 （形如：5x1x8x64@1x1500x8x64 ）。
KVCache動(dòng)態(tài)轉靜態(tài)
由于模型結構是基于Transformer模型結構構建的，KVCache的使用是比較常見(jiàn)的，對于使用KVCache的模型，我們的做法是，將KVCache作為模型的輸入和輸出。但在推理過(guò)程中KVCache會(huì )不斷增長(cháng)，呈現一定的動(dòng)態(tài)性，這就使得有一些數據需要PCIe重復搬運；內存管理比較復雜。通過(guò)padding將動(dòng)態(tài)模型轉換為輸入和輸出為固定大小的靜態(tài)模型，并且在內部利用Slice和Concat算子自動(dòng)完成數據拼接，減少外部?jì)却嫣幚韽碗s性，而且通過(guò)優(yōu)化runtime過(guò)程避免多余的PCIe搬運。
Permute算子消除
完成上面兩步優(yōu)化后，通過(guò)觀(guān)察Profile時(shí)序圖和分析final.MLIR文件發(fā)現，這一部分的網(wǎng)絡(luò )被分為了兩個(gè)LayerGroup，之間還存在許多Global Layer，這些都會(huì )導致額外的數據搬運。追溯這一現象的原因，很大可能是由于Slice、Concat和Permute三種操作的存在阻斷了LayerGroup的劃分。而且Permute、Concat和Slice的GDMA操作帶寬利用率低，浪費了較多時(shí)間。下面從這三個(gè)算子入手進(jìn)行優(yōu)化。
在Transformer類(lèi)的模型里，由于數據需要維度翻轉整理，造成模型里有很多的Permute操作，但是通過(guò)算子的實(shí)現，一些運行前后的Permute是可以相互抵消的。如下圖所示，MatMul算子輸入和輸出的Permut是可以消除的，主要步驟是將輸入的Permute算子移動(dòng)到輸出，MatMul利用TPU指令的特性，實(shí)現轉置的矩陣乘法，并在輸出處與原來(lái)的Permute抵消掉。最終結果下圖右側所示。
Slice+Concat算子融合
Slice和Concat本質(zhì)上是將已經(jīng)計算好但放置位置錯誤的結果進(jìn)行截取或搬運。如果我們能提前知道結果應該放到哪里，就可以完全去掉這兩種操作。下面是一個(gè)典型的Slice+Concat的Pattern。通過(guò)分析右圖，可以看到為了將數據排到前面，Slice將Load后的數據進(jìn)行搬運，之后Concat將MatMul的結果搬運到Slice后數據后面。其實(shí)這兩次搬運如果提前知道了放置位置，是可以去掉的。
下圖示意了簡(jiǎn)單的優(yōu)化方法，在Load操作中引入一個(gè)Offset參數，表示數據在Global中的偏移，相當于在Load時(shí)直接做Slice, 減少了重復數據搬運，另外直接將Concat的輸出地址分配給Load，將數據直接寫(xiě)在Concat的輸出地址，省去Concat的GDMA搬運時(shí)間。
最終效果如下：
可以對比下優(yōu)化前后的final.MLIR文件
比較上面兩圖，可以看到Concat和Slice, Permute大部分都去掉了。從下面的Profile也可以看出明顯的提升：
優(yōu)化結果
經(jīng)過(guò)上述優(yōu)化，模型的運行時(shí)間由原來(lái)的23.143ms變?yōu)?3.589ms。為方便定量分析，下面提供了優(yōu)化前后的算子性能統計結果。?
優(yōu)化前性能統計
優(yōu)化后性能統計
部署代碼
部署代碼目前已經(jīng)開(kāi)源(https://github.com/JKay0327/whisper-TPU_pyd)。本W(wǎng)hisper應用整體是由多個(gè)環(huán)節串聯(lián)起來(lái)的，包括前處理、Encoder、Language Detect Decoder、以及主循環(huán)中的Decoder迭代過(guò)程。上面優(yōu)化的主要是針對主循環(huán)中的Decoder模型進(jìn)行的。具體運行過(guò)程如下圖所示。
使用方法如下：
下面是實(shí)際的運行結果展示：
結論
本文是對在Whisper模型應用過(guò)程中的總結，說(shuō)明了在模型優(yōu)化過(guò)程中采用的各種思路和方法，最終將模型的性能翻倍。Whisper模型是一個(gè)很有價(jià)值的應用，可以實(shí)現各種語(yǔ)音任務(wù)，期待大家借助算能產(chǎn)品進(jìn)行更多功能的開(kāi)發(fā)。