地址標準化服務AI深度學習模型推理優化實踐
導讀
深度學習已在面向自然語言處理等領域的實際業務場景中廣泛落地,對它的推理性能優化成為了部署環節中重要的一環。推理性能的提升:一方面,可以充分發揮部署硬件的能力,降低用戶響應時間,同時節省成本;另一方面,可以在保持響應時間不變的前提下,使用結構更為復雜的深度學習模型,進而提升業務精度指標。
本文針對地址標準化服務中的深度學習模型開展了推理性能優化工作。通過高性能算子、量化、編譯優化等優化手段,在精度指標不降低的前提下,AI模型的模型端到端推理速度最高可獲得了4.11倍的提升。
1. 模型推理性能優化方法論
模型推理性能優化是AI服務部署時的重要環節之一。一方面,它可以提升模型推理的效率,充分釋放硬件的性能。另一方面,它可以在保持推理延遲不變的前提下,使得業務采用復雜度更高的模型,進而提升精度指標。然而,在實際場景中推理性能優化會遇到一些困難。
1.1 自然語言處理場景優化難點
典型的自然語言處理(Natural Language Processing, NLP)任務中,循環神經網絡(Recurrent Neural Network, RNN)以及BERT[7](Bidirectional Encoder Representations from Transformers.)是兩類使用率較高的模型結構。為了便于實現彈性擴縮容機制和在線服務部署的高性價比,自然語言處理任務通常部署于例如Intel? Xeon?處理器這樣的x86 CPU平臺。然而,隨著業務場景的復雜化,服務的推理計算性能要求越來越高。以上述RNN和BERT模型為例,其在CPU平臺上部署的性能挑戰如下:
- RNN
循環神經網絡是一類以序列(sequence)數據為輸入,在序列的演進方向進行遞歸(recursion)且所有節點(循環單元)按鏈式連接的遞歸神經網絡。實際使用中常見的RNN有LSTM,GRU以及衍生的一些變種。在計算過程中,如下圖所示,RNN結構中每一次的后級輸出都依賴于相應的輸入和前級輸出。因此,RNN可以完成序列類型的任務,近些年在NLP甚至是計算機視覺領域被廣泛使用。RNN相較于與BERT而言,計算量更少,模型參數共享,但其計算時序依賴性會導致無法對序列進行并行計算。
RNN結構示意圖
- BERT
BERT[7]證明了能夠以較深的網絡結構在大型數據集上完成無監督預訓練(Unsupervised Pre-training),進而供給特定任務進行微調(finetune)的模型。它不僅提升了這些特定任務的精度性能,還簡化了訓練的流程。BERT的模型結構簡單又易于擴展,通過簡單地加深、加寬網絡,即可獲得相較于RNN結構更好的精度。而另一方面,精度提升是以更大的計算開銷為代價的,BERT模型中存在著大量的矩陣乘操作,這對于CPU而言是一種巨大的挑戰。
BERT模型結構示意圖
1.2 模型推理優化策略
基于上述推理性能挑戰的分析,我們認為從軟件棧層面進行模型推理優化,主要有如下策略:
- 模型壓縮:包括量化、稀疏、剪枝等
- 特定場景的高性能算子
- AI編譯器優化
量化
模型量化是指將浮點激活值或權重(通常以32比特浮點數表示)近似為低比特的整數(16比特或8比特),進而在低比特的表示下完成計算的過程。通常而言,模型量化可以壓縮模型參數,進而降低模型存儲開銷;并且通過降低訪存和有效利用低比特計算指令(如Intel? Deep Learning Boost Vector Neural Network Instructions,VNNI),取得推理速度的提升。
給定浮點值,我們可以通過如下公式將其映射為低比特值:
其中和是通過量化算法所得。基于此,以Gemm操作為例,假設存在浮點計算流程:
我們可以在低比特域完成相應的計算流程:
高性能算子
在深度學習框架中,為了保持通用性,同時兼顧各種流程(如訓練),算子的推理開銷存在著冗余。而當模型結構確定時,算子的推理流程僅是原始全量流程個一個子集。因此,當模型結構確定的前提下,我們可以實現高性能推理算子,對原始模型中的通用算子進行替換,進而達到提升推理速度的目的。
在CPU上實現高性能算子的關鍵在于減少內存訪問和使用更高效的指令集。在原始算子的計算流程中,一方面存在著大量的中間變量,而這些變量會對內存進行大量的讀寫操作,進而拖慢推理的速度。針對這種情況,我們可以修改其計算邏輯,以降低中間變量的開銷;另一方面,算子內部的一些計算步驟我們可以直接調用向量化指令集,對其進行加速,如Intel? Xeon?處理器上的高效的AVX512指令集。
AI編譯器優化
隨著深度學習領域的發展,模型的結構、部署的硬件呈現出多樣化演進的趨勢。將模型部署至各硬件平臺時,我們通常會調用各硬件廠商推出的runtime。而在實際業務場景中,這可能會遇到一些挑戰,如:
- 模型結構、算子類型的迭代的速度會高于廠家的runtime,使得一些模型無法快速基于廠商的runtime完成部署。此時需要依賴于廠商進行更新,或者利用plugin等機制實現缺失的算子。
- 業務可能包含多個模型,這些模型可能由多個深度學習框架訓得,此外模型可能需要部署至多個硬件平臺。此時需要將這些格式不同的模型轉化至各個硬件平臺所需的格式,同時要考慮各推理框架實現的不同導致的模型精度性能變化等問題,尤其是像量化這類對于數值差異敏感度較高的方法。
AI編譯器就是為了解決上述問題而提出的,它抽象出了多個層次來解決上述的一些問題。首先,它接受各個前端框架的模型計算圖作為輸入,并通過各類Converter轉化生成統一的中間表示。隨后,諸如算子融合、循環展開等圖優化pass會作用至中間表示,以提升推理性能。最后,AI編譯器會基于優化后的計算圖進行面向特定硬件平臺的codegen,生成可執行的代碼,這過程中會引入諸如stitch、shape constraint等優化策略。AI編譯器有很好魯棒性、適應性、易用性,并且能夠收獲顯著優化收益。
本文中,阿里云機器學習平臺PAI團隊聯合英特爾數據中心軟件團隊、英特爾人工智能和分析團隊、達摩院NLP地址標準化團隊,針對地址標準化服務的推理性能挑戰,合作實現了高性能的推理優化方案。
2. 地址標準化介紹
公安政務、電商物流、能源(水電燃)、運營商、新零售、金融、醫療等行業在業務開展的過程中往往涉及大量地址數據,而這些數據往往沒有形成標準結構規范,存在地址缺失、一地多名等問題。隨著數字化的升級,城市地址不標準的問題愈加凸顯。
地址應用現存問題
地址標準化[2](Address Purification)是阿里巴巴達摩院NLP團隊依托阿里云海量的地址語料庫,以及超強的NLP算法實力所沉淀出的高性能及高準確率的標準地址算法服務。地址標準化產品從規范地址數據、建立統一標準地址庫的角度出發,提供高性能地址算法。
地址標準化優勢
該地址算法服務能自動地標準化處理地址數據,可有效地解決一地多名,地址識別,地址真偽辨別等地址數據不規范、人工治理耗時耗力、地址庫重復建設問題,為企業,政府機關以及開發者提供地址數據清洗,地址標準化能力,使地址數據更好的為業務提供支持。地址標準化產品具有如下的幾個特點:
- 準確率高:擁有海量地址語料庫以及超強的NLP算法技術,并持續優化迭代,地址算法準確率高
- 超強性能:積累了豐富的項目建設經驗,能夠穩定承載海量數據
- 服務全面:提供20多種地址服務,滿足不同業務場景需求
- 部署靈活:支持公共云、混合云、私有化部署。
本次優化的模塊屬于地址標準化中的搜索模塊。地址搜索是指用戶輸入地址文本相關信息,基于地址庫和搜索引擎,對用戶輸入的地址文本進行搜索和聯想,并返回相關興趣點(Point of Interest,POI)信息。地址搜索功能不僅能夠提升用戶數據處理體驗,同時也是多個地址下游服務的基礎,如經緯度查詢、門址標準化、地址歸一等,因此在整套地址服務體系中起到了關鍵作用。
地址服務搜索體系示意圖
具體而言,本次優化的模型是基于多任務地理預訓練語言模型底座產出的多任務向量召回模型和精排模型。
多任務地理預訓練語言模型底座在掩碼語言模型 (Masked Language Model, MLM) 任務的基礎上結合了相關興趣點分類與地址元素識別(省、市、區、POI 等),并通過元學習(Meta Learning)的方式,自適應地調整多個任務的采樣概率,在語言模型中融入通用的地址知識。
多任務地址預訓練模型底座示意圖
多任務向量召回模型基于上述底座訓練所得,包含雙塔相似度、Geohash (地址編碼) 預測、分詞和 Term Weighting (詞權重) 四個任務。
多任務向量召回模型示意圖
作為計算地址相似度匹配的核心模塊,精排模型則是在上述底座的基礎上,引入了海量點擊數據和標注數據訓練訓練所得[3],并通過模型蒸餾技術,提升了模型的效率[4]。最終用應用于召回模型召回的地址庫文檔重新排序。基于上述流程訓練得到的4層單模型能夠在CCKS2021中文NLP地址相關性任務[5]上獲得較12層基線模型更好的效果(詳見性能展示部分)。
精排模型示意圖
3. 模型推理優化解決方案
阿里云機器學習平臺PAI團隊推出的Blade產品支持以上提及的所有優化方案,提供了統一的用戶接口,并擁有多個軟件后端,如高性能算子、Intel Custom Backend、BladeDISC等等。
Blade模型推理優化架構圖
3.1 Blade
Blade是阿里云機器學習PAI團隊(Platform of Artificial Intelligence)推出的通用推理優化工具,可以通過模型系統聯合優化,使模型達到最優推理性能。它有機融合了計算圖優化、Intel? oneDNN等vendor優化庫、BladeDISC編譯優化、Blade高性能算子庫、Costom Backend、Blade混合精度等多種優化手段。同時,簡潔的使用方式降低了模型優化門檻、提升了用戶體驗和生產效率。
PAI-Blade支持多種輸入格式,包括Tensorflow pb、PyTorch torchscript等。對于待優化的模型,PAI-Blade會對其進行分析,再應用多種可能的優化手段,并從各種優化結果中選取出加速效果最明顯的為最終的優化結果。
Blade優化示意圖
為了在保證部署成功率的前提下獲得最大的優化效果,PAI-Blade采取了“圈圖”的方式進行優化,即:
- 將待優化子計算圖中,能夠被推理后端/高性能算子支持的部分轉化至相應的優化子圖;
- 無法被優化的子圖回退(fallback)至相應的原生框架(TF/Torch)執行。
Blade圈圖示意圖
Blade Compression是Blade推出的面向模型壓縮的工具包,旨在協助開發人員進行高效的模型壓縮優化工作。它包含了多種模型壓縮功能,包括量化、剪枝、稀疏化等。壓縮后的模型可以便捷地通過Blade實現進一步優化,以獲得模型系統聯合的極致優化。
量化方面,Blade Compression:
- 提供了簡潔的使用接口,通過調用幾個簡單api,即可完成量化改圖、校準(calibration)、量化訓練(Quantization-aware Training,QAT)、導出量化模型等步驟。
- 提供了多種后端的支持,通過config文件的配置,即可完成面向不同設備、不同后端的量化過程。
- 集成了PAI-Blade團隊在實際生產業務中自研的各種算法,以獲得更高的量化精度。
同時,我們提供了豐富的原子能力api,便于對特定情況進行定制化開發。
Blade Compression示意圖
BladeDISC是阿里云機器學習平臺PAI團隊推出的面向機器學習場景的動態shape深度學習編譯器,是Blade的后端之一。它支持主流的前端框架(TensorFlow、PyTorch)與后端硬件(CPU、GPU),同時也支持推理以及訓練的優化。
BladeDISC架構圖
3.2 基于Intel? Xeon?的高性能算子
神經網絡模型中的子網絡通常具有長期的通用性和普遍性,如 PyTorch 中的 Linear Layer 和Recurrent Layers 等,是模型建構的基礎模塊,負責著特定的功能,通過這些模塊的不同組合得到形形色色的模型,并且這些模塊也是AI編譯器重點優化的目標。據此,為了得到最佳性能的基礎模塊,從而實現性能最佳的模型,Intel針對X86架構對這些基礎模塊進行了多層次優化,包括使能高效的AVX512指令、算子內部計算調度、算子融合、緩存優化,并行優化等等。
在地址標準化服務中,經常會出現Recurrent Neural Network (RNN) 模型,并且RNN模型中最影響性能的模塊是LSTM或GRU等模塊,本章節以LSTM為例,呈現在不定長且多batch的輸入時,如何實現對LSTM的極致性能優化。
通常,為了滿足不同用戶的需求和請求,追求高性能和低成本的云上服務會將不同的用戶請求進行Batch,以實現計算資源的最大化利用。如下圖所示,總共有3條已經被embedding的句子,并且內容和輸入的長度是不相同的。
原始輸入數據
為了使得LSTM計算的更高效,需要對Batched input采用PyTorch的pack_padded_sequence()函數進行padding和sort,得到下圖所示,一個paddding的數據tensor,一個描述數據tensor的batch size的tensor,一個描述數據tensor的原始序號tensor。
原始輸入數據
到目前為止,已經準備好了LSTM的輸入,對于LSTM的計算過程如下圖所示,對輸入的tensor進行分段批量計算,及跳過零值計算。
LSTM針對輸入的計算步驟
更深入的LSTM的計算優化如下圖17所示,公式中的矩陣乘部分進行了公式間計算融合,如下圖所示,原先4次矩陣乘轉換成1次矩陣乘,并且采用AVX512指令進行數值計算,以及多線程并行優化,從而實現高效的LSTM算子。其中,數值計算指的是矩陣乘和后序的elementwise的元素操作,針對矩陣乘部分,本方案采用的是oneDNN庫進行計算,庫中具有高效的AVX512 GEMM實現,針對elementwise的元素操作,本方案對其采用AVX512指令集進行算子融合,提升了數據在緩存中的命中率。
LSTM計算融合[8]
3.3 推理后端 Custom Backend
Intel custom backend[9]作為Blade的軟件后端,強有力地加速著模型量化和稀疏的推理性能,主要包含三個層面的優化。首先,采用Primitive Cache的策略對內存進行優化,其次,進行圖融合優化,最后,在算子層級,實現了包含稀疏與量化算子在內的高效算子庫。
Intel Custom Backend架構圖
低精度量化
稀疏與量化等高速算子, 得益于Intel? DL Boost加速指令集,如VNNI指令集。
VNNI 指令介紹
上圖為VNNI 指令, 8bits可以使用AVX512 BW三個指令來加速, VPMADDUBSW 先對2對由8bits組成的數組做乘法與加法, 得到16bits數據, VPMADDWD將相鄰數據加總起來,得到32bits數據, 最后VPADDD加上一個常數, 此三函數可組成一個AVX512_VNNI,此指令可用來加速推理中的矩陣相乘。
圖融合
除此之外,Custom Backend中也提供了圖融合,例如矩陣相乘后不輸出中間態臨時Tensor,而是直接運行后面指令,即將后項的post op與前級算子進行融合,如此減少數據搬運以減少運行時間,下圖為一個范例,紅框內的算子融合后可消除額外的數據搬移,成為一個新的算子。
圖融合
內存優化
內存分配與釋放會與操作系統進行通信,從而導致運行時的延時增加,為了減少這部分的開銷,Custom Backend中增加了Primitive Cache的設計,Primitive Cache用于緩存已經被創建的Primitive,使得Primitive不能被系統回收,減少了下一次調用時的創建開銷。
同時為耗時較大的算子建立了快取機制,以加速算子運行,如下圖所示:
Primitive Cache
量化功能如之前所說,模型大小減小后,計算與存取的開銷大幅減少,從而性能得到巨大的提升。
4. 整體性能展示
我們選取了地址搜索服務中典型的兩個模型結構來驗證上述優化方案的效果。測試環境如下所示:
- 服務器型號:阿里云 ecs.g7.large,2 vCPU
- 測試CPU型號:Intel? Xeon? Platinum 8369B CPU @ 2.70GHz
- 測試CPU核數:1 vCPU
- PyTorch版本:1.9.0+cpu
- onnx版本:1.11.0
- onnxruntime版本:1.11.1
4.1 ESIM
ESIM[6]是一種專為自然語言推斷而生的加強版LSTM,它的推理開銷主要來自于模型中的LSTM結構。Blade利用Intel數據中心軟件團隊開發的高性能通用LSTM算子對其進行加速,替換PyTorch module中的默認LSTM (Baseline)。本次測試的ESIM中包含兩種LSTM結構,單算子優化前后的性能如表所示:
LSTM結構 | 輸入shape | 優化前RT | 優化后RT | 加速比 |
LSTM - A | 7x200 | 0.199ms | 0.066ms | +3.02x |
202x200 | 0.914ms | 0.307ms | +2.98x | |
LSTM - B | 70x50 | 0.266ms | 0.098ms | +2.71x |
202x50 | 0.804ms | 0.209ms | +3.85x |
LSTM單算子優化前后推理性能
優化前后,ESIM端到端推理速度如表 所示,同時優化前后模型的精度保持不變。
模型結構 | ESIM[6] | ESIM[6]+Blade算子優化 | 加速比 |
RT | 6.3ms | 3.4ms | +1.85x |
ESIM模型優化前后推理性能
4.2 BERT
BERT[7]近年來在自然語言處理 (NLP) 、計算機視覺(CV)等領域被廣泛采納。Blade對該種結構有編譯優化(FP32)、量化(INT8)等多種手段。
速度測試中,測試數據的shape固定為10x53,各種后端及各種優化手段的速度性能如下表所示。可以看到,blade編譯優化后或INT8量化后的模型推理速度均優于libtorch與onnxruntime,其中推理的后端是Intel Custom Backend & BladeDisc。值得注意的是,經過量化加速后的4層BERT的速度是2層BERT的1.5倍,意味著可以在提速的同時,讓業務用上更大的模型,獲得更好的業務精度。
地址BERT推理性能展示
精度方面,我們基于CCKS2021中文NLP地址相關性任務[5]展示相關模型性能,如下表所示。達摩院地址團隊自研的4層BERT的macro F1精度要高于標準的12層BERT-base。Blade編譯優化可以做到精度無損,而經過Blade Compression量化訓練后的真實量化模型精度要略高于原始浮點模型。
模型結構 | macro F1(越高越好) |
12層BERT-base | 77.24 |
地址-4層BERT | 78.72(+1.48) |
地址-4層BERT + Blade編譯優化 | 78.72(+1.48) |
地址-4層BERT + Blade量化 | 78.85(+1.61) |
地址BERT相關精度結果