AICompiler動態(tài)shape編譯框架
移動互聯(lián)網(wǎng)的興起，不僅產(chǎn)生了海量數(shù)據(jù)，也對人機(jī)交互有了新的定義。企業(yè)如何動態(tài)處理不同規(guī)格圖片數(shù)據(jù)，如何更靈活處理不同長度的對話語料等等，提升企業(yè)運(yùn)營效率，爭取更多的商業(yè)機(jī)會和流量，成為眾多企業(yè)探索的熱門技術(shù)應(yīng)用。
近期，阿里云機(jī)器學(xué)習(xí)PAI團(tuán)隊(duì)全新上線一套Dynamic Shape Compiler框架，不僅作為AICompiler技術(shù)棧中原有的Static Shape Compiler框架的重要補(bǔ)充，更是增加了Compiler在企業(yè)級數(shù)據(jù)處理應(yīng)用的無限可能，在提升數(shù)據(jù)處理效率的同時(shí)，大幅提升AI工程化效率。先來看看案例和效果數(shù)據(jù)。
性能結(jié)果。
TensorFlow語音識別業(yè)務(wù)示例
以某業(yè)務(wù)方的語音識別模型為例。過往為這個業(yè)務(wù)提供的優(yōu)化主要基于Static Shape Compiler，shape變化范圍較大，只能采用脫機(jī)預(yù)編譯的方式來完成部署，部署過程較為繁瑣。
下圖展示了基于Dynamic Shape Compiler在不同batchsize下的實(shí)際性能結(jié)果，其中縱軸為latency的提升倍數(shù)。整體編譯次數(shù)從之前的幾千次降低到1次。從數(shù)字上來看，在只有一次編譯的較小編譯開銷下，性能十分接近Static Shape Compiler的性能優(yōu)化結(jié)果。

TensorFlow廣告推薦業(yè)務(wù)示例
下面這個例子是某廣告推薦業(yè)務(wù)方的推理模型，在在線系統(tǒng)中，預(yù)估時(shí)的input shape變化非常頻繁，比如：用戶畫像卷標(biāo)個數(shù)可能不同；用戶的歷史行為序列的長度會不同；召回廣告的集合大小會不同；廣告所屬類目的數(shù)量也會不同。這些變量會最終導(dǎo)致請求預(yù)估服務(wù)的input shape也是時(shí)刻變化的。
過往業(yè)務(wù)方的同學(xué)需要通過batching/手工干預(yù)圈圖等方式才能將編譯次數(shù)控制到可接受范圍內(nèi)，造成每次模型迭代后部署過程較為繁瑣。從性能結(jié)果上看，對比基于XLA優(yōu)化的性能，Dynamic Shape Compiler基本接近甚至超過Static Shape Compiler的性能優(yōu)化結(jié)果。

TensorFlow語音合成業(yè)務(wù)示例
以某業(yè)務(wù)方的TTS模型為例，具體看一下實(shí)際業(yè)務(wù)中對優(yōu)化工具對動態(tài)shape支持的需求情況。在這個業(yè)務(wù)模型里，用戶的輸入sequence length輸出sequence length都可能發(fā)生變化。此外，由于TTS類算法本身需要在推理過程中引入隨機(jī)數(shù)的特點(diǎn)，即使輸入同一條樣本，內(nèi)部子圖的shape也會發(fā)生變化。基于static shape compiler來優(yōu)化這個模型的話，輸入數(shù)據(jù)數(shù)量，以及累積編譯次數(shù)的變化曲線。在這個業(yè)務(wù)模型中每個shape的編譯開銷約為20s左右，可以看到在經(jīng)過幾百輪迭代之后，編譯cache都還沒有收斂趨勢，根據(jù)理論shape變化范圍測算總編譯時(shí)間至少在10個小時(shí)以上，這類業(yè)務(wù)如果使用XLA等靜態(tài)shape編譯程序的話，無法透明的實(shí)現(xiàn)可商用的部署。AICompiler里面的dynamic shape compiler組件很好的解決了這一問題，在一次編譯的前提下幫助用戶獲得平均2X的性能收益，目前業(yè)界尚無其它AI編譯程序能夠?qū)崿F(xiàn)類似的效果。

PyTorch公式識別業(yè)務(wù)示例
下圖是一個PyTorch業(yè)務(wù)的例子，對比的Baseline是基于libTorch執(zhí)行導(dǎo)出后的TorchScipt腳本，可以看到AICompiler對PyTorch業(yè)務(wù)提供了同樣的編譯優(yōu)化能力。

本文主要介紹這套動態(tài)shape編譯框架，對更多技術(shù)細(xì)節(jié)興趣的讀者可以參考DISC: A Dynamic Shape Compiler for Machine Learning Workloads.
從PAI團(tuán)隊(duì)三年前啟動深度學(xué)習(xí)編譯程序方向的工作以來，「Dynamic Shape」問題一直是阻礙實(shí)際業(yè)務(wù)落地的嚴(yán)重問題之一。包括XLA在內(nèi)的主流深度學(xué)習(xí)框架，都是基于Static Shape語義的編譯程序框架。即，just-in-time運(yùn)行的編譯程序，在運(yùn)行時(shí)捕捉待編譯子圖的實(shí)際輸入shape組合，并且為每一個輸入shape組合生成一份編譯結(jié)果。
Static Shape Compiler的優(yōu)勢顯而易見，編譯期完全已知靜態(tài)shape信息的情況下，Compiler可以作出更好的優(yōu)化決策并得到更好的CodeGen性能，同時(shí)也能夠得到更好的顯存/內(nèi)存優(yōu)化plan和調(diào)度執(zhí)行plan；然而，Static Shape Compiler的缺點(diǎn)也十分明顯，具體包括：
? 編譯開銷的增加。對于訓(xùn)練業(yè)務(wù)，編譯開銷導(dǎo)致訓(xùn)練迭代速度不穩(wěn)定，訓(xùn)練初期顯著負(fù)優(yōu)化，甚至整個訓(xùn)練過程的時(shí)間開銷負(fù)優(yōu)化；對于Inference業(yè)務(wù)，很多業(yè)務(wù)實(shí)際部署和迭代時(shí)不允許出現(xiàn)性能抖動，而脫機(jī)的預(yù)編譯預(yù)熱又會使得部署的過程變復(fù)雜。
? 內(nèi)存顯存占用的增加。除編譯開銷的問題之外，當(dāng)shape變化范圍特別大的時(shí)候，編譯緩存額外占用的內(nèi)存顯存，經(jīng)常導(dǎo)致實(shí)際部署環(huán)境下的內(nèi)存/顯存OOM，直接阻礙業(yè)務(wù)的實(shí)際落地。
? 對于一部分業(yè)務(wù)場景，shape變化范圍可能非常大甚至是趨于無窮的，比較常見的包括廣告推薦類業(yè)務(wù)中常見的稀疏化模型，還有例如分布式訓(xùn)練下的embedding切片等等。在這種情況下，編譯緩存永遠(yuǎn)也無法收斂，用戶也就不可能通過compiler獲取到性能收益了。
? 上述問題在部分情況下，可以通過人工干預(yù)Compiler的圈圖過程來緩解，即，將shape變化劇烈的子圖排除在編譯范圍之外。然而，這種解決辦法對用戶非常不友好，大大降低了Compiler應(yīng)用的通用性和透明性，這要求做部署和優(yōu)化的同學(xué)同時(shí)對模型結(jié)構(gòu)和compiler非常了解，且每一次模型結(jié)構(gòu)迭代時(shí)，都需要花費(fèi)額外的工作量來調(diào)整圈圖，獲得可以接受的性能效果。

關(guān)于這一問題，曾經(jīng)出現(xiàn)過若干類解決方案，包括，對Compiler在圈圖過程中的自動化干預(yù)；在編譯期內(nèi)部自動對變化維度做bucketing補(bǔ)齊并將子圖計(jì)算結(jié)果做自動的slicing。然而這些解決方案都存在各自的局限，例如前者只能適配于小部分子圖shape變化劇烈的情況，后者在很多模型上都無法得到自動slicing的完備數(shù)學(xué)推導(dǎo)。
為徹底解決這一問題，選擇基于MLIR（Multi Layer Intermediate Representation），結(jié)合團(tuán)隊(duì)過往對AICompiler中積累的部分經(jīng)驗(yàn)，打造一套完備支持Dynamic Shape語義的AI編譯程序，希望能夠徹底解決深度學(xué)習(xí)編譯程序在這部分對靈活性要求較高的業(yè)務(wù)中無法落地應(yīng)用的問題。
整體架構(gòu)
Dynamic Shape Compiler的整體架構(gòu)，及其在AICompiler中的上下文關(guān)系如下圖所示。

Compiler部分
MLIR Infrastruction
MLIR是由Google在2019年發(fā)起的項(xiàng)目，MLIR 的核心是一套靈活的多層IR基礎(chǔ)設(shè)施和編譯程序?qū)嵱霉ぞ邘?#xff0c;深受 LLVM 的影響，并重用其許多優(yōu)秀理念。
這里選擇基于MLIR的主要原因包括：
? 比較豐富的基礎(chǔ)設(shè)施支持，使得完成編譯程序的常規(guī)開發(fā)工作更為便捷，效率更好。TableGen，以及編寫常規(guī)pattern matching的graph optimization pass的簡化等。
? Open for Extension的模塊化設(shè)計(jì)架構(gòu)，這里的核心是其Dialect抽象的設(shè)計(jì)。除Dialect的concept本身，在架構(gòu)設(shè)計(jì)上，基于LLVM在傳統(tǒng)編譯期領(lǐng)域的成功經(jīng)驗(yàn)，MLIR團(tuán)隊(duì)還是展現(xiàn)出了老練的架構(gòu)設(shè)計(jì)能力，將整個MLIR架構(gòu)的設(shè)計(jì)變得很具模塊化。
? MLIR的膠水能力，使得其可以比較靈活方便地與已經(jīng)存在的優(yōu)化手段進(jìn)行集成，而非拒斥。
具體實(shí)現(xiàn)
MLIR框架的上述特性，使得可以比較方便的有選擇性的leverage部分小區(qū)已有組件，避免完全的重新造輪子，也一定程度上避免從頭徹底重構(gòu)XLA代碼帶來的巨大工作量。
這里根據(jù)過往對AI編譯程序的理解，選擇了4層比較主要的中間層抽象，包括：
? DHLO Dialect：能夠完備表達(dá)動態(tài)shape語義的操作數(shù)層計(jì)算圖抽象，主要作用是能夠用有限數(shù)量的操作數(shù)類型，描述不同前端框架的大量操作數(shù)定義，且表達(dá)足夠靈活。
? DLHLO Dialect：引入Buffer語義的計(jì)算圖抽象，用于在編譯程序流程中進(jìn)行內(nèi)存/顯存的管理和優(yōu)化。
? Loop Dialect：用于將操作數(shù)層的計(jì)算描述基于Loop等展開為指令集的計(jì)算描述，在這一層上完成了操作數(shù)fusion的CodeGen。
? GPU Dialect：為GPU編程模型中的kernel launching及各種底層原語提供中間層抽象。
下圖展示了基于MLIR的Loop Dialect等基礎(chǔ)設(shè)施，在CodeGen中實(shí)現(xiàn)最簡單的Input fusion的基本原理。對比XLA中只有高層的HLO和底層的llvm兩層中間表示，MLIR提供的Loop Dialect抽象可以直接在中間層完成fusion，很好的簡化了開發(fā)的復(fù)雜度。

這次不再贅述Compiler部分其它各個模塊的具體實(shí)現(xiàn)細(xì)節(jié)，移步MLIR小區(qū)中發(fā)起的相關(guān)細(xì)節(jié)討論：RFC，以及會議討論。
此處想著重介紹下對比于XLA，Dynamic Shape Compiler需要額外考慮的一些問題，包括：
? DHLO IR，在XLA的HLO IR基礎(chǔ)上，擴(kuò)展了一套具有完備動態(tài)shape表達(dá)能力的IR。靜態(tài)場景下，HLO IR中的shape表達(dá)會被靜態(tài)化，所有的shape計(jì)算會被固化為編譯時(shí)常量保留在編譯結(jié)果中；而在動態(tài)shape場景下，IR本身需要有足夠的能力表達(dá)shape計(jì)算和動態(tài)shape信息的傳遞。
? Placer模塊，對于Dynamic Shape Compiler來說，計(jì)算可以分為shape計(jì)算和data計(jì)算兩類，對于GPU backend而言，通常shape計(jì)算的計(jì)算量較小，launch和拷貝開銷相比較大因此通常更適合在host側(cè)完成計(jì)算。實(shí)現(xiàn)了一個簡單的單卡分圖策略，對host側(cè)和device側(cè)計(jì)算執(zhí)行不同的lowering pipeline。
? Buffer管理及Buffer優(yōu)化模塊，有別于靜態(tài)Shape編譯期能夠比較容易通過liveness分析，實(shí)現(xiàn)Buffer的復(fù)用等優(yōu)化，而在動態(tài)shape語境下，由于Buffer Size未知編譯期則不容易做到完全一致的優(yōu)化。目前使用的是動態(tài)的Buffer申請和釋放，優(yōu)化申請和釋放的時(shí)間點(diǎn)，同時(shí)后臺使用應(yīng)用層包含Cache的Allocator，來達(dá)到性能和靈活性之間的平衡。后續(xù)可考慮在IR中充分表達(dá)Shape Constraint信息的情況下，來嘗試在編譯期做精細(xì)的Buffer復(fù)用優(yōu)化。
此外，注意到在動態(tài)shape語境會為編譯期的性能performance帶來一些有趣的新挑戰(zhàn)：
? 部分優(yōu)化決策后置到運(yùn)行期，以Implicit Broadcast為例，目前主流的前端AI框架都支持implicit broadcast語義，而在動態(tài)shape語義下，編譯期無法充分知道LHS/RHS是否需要執(zhí)行Broadcast操作。為保證完備性，如果所有情況下都穩(wěn)妥的執(zhí)行Broadcast計(jì)算的話，則會帶來比較嚴(yán)重的冗余計(jì)算/Fusion顆粒度等問題。其它與之類似問題還包括GPU Kernel的Launch Dimension選擇等，解決這一問題的做法是編譯期做多版本編譯，運(yùn)行期根據(jù)實(shí)際shape來選擇最優(yōu)實(shí)現(xiàn)，保證靈活性的同時(shí)，緩解靈活性帶來的性能損耗。
? Shape約束信息的使用，發(fā)現(xiàn)在Dynamic Shape Compiler中，即使Tensor的Shape信息未知，但Shape之間的約束信息，例如兩個Tensor之間的某兩個維度的size是否相等等信息，仍然會對編譯結(jié)果的性能產(chǎn)生比較重要的影響。主要原因包括：在圖層面，這些信息帶來了更大的圖優(yōu)化空間，而在CodeGen層面，這些信息能夠更有效的指導(dǎo)低層Lowering做CSE等傳統(tǒng)編譯程序優(yōu)化，減少冗余的計(jì)算指令數(shù)。
多前端框架支持
隨著近年來PyTorch用戶數(shù)量的持續(xù)增加，對PyTorch作業(yè)的性能優(yōu)化需求也正在變得越來越重要。AICompiler框架在設(shè)計(jì)時(shí)，包含了擴(kuò)展支持不同前端框架的考慮。
從IR lowering的角度，這里選擇相比于HLO更具泛化表達(dá)能力的DHLO Dialect作為不同前端框架的統(tǒng)一接入IR，而在PyTorch側(cè)選擇用戶部署時(shí)導(dǎo)出的TorchScript IR，通過實(shí)現(xiàn)一個輕量的Converter將TorchScript轉(zhuǎn)換為DHLO IR實(shí)現(xiàn)了對PyTorch Inference作業(yè)的覆蓋。MLIR相對完備的IR基礎(chǔ)設(shè)施也為Converter的實(shí)現(xiàn)提供了便利。
RAL (Runtime Abstraction Layer）
除編譯本身的問題之外，還面臨其它一些問題，例如如何將編譯的結(jié)果能夠配合TensorFlow/LibTorch等宿主在各自的運(yùn)行環(huán)境上下文中執(zhí)行起來，如何管理運(yùn)行時(shí)IR層不易表達(dá)的狀態(tài)信息等等。希望為不同的運(yùn)行時(shí)環(huán)境實(shí)現(xiàn)一套統(tǒng)一的Compiler架構(gòu)，為此引入了運(yùn)行時(shí)抽象層，即RAL層。RAL層主要負(fù)責(zé)解決如下問題：
Compile Once and Run Anywhere
RAL實(shí)現(xiàn)了多種運(yùn)行環(huán)境的適配支持，用戶可以根據(jù)需要進(jìn)行選擇。
? 全圖編譯，獨(dú)立運(yùn)行。當(dāng)整個計(jì)算圖都支持編譯時(shí)，RAL提供了一套簡易的runtime以及在此之上RAL Driver的實(shí)現(xiàn)，使得compiler編譯出來結(jié)果可以脫離框架直接運(yùn)行，減少框架overhad，比如在支持某語音ASR模型（類transformer網(wǎng)絡(luò)）推理優(yōu)化時(shí)，使用全圖編譯將框架開銷從TF的22ms減小到4ms；
? TF中子圖編譯運(yùn)行。RAL目前實(shí)現(xiàn)了TF Driver，可以支持在訓(xùn)練/推理場景中對圈出的子圖進(jìn)行編譯執(zhí)行；
? Pytorch中子圖編譯運(yùn)行。RAL目前實(shí)現(xiàn)了Libtorch Driver，可以支持在推理場景中對圈出子圖進(jìn)行編譯執(zhí)行；
以上環(huán)境中在諸如資源（e.g. memory）管理，API語義等上存在差異，希望能夠引入一層抽象對compiler側(cè)屏蔽這些差異。RAL通過抽象出一套最小集合的API （RAL中稱為Driver），并清晰的定義出它們的語義，這樣compiler和runtime就可以在一定層度上隔離開來，簡化compiler的開發(fā)，同時(shí)通過提供這套API在不同環(huán)境下的實(shí)現(xiàn)，來達(dá)到在不同的環(huán)境中都能夠執(zhí)行編譯出來的結(jié)果的目的。
Stateless編譯
dynamic shape compiler完成一個計(jì)算圖的編譯之后，編譯的結(jié)果可能被多次執(zhí)行，而有些op的執(zhí)行是帶狀態(tài)的：
? 在device（e.g. gpu）上執(zhí)行時(shí)，對const op希望只在第一次執(zhí)行時(shí)加載并常駐device，而不是每次都引入一次host-to-device的拷貝；
? 對于需要根據(jù)具體shape信息進(jìn)行tuning的op （e.g. gemm/conv），tuning cache需要一個地方存儲；
RAL將資源初始化等帶狀態(tài)的部分抽取出來，封裝成context來管理生命周期。在代碼生成的過程中，通過簡單的注入context，將狀態(tài)量隱藏在context之后，使得compiler側(cè)看到的是一個純計(jì)算的過程。無狀態(tài)的設(shè)計(jì)一方面簡化了代碼生成的復(fù)雜度，另一方面也更容易支持多線程并發(fā)執(zhí)行（比如推理）的場景，同時(shí)在錯誤處理，回滾方面也更加容易支持。
對用戶透明的編譯模式切換
對于Dynamic Shape Compiler在AICompiler中的定位是：與原Static Shape Compiler并列的一套框架，在允許適度犧牲性能的情況下，提供對于強(qiáng)Dynamic Shape類業(yè)務(wù)的通用透明支持。
然而從用戶的角度來說，通常并不容易判斷一個Workload的更適合Dynamic Shape Compiler還是Static Shape Compiler，為此結(jié)合接耦和全量打開[link]中的工作，設(shè)計(jì)了一套編譯模式自動切換的狀態(tài)機(jī)。其基本思路是，在任務(wù)初期先選擇較為安全的Dynamic Shape Compiler，結(jié)合后臺編譯讓用戶能夠在運(yùn)行時(shí)盡早得到有性能提升的編譯執(zhí)行，并在后續(xù)執(zhí)行過程中結(jié)合資源的實(shí)際占用情況和實(shí)際運(yùn)行時(shí)的shape變化范圍來有選擇性的切換到Static Shape Compiler的執(zhí)行。
兩套compiler在運(yùn)行時(shí)的切換關(guān)系如下圖所示：

阿里云機(jī)器學(xué)習(xí)PAI平臺面向企業(yè)客戶及開發(fā)者，提供輕量化、高性價(jià)比的云原生機(jī)器學(xué)習(xí)平臺，涵蓋交互式建模、拖拽式可視化建模、分布式訓(xùn)練到模型在線部署的全流程覆蓋，百余種落地場景，全面提升機(jī)器學(xué)習(xí)工程效率。目前，PAI AICompiler已經(jīng)集成在阿里云機(jī)器學(xué)習(xí)PAI的通用推理優(yōu)化工具PAI-Blade敏捷版中，用戶可以可以參考開發(fā)文檔來快速體驗(yàn)。

總結(jié)

以上是生活随笔為你收集整理的AICompiler动态shape编译框架的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。