差不多開始,可能十一點五十包了各位觀眾大家好然後我很高興有這個機會在QPcon給大家介紹我們字節跳動團隊是AML組下面的拜他ARL團隊的一個主要的工作然後首先簡單介紹一下我自己我名字叫朱鴻宇然後在加入字節跳動之前是多倫多大學的一個博士生然後導師是Genevy然後做的主要是Machine Learning System的一個方向然後去年加入的字節跳動AML組然後現在也主要是在做AI編譯和訓練優化方面的一些這個工作然後這次分享我會大概分幾個部分展開第一個是關於我們拜他ARL這個開源項他背後的一個Motivation還有一些總體的設計然後第二部分的話會展開一些這個技術還有設計細節然後這裡會有一些比較硬核的東西然後第三個部分會展示一些我們的對外的一個用戶的接口還有一些在性能上的表現然後還會給大家做一個現場的一個小的一個demo首先可能有不像20個人會問這個問題就是拜他ARL是不是一個ARL這個名字可能會讓人產生一些誤解但是拜他並不是一套ARL它是一個含含繁蓋的範圍要廣的很多如果我用一句話來概括的話拜他ARL是我們的一整套從框架到硬件的一個AI編譯的一個解決方案那麼這裡順帶出來一個AI編譯的一個大的一個背景那整個AI的這個編譯呢主要是由台上下有兩方面的一個總體做的一個驅動上有的方面的驅動呢它就是應用端那應用端我們現在要面對處理的AI workflow的場景現在是變得越來越複雜那這裡面不僅僅是說我們的模型結構變得很多然後還有比如說這個不同的學習這個積極學習的這個框架像Python去Tensor Flow這樣子還有訓練和推理的這種場景還有各種就是非常細的這種業務場景所以我們要處理的這種AI workflow是很複雜的然後另外一個是硬件端硬件端也是比較複雜的硬件端它除了現在那個GPU的它本身自己在疊在之外然後它還有很多外面各種開發的一些高效的算子庫也是層出不窮這兩個大趨勢來導致說這幾個AI的編譯變成一個很重要的方向然後它也是我們BadIR所需要面對的挑戰的一個主要的來源然後接下來我們展開一下它的BadIR我們設計了一個總體的架構然後BadIR本身是基於MLR的一個生態來開發它的上游是LVM然後包含的大概就是這個圖中的三個藍色的比較深的這個藍色的三塊然後就是前端還有中端的那個compiler還有一個round time的三級的部分然後最上面那個部分就是front end的前端然後它往上面的主要是接入我們現在主流的幾個積學系的這個框架然後就是這三個我們請現在目前支持的就是這三個拜託是Tensor Flow Onyx然後往下是輸出一個統一的MLTO的一個格式的一個IR然後第二部分是compilercompiler它的輸入就是一個MLTO的一個IR然後進行一些編譯優化那這裡面就有很多我們自己團隊所開發的一些優化的這個PASS然後之後往下它的輸出就是一些編譯的結果然後比如說會有一些device的這種Banerys然後或者是CIA家的還有Cuda的原代碼然後也會有這個LVM、L格式的這樣子可以在GPU端運行的這個東西然後還有就是我們一個自定義的一個主極端的round time所接受的一個IR叫BYRE這是編譯的結果然後第三部分就是round time的round time的它的輸入也就是compiler的一個產出然後就是主要是我們的BYRE的格式的一個接口然後它可以去直接執行我們編譯的空間的一個產物然後也可以支持我們接入我們那個公司內部開發的一些高效的這個算子庫或者是一些第三方的算子庫然後除此之外round time的話它可以支持訓練可以推理場景也有我們也把它接入了我們自己公司內部的推理引擎然後網上也可以跟Frameworks、TensorFlow、Pytorch就先去做對解然後這三個部分他們是可以各自獨立出來工作的也可以合在一塊工作然後每個的話都有自己對應的這個CIA家和Pytor的一個接口然後稍微說主要的優勢我們自己認為的BYRE它主要的優勢大概兩個方面第一個就是可擴展性或者說比較靈活性就是我們想像這也是編譯器主要解決的問題我們想像比如說當我們有一種新的backend比如說最近Flash Tension 2這樣子的一個backend後端然後我們想接入我們的這樣一個系統如果我們沒有BYRE IR的話我們可能要對每一個前端我們都要去試配一遍這樣子的後端但是我們有了BYRE IR以後我們就可以只做一次試配然後讓所有的前端全部享受到這樣優化的結果那也就是說我們把原來從平方級別的一個開發複雜度變成了一個線性的級別然後是這樣子然後就是我們的各個部分不見可以獨立工作這也是可擴展性的一個體現然後第二個方面的話主要就是性能方面的優勢那性能方面其實來源於兩個一個是我們繼承了上游的所定義的一些IR的格式所以我們可以同時享受到上游的優化特斯的一些結果和我們自己的一個優化的一個收益然後另外一個是我們的後端我們是支持一個混合的這個這專 time 執行的時候它是一個混合後端所以可以讓不同的後端發揮各自的優勢然後讓總體的一個效果做得比較好然後第二部分是一些設計還有一些技術的一些細節首先是介紹一下前端的這個部分然後因為現在開源的生態裡面的就是Petorch的支持會相對多一點所以本次talk我就主要以Petorch為例子來講這樣一些細節然後這個圖中後面會有很多這樣子的圖示這個圖裡面塗紅的部分是我們和上游混合就是mix的一個結果就上游也有這樣pass然後我們在其上進行了一些擴展對 然後就是這樣一個表示然後綠色是純粹地來自於上游的一些優化然後藍色的部分是純粹我們大家自己所開發的這樣的一些細節然後後面會沿用這樣子的這個圖例然後Petorch方面的話首先Petorch2.0它引入了一個Torch Dynamo它可以把這個用戶寫到這個Petorch代碼它就做一個編譯的這樣一個事情把它可以轉化為fx的這個圖然後整個那個拜達的這個前端也是跟這個fx是兼容的然後它把fx經過了這個Torch MLR也是上游提供的一個工具轉化為Torch Dialog然後經過我們自己所開發的這樣一個Torch2 MLTO這樣一整個模塊然後把它轉成那個MLTO格式然後這個Torch2 MLTO模塊也是我們團隊的主要的一個貢獻然後應該是去年去年這一波模塊被併入了上游所以現在上游的這個模塊現在是主要由我們團隊來進行開發和貢獻和維護然後這裡是一個前端的就是AR經過轉換的一個例子左上就是大家通常就會寫的那種Petorch的一個代碼然後轉換成fx圖然後變成一個Torch Dialog最終它前端的產物就是右下角的這樣一個MLTO格式的這樣子一個AR然後在這個地方我們是可以有一些優化可以做的我這邊就簡單的介紹一個就是在Torch2這個fx轉換的圖中這個對於Transpose的一個消除然後我們知道訓練它本身它是一個前向反向就是反覆疊帶的這樣一個模型那這個過程在fx裡面它會變成兩張圖一張前向圖然後還有一張反向圖然後這個過程是這樣就是說我用戶就定義前向的邏輯然後它Petorch它會有一個自動的AutoGrad這樣子一個模塊它去生成那個反向的邏輯那這個中間我們就發現其實我們可以定義前向圖它的每個OP是否要進行一個重計算來去改變它生成的那個反向圖的這樣子一個結果那基於這樣子的話我們可以做一些Transpose的這種消除這樣子然後這裡我就畫了一個例子來表示一個Transpose的這樣一個性能問題這個是我們發現編譯當中我們很經常需要遇到經常需要解決的這樣的一個性能瓶頸然後假設我們比如說有左邊的這樣一個計算圖它默認的狀態下就是所有Petorch都不進行重計算那這樣子我在生成那個反向的時候它為了適配比如說這樣子一個BatchMount這樣子一個Layout一個要求它的編譯器就會加一些這種Transpose的這個算子那這個東西的話在後面可能會比較耗時而且沒有什麼意義對 它不進行實際的這個操作計算那這個時候如果說我們指定說前端的這兩個Transposere-shape它是需要進行重計算的話我們就會得到右邊的那樣子一個圖然後我們可以看到右邊的那三個就是Transposere-shape這個Transpose大家如果肉眼看了知道這三個是Fuse在一塊就自動地消掉所以這樣子就達到了我們這個Transpose一些算子消除的這樣子一個目的OK然後接下來會介紹一下Compiler的這樣子一個部分對 這個圖上有一點複雜當時我大概分就是大概這個大家可以看出來主要這裡面有三個部分首先就是左邊的這一條路徑左邊的這一條從到Cat到AirTemplate這個是相對獨立的一條編譯路徑這個是我們自己所開發的跟第三方這個算子庫所對接的這樣一個路徑然後我們選擇的是Meta所開發的這個AirTemplate進行了一個對接然後中間的那一條路徑就是我們的主體的一個部分就是這個GPU-CPU的一個PathMath然後這一塊的話它的最主要的核心的部分其實就是在上面的那個Lynnog的那個部分然後它主要是把Metro經過Lynnog的Tiling的Fusion的功能然後變成底下的那個由變成底下的這個SEF或是Five或者是Vector這樣一個基本的操作那從這個之後我們可以選擇說我們是要生成一個GPU的庫拉代碼也好我們是要生成一個Hose關上面跑的C代碼也好都是可以的然後最右邊的那一條就是我們生成一個我們自己主機端所進行的一個執行的一個格式那那個的話就比較簡單的這次的分享就不會具體展開然後可以稍微展開一下左邊的這兩條然後首先是接入AI Template這樣一個部分AI Template我不知道大家熟悉不熟悉它是Metro公司開發的一個就是把Pytorch的那個Newer Network模型直接轉成庫拉代碼的這樣一個推理的框架它內部會有一個算子庫這個算子庫是基於CutLots所開發的CutLots也是NV的一個就是那個做巨人懲罰或者是康復的這樣一個模板是這樣的一個庫那個這樣一個TimePlay的這樣一個庫然後它需要你用戶去做一些這個這個tuning然後把它變成一個這個算子庫它在GM或者是康伍類的算子上是表現的會比較高效一些那AI Template本身它雖然它自己定義是一個推理框架但是其實它的算子庫扣出來給訓練用其實是沒有任何的問題所以為了揭露它的話我們就自定義了一個新的這個Dialect叫Cut然後一個Cut它就對應一個AI Template當中的一個算子然後我們實現了一個MLTO2Cut這樣一個流程它的產物呢是一個它會把部分的MLTO的圖這個算子轉化成Cut算子然後變成一個Cut的這個最後變成一個Cut的這個MLTO加一堆Cut紫圖的這樣子一個格式然後每張Cut的紫圖會去生成一個Cut執行的文件那這一步路徑主要處理一些這種計算密集的像GM或者是Cost這樣子的OP然後除把這些OP之外呢主要就是一些IO1的這樣子off然後這樣子off我們會選擇走我們自己的一個Pulling & Fusion這樣子的一個拍檯然後MLTO2Cut然後對它主要就是把MLTO轉化成Cut算子然後還有一個功能就是它會消除掉一些沒必要的也是Transforce算子就比如說左邊的這張圖它為了適配這個Dought的Layout它會要就是我需要這個前後加兩個Permutation然後但是這個這幾個是可以House在一塊變成以整個那個叫做CutGM RCR Permutation 0231這樣子的一個算子就是在ATI Play當中它是有這樣的算子那這個優化配合我們前端剛才所說的那個就是消除Transforce的優化實際上這部分帶來性能收益也是比較明顯的它大概也能夠在不同的模型下會有10%到20%這樣子的一個收益然後接下來就是這個GPU CPU Pipeline然後這裡面最重要的部分其實就是Lenoge這個Dialect然後我這邊簡單介紹一下Lenoge它也是上游所提供的一個AR的抽象它的一個精髓就是說我在一個很High Level的一個AR上面我就可以同時完成Tiling和Fusion這兩件事情就是如果大家熟悉算子生成的話可能就知道Tiling和Fusion這兩個基本上是最核心的決定我們最後Code性能的兩個事情然後就比如說這樣子一個例子我們上面有兩個OP它連在一塊就是前面是後面的Producer然後我經過Lenoge轉化以後我就可以生成Lenoge.Generic這樣一個算子然後它這邊的Index Maps還有Eternal Type這個就決定了我們的Tiling的一個邏輯然後在下面的這個Loaf包底它實際上就是決定了這樣一個Fusion的一個邏輯所以它是同時完成了TilingFusion兩件事情那這樣子的話我們在後端進行Code鍵的時候就會相對的簡單和直接那這個是我們拜託在這個Lenoge上面做的一些擴展的這些例子其實做了很多東西就是不只是這裡列的比如說我們添加了很多算子讓他們也支持Lenoge.Tiling的Fusion然後我們還自己要去適配這些算子做一些transformation然後我們也對上游的一些Fusion Transformation去做了一些改進還添加了一些其他的東西然後我就不展開這些了但做那麼多事情核心目標就是右邊這兩個然後一個就是盡可能的去壓榨我融合的算子的這個收益然後另外一個就是盡可能讓我融合好的那個算子的性能打高就是它Bound在IO上的話我要盡可能讓它極大地去壓榨它的IO的那個貸款那從最後的結果上來看基本上我們是做到這兩件事情也就是一個是把算子融合壓榨到極致就是把IOBound的算子融合壓榨到極致另外一個是把它的那個貸款壓榨到極致OK然後接下來是一個對ReduceOP所做的優化這個不是一個很Trivial的一個IOBound的算子所以我單拉出來稍微展開一下這個也是我們關於在Linoc Extension上的一個具體的例子它本身Reduce算子本身它是一個Bound在IO上的一個算子然後所以我們要去把它做Fusion但是相對於其他的Element之外的這個算子它又有不同因為它有一個Reduction維度這一條天然維度它天然是不支持Fusion的不支持切分的就是不支持去做Toweling的所以在實際的情況中如果我們遇到像右邊這樣子的就是它的Parallel維度比較小但Reduce維度很長這樣碰到右邊這樣的算子我們就會再編一首遇到問題就是我們的併行度不夠併行度不夠怎麼會導致就是我們很難就是沒辦法直接就完全壓榨到它的Drum的帶寬所以當解決方法idea上也很簡單就是說我們把Reduce的維度也把它打碎讓它去併行然後這樣子的話生成出來的算子就會有足夠的併行度但是為了支持這件事情我們是做了很多的這樣一個擴展然後這裡就展示了整個Reduce OP的Toweling的過程然後比如說我們首先會切分Reduction的維度然後保證還要保證我們的讀取的數據的連續性就是他們在內存當中連續性這樣子的一些東西然後會根據併行度來判斷是否要切Reduction維度然後我們會根據比如說Block Size 256去劃分Parallel的這樣一個維度然後最後我們會做一些Panning就是說我們Toweling好的這樣子一個每個Towel上的邏輯它可能我切出來的維度不是很規則那這個時候我要做一些Panning讓它變得規則一些對就基本上就是我們在這個就很多細節工程的問題idea上其實不複雜最後的效果其實也是比較理想的最後我們在Profile的一個結果比較理想然後最後的話就是大概介紹一下這個round time的這樣的一個部分然後這部分的這個邏輯相對比較簡單就是我就不就這主要介紹一下對外接口然後這個BRT就代表round time對外接口是分兩個部分一部分是給中端用戶中端用戶的話其實就是在運行的時候你需要加載進來兩個事情一個是你的模型邏輯模型邏輯的話那個就是一個比外格是這樣的一個就是我們編譯的一個產物然後如果你在做influence的話你可能會要把自己的位置帶進來然後就是一個運行時候的數據的邏輯就是我們的Io的數據可以帶進這部分的話Io的數據我們是可以直接跟Touch的Tensor的這個接口進行對接的然後另外一塊就是這個除了上述這個就是那個中端用戶的然後我們還有就是對於硬件後端的這樣一個接口也是三組然後第一組是跟Provider相關主要就是說我後端的硬件的那些對於各種算子的實現就可能簡單如指正成法可能複雜如一堆就是非常複雜的也有可能它就是一個Memory copy就是提供這樣一些接口實現然後另外一個就是Work-U相關那這個就是定義一個執行的執行序列就相當於Cuda Stream就這樣子一個抽象然後第三個就是內存分配的一些相關的一些邏輯然後最後簡單展示一些這個兩個用戶的接口一些例子然後這個是前單的我們的一個例子它這個Torch 2.0呢是它引入了Torch Dynamo這樣一個就是Torch.compileTorch Dynamo引入了這樣的一個接口那這樣的接口的話它會集成了就是我們在適配這樣的接口實現了一個Bata IR Compile FX這樣一個接口它是集成了我們整個端道端的一個編譯的整個流程然後運行的就是用起來很簡單這把傳承Backend然後正常執行就完事對 就是它後面會進行編譯然後最後執行的時候第一次執行是編譯過程後面的都是我們的實際的Runtime所運行的這個過程然後這裡是BAT就是Runtime所對外的用戶代碼的一些例子然後使用的時候也是我首先把錢把Compile 編譯好的這個BRE加載進來然後我們的數據是可以直接跟我們的Torch.tensor的那個進行交互的所以我們其實是在就是在Petor的層面可以直接拿到我們的這個執行的一個結果然後這個是我們做Inference的一個例子也是我後面Demo的會Demo code它的主要的一個邏輯然後我們可以這個我們可以直接從比如說像Huggingface這樣子拿到我們的這個大模型然後經過我們的那個Compile然後生生可視性文件最後去在基於BRT去做一個執行然後這個地方我們有一個小的Demo稍微對 這個是之前搭好的一個環境然後比如說我現在跑的是一個Nano GPT的一個一個Training的這樣的一個Demo然後它前面會經過一些編譯的這樣的一個過程大概會有一個幾十秒大家稍微有點耐心然後這個就是在生成Torwar的時候的那個編譯產物然後現在開始是Bacore的時候的編譯產物對 然後這邊最後底下的話就是我們一個打印出來一些輸出的時間或者是它前後向的一個執行的時間對 大概是這樣子一個情況然後比如說我們可以試一下那個Lama這個模型它實際中的話第一次的編譯可能結果會有些長所以我們為了這次Demo我們做了一些這個編譯時間上面的一些這個優化就讓它的那個速度比如說做了一些Cache或者說做了一個其他症優化讓它的速度會稍微快一點不用加等等很長時間對 Lama這個模型比較大它這樣子它跑出來的結果然後這個時間它本身是就是說它會加雜一些Pytorch的它的前後向的處理對 所以它比我們圖的執行時間會稍微長一點然後再修一下這個Lama的一個Inference這樣一個結果對 Inference它沒有那些就是前後向的處理所以比剛才的Training的那個佛臥的其實會稍微快一點對 OK這個就是我們的一個Live demo的一個展示然後然後我們的那個Link我們這個Demo代碼的Link都放在那個上面那個大家如果有感興趣的話可以自己這個在在那個線下回家可以自己去嘗試去玩這樣的東西OK然後大概總結一下我們在幾個常見的大圓模型結構上面一個性能的一個表現然後這些當然都是單卡上面的一個性能然後對比是跟那個Pytorch自帶的那個Inductor它的後端做著做的一個對比然後所有數字都是基於NVA卡的這樣子的一個結果那我們基本上在GVT還有OPT還有像Bloom這樣子的模型結構上其實提升效果是相當大然後有的會達到一些兩倍的這樣的一個情況然後在Lama上面基本上是打平的這樣子一個狀態最後總結一下就是我們在本次分享中如果大家前面有很多技術系列都忘了的話這幾個大家帶回家就可以了然後一個就是我們提出了BatR這樣子的一個從框架到硬件的一個編譯的解決方案那它主要就是幾個特性應該是它的可擴展性我們支持了幾種三種不同的方案子還有接入了很多的那個backends然後它的前端compiler和後端是可以單個自獨立工作的然後我們也是在執行的時候支持和後端的這樣子的一個做法然後讓不同的後端發揮各自的優勢然後達到一個最後的草的性然後這裡是我們的網頁還有現在開源出來的github的repo然後那個OK然後我這邊分享就結束那個大家有什麼問題想問一下這一套的優化主要是針對是GPU卡還是在設備端上也有不錯的效果比如像手機這樣的設備主要還是在GPU卡上因為我們是做的訓練的一個優化對然後手機端的話後面也可以這個其實可以往上接的對就是後面會做一些這樣的擴展就是看這個業務的這個需求我看到後面有張VBT performance這塊對吧然後是這塊有沒有和TVM這塊的比較TVM的話沒有那你個人覺得這兩個兩個產品吧或兩個編譯器產品它的補充點或者它的優缺點在哪裡就我自己使用TVM的結果就是來看它更多注重在CodeGEN在生成後端的這樣子一個就這只是我自己因為我自己使用TVM的經那個比較有限吧就是我當時用TVM的時候主要是用來做CodeGEN對然後我感覺的話就是MLR會更它的優勢其實在於一些fusion的這種這種功能就Linoc那個功能還是很強的它可以喊比我覺得比TVMTVM本身的抽象要要更合適一些這是我的理解對好謝謝