AI推理与Compiler
AI推理與Compiler
AI芯片編譯器能加深對AI的理解, AI芯片編譯器不光涉及編譯器知識,還涉及AI芯片架構和并行計算如OpenCL/Cuda等。如果從深度學習平臺獲得IR輸入,還需要了解深度學習平臺如Tensorflow、TVM等。
編譯器領域的知識本身就非常艱深,和AI模型本身的關系也不是特別緊密,很難將AI建模作為發展方向,可以多關注GPGPU Architecture。即使AI芯片過氣了,GPGPU還是會長盛不衰。
OneFlow是有其獨特的設計理念和技術路線的。目前市面上已有的眾多開源框架,用戶最多的是PyTorch和TensorFlow,除此之外還有各大公司的自研框架:PaddlePaddle、MXNet、MindSpore、MegEngine等等。其中TensorFlow突出的特點是最完備,推理、serving、XLA、tensorboard等一應俱全,工業部署的主流還是TensorFlow;PyTorch的特點是易用性最好,eager執行、動態圖、跟python的交互方式等等。
本文以 Batch Norm 為例介紹了推理計算圖的具體實現,以及 MatMul 在 CPU 上的優化細節。作為 CPU 推理優化的基石,最優的推理計算圖是實現高性能 CPU 推理的前提條件,極致性能的 MatMul 計算基礎算子將為實現卷積計算中的 Im2col 和 Winograd 提供性能保障。
概述
天元(MegEngine)深度學習框架憑借「訓練與推理一體化」的獨特范式,能夠極大程度上(90%)節省模型從研發到部署的整體成本,降低轉換難度,真正實現小時級轉化;同時,天元(MegEngine)在 CPU 推理方面所做的大量優化工作,也使得開發者在推理時能夠發揮出處理器的最佳性能。
CPU 推理
CPU 推理的性能優化至關重要,經過優化的推理性能可以較未經優化的原始性能提升數十倍。
在進行模型推理階段的優化時,首先需要對模型的計算圖進行優化,以避免冗余的計算與訪存,確保計算圖在推理時是最優的;其次,在大多 CV 相關模型中,卷積的計算比重最高,達到模型總計算量 90% 以上,因此對模型推理的優化主要聚焦在對卷積計算的優化上。
卷積計算的實現方式有 3 種:direct 卷積,Im2col 卷積以及 Winograd 卷積。
? direct 卷積:根據卷積的計算公式直接對 FeatureMap 上進行滑窗計算;
? Im2col 卷積:根據卷積計算需要在輸入通道上進行 reduce sum 的特點,將卷積運行轉化為 MatMul 計算;
? winograd:在保證計算無誤的前提下,使用加法替代乘法,達到優化卷積乘法計算量的目的,在中間過程需要使用 MatMul 進行計算。
以 Im2col 或 Winograd 方法進行的卷積計算,頻繁使用到 MatMul,對 MatMul基礎算子的優化尤為重要。
本文將首先介紹模型優化中的圖優化,然后介紹基礎算子 MatMul 在 CPU 上的優化方法。
推理計算圖優化
在訓練階段定義模型的計算圖,主要是為了滿足模型參數的訓練需求。當訓練結束,模型參數固定后,對計算圖的進一步優化能夠幫助模型推理更加高效。
推理和訓練的計算圖是一張有向無環圖(DAG),在天元中,開發者能夠以類似 LLVM 的方式對 DAG 計算圖定義許多優化方法,這里簡稱 OptPass。OptPass 可以根據用戶的配置有選擇性的加入到圖優化的 OptPass 列表中,從而幫助用戶靈活地為圖優化定義 OptPass。
計算圖優化
天元定義了多個為推理進行計算圖優化的 OptPass,開發者使用這些 OptPass 之后,將得到一張用于推理的最優計算圖。下面以 MobileNetV1 中, Convolution+Batch Norm+Relu 這樣的典型結構經過計算圖優化之后 Fuse 為 ConvBias 的過程為例,介紹天元的計算圖優化過程。
由于 Batch Norm 在推理階段除了輸入 Tensor 外都是常數,可以簡化為多個 elemwise 的組合,天元實現了一個 ConvertBatchNormToElemwisePass 的 Optpass,這個 OptPass 將模型中的所有 Batch Norm 轉化為 Elemwise,具體轉化原理如下:
緊接這一過程,天元將運行 OptPass ParamRedistributePass,該 OptPass 會將上述 Batch Norm 轉化而來的 Elemwise 中的 scale 融合到 convolution 的權重中,具體實現原理如下:
天元將運行 OptPass FuseConvBiasNonlinPass,該 OptPass 會將計算圖中的 Convolution+Elemwise 轉化為天元內部實現的 ConvBias Op 中,還會設置 ConvBias Op 中 NonelineaMode 參數。
如此便完成了從 Convolution+Batch Norm+Relu 到 ConvBias 的轉換,整體轉換過程如下圖:
實驗驗證圖優化之后的性能
在推理之前,如果要對完成訓練的模型進行圖優化,則需要在模型 dump 的期間進行,當然,也可以在 SDK 運行模型之前進行。下面是在模型 dump 時進行圖優化的代碼:
from megengine.jit import trace@trace(symbolic=True)def pred_fun(data, *, net): net.eval() pred = net(data) pred_normalized = F.softmax(pred) return pred_normalized # 使用 trace 類的 trace 接口無需運行直接編譯pred_fun.trace(data, net=xor_net)# 使用 trace 類的 dump 接口進行部署pred_fun.dump(“xornet_deploy.mge”, arg_names=[“data”], optimize_for_inference=True, enable_fuse_conv_bias_nonlinearity=True)
復制代碼
上面的 optimize_for_inference=True 將在 dump 模型時候針對 inference 進行優化, enable_fuse_conv_bias_nonlinearity=True 將在模型中進行 Op Fuse 的優化,此外天元還支持其優化參數,具體可見天元的文檔。
下表展示的,是在實驗過程中對模型進行圖優化前、后,模型運行性能的測試對比。
圖優化對模型性能的提升效果顯著,具體提升比例由模型自身決定。
MatMul 優化
MatMul 作為卷積運算的基礎算子,會頻繁地被 Im2col、Winograd 以及 FullyConnect 使用。在天元中,MatMul 既被封裝為單獨的 Op,也可以作為單獨的 algo 供卷積的實現使用。
MatMul 也是計算最為密集的算子,因此天元進行了極致的優化。
優化
MatMul 是線性代數中的矩陣乘,假設矩陣 A 大小為 MK,矩陣 B 大小為 KN,則得到矩陣 C 大小為 M*N,其中 C 的每個元素的計算公式如下:
在 MatMul 的計算中乘法和加法的計算量為 2MNK (計算 C 中每個元素時,加法和乘法計算量分別為 K,C 的總元素個數為 MN),訪存量為 2MNK (計算每個 C 中元素需要 2K 訪存)+ 2MN(整個 C 矩陣讀一次和寫一次)。由于計算量固定(排除 Strassen),所以只能優化訪存,使得乘法和加法運算達到處理器的極限性能,從而實現 MatMul 的最佳性能。
MatMul 分塊
關于減少 MatMul 計算時的訪存量,最有效的方法是對 MatMul 的計算進行分塊,并將分塊之后的數據保存在 CPU 的 Cache 中。
如下圖所示,將 A 按照 mr 進行行分塊,將 B 按照 nr 進行列分塊,計算時將需要用到的分塊保存在 CPU 的各級 Cache 中,從而極大的減少了內存的讀寫。
進一步細化上面的分塊計算過程如下圖所示,A 中的一個行塊都要重復地和 B 中的每一個列塊進行小分塊的 MatMul,寫入到 C 的小塊中,為了使得分塊盡量的大 (如上所述,能夠減少內存訪存量),Cache miss 率盡量低,因此需要根據 CPU 的 Cache 結構特點(速度:L1D>L2>L3,容量 L1D<L2<L3) 來分配 Cache 和數據塊之間的存放關系,天元中進行的分配如下:
· 將下圖中紅色 (訪問重復次數最多的 A 的行塊,計算時需要的 B 的一個列塊以及計算結果的 C 的小塊) 部分都保存在 L1 中。
· 由于計算完每一個 C 的行塊,都需要重復遍歷一次整個 B 矩陣,因此將 B 存放在 L2 中使得每次讀取 B 的一個列塊都是從 L2 中讀取,代價最小。
· 將重復訪問率最高的 C 的累加中間結果保存在速度最快的 CPU 處理器的寄存器中。
通過上面的分配策略,并結合 CPU 中資源(寄存器數量,L1D 和 L2 的大小),便可以確定最佳的 MatMul 計算中的 Nr,Kr:
· 可以根據 CPU 處理器的寄存器數量得到 mr 和 nr 的具體大小,寄存器容量 > mr*nr
· 根據 L1D Cache 的大小結合 mr 和 nr 計算出 Kr,Kr=L1D/(mr+nr)
· 再根據 L2 的大小計算出 B 矩陣中的 Nr,Nr=(L2-L1D)/Kr
在上面計算 N 時,用 L2-L1D 的原因是,由于當前 CPU 使用的 Cache 是 Inclusive 的,且 L2 是指令緩存和數據緩存合并的,另外上面 M 沒有明確限制。
在得到上面最佳 Nr、Kr、mr 和 nr 之后,進一步便可以首先對 MatMul 計算中的 N、K 進行 Nr 和 Kr 分塊,然后在 Nr、Kr 的基礎上再進行 mr 和 nr 分塊。如此,MatMul 計算中的計算訪存比達到最高,且 CPU 處理器的資源也得到充分利用。
經過分塊之后,由于在最內層的計算 Kernel 為 A(mrKr)B(Krnr)=C(mrnr)的分塊矩陣乘,決定了整個 MatMul 的計算性能,因此需要極致地優化。
Kernel 計算
最核心的計算 Kernel 進行的是尺寸為 (mrKr)x(Krnr)–>(mr*nr) 的小尺寸 MatMul,其計算示意圖如下:
如上圖所示,Kernel 在計算時會讀取 A 中一列, B 中一行,進行矩陣乘,得到 大小為 mr*nr 的 C,然后和原來 C 中的值相加,如此循環 Kr 次,完成該 Kernel 的計算。
在該 Kernel 的計算過程中乘法和加法的計算量為 2mrnrKr,訪存量為 (mr+nr)Kr+2mrnr,可以根據處理器來判斷該計算能否隱藏數據的訪存。下面以 ARM cortex A76 為例進行分析,根據 A76 的數據手冊得到:
· FP32 SIMD Load throughput=2,即單周期可以 load 8 個 float 數據
· FP32 SIMD FMLA throughput=2,即單周期可以進行 16 個乘加運算
當 Kernel 的尺寸 mr=8、nr=12、Kr=256,計算量為 49152 次乘加運算,訪存量為 5312 個 float 數據時,該計算訪存量為 9.25,大于處理器的計算訪存比 2。因此可以得出結論,如果 A 和 B 均在 L1 中,則該 Kernel 的計算不會因為數據的 Load 被阻塞,所以計算單元能夠發揮出處理器的最佳性能。
雖然在對 MatMul 進行分塊時,已經計劃將 A 和 B 這樣的小矩陣置于 L1D Cache 中,但是數據在真正運行時卻不一定都在 L1D 中,原因在于,B 矩陣的列塊在原來的大矩陣中內存并不連續,其次 A 中的一列由于內存不連續,也不能使用 SIMD 進行 load,為此需要對 A 和 B 進行數據 PACK。
MatMul 數據 PACK
上文 kernel 在計算過程中,需要同時讀取 A 矩陣 mr 行數據,而每行數據之間內存不連續,因此這種訪問對 Cache 不友好,同樣,在讀取矩陣 B 中 nr 列的時候也存在數據讀取不連續的問題,加之 A 的所有行塊和 B 中的所有列塊將被讀取多次,因此可以通過對 A 和 B 提前進行數據 PACK, 實現在后續計算中頻繁多次訪問數據時是連續的(只在第一次 PACK 時對 Cache 不友好),進而獲得巨大收益。
對矩陣 A 進行數據 PACK 是將 A 中 mr 行數據的相同列拷貝到一起,如上圖中將 A PACK 到 A’ 的步驟。重復完所有 A 中的行塊便完成了 A 矩陣的數據 PACK。B 矩陣的 PACK 操作是,將 nr 列數據拷貝到連續的內存地址中,對應上圖 B PACK 到 B’ 的過程。
實 驗
按照文介紹方式方式,天元在 X86 和 ARM 上分別對 MatMul 進行了優化。下表展示了 ARM64 上的性能測試結果,實驗平臺為 kirin 980。
首先,對該處理器進行分析可以看到,其主頻為 2.6 GHz,每個周期能夠進行 16 次乘加計算,因此其理論計算峰值為 16*2.6=41.6 Gflops。
可以看到,經天元優化的 MatMul 計算,發揮出了該處理器 90% 以上的計算性能。
OneFlow
- OneFlow的設計思路,獨特性
- OneFlow的特色一:Actor機制
- OneFlow的特色二:SBP機制
- 對人工智能/深度學習未來的看法
- 對開源的感想和總結
OneFlow是獨特的
可以說完備性和易用性這兩極分別被TF和PyTorch占據,OneFlow作為一個小團隊打造的框架,如果單純的模仿別的框架,跟TF比完備性,跟PyTorch比易用性,那是絕對沒戲的。
但深度學習框架還有第三極:性能。OneFlow的特點就是追求極致的性能,而且是分布式多機多卡環境下的橫向擴展性。OneFlow的核心設計理念就是從分布式的性能角度出發,打造一個多機多卡體驗就像使用單卡一樣容易,而且速度最快的深度學習框架。
深度學習是吞沒算力的巨獸。多機多卡的理想很豐滿,現實很骨感,用戶會遇到多機多卡加速比不高,得不償失,用戶會遇到參數梁巨大時,現有框架不支持模型并行而無法訓練。為解決此類問題,業界不僅僅改進深度學習框架自身,還研發了多種第三方插件,譬如NCCL, Horovod, BytePS,HugeCTR,Mesh-tensorflow, Gpipe等等。但是,仍只滿足一小部分需求。
分別介紹兩點OneFlow相較于其它框架的獨特設計,來體現OneFlow是如何看待分布式場景下的深度學習訓練的。
Actor——用一套簡潔的機制解決所有令人頭禿的技術難題
(關鍵特性:去中心化調度流水線數據搬運是一等公民傳輸被計算掩蓋控制邏輯被執行邏輯掩蓋)
在OneFlow的設計中,分為Compiler和Runtime兩個時期,Compiler把用戶定義的網絡、分布式環境信息等編譯成一個靜態圖的中間表示(稱之為Plan);Runtime時期,各個機器根據Plan里的Actor描述信息真實地創建屬于機器的眾多Actor,整個深度學習訓練期間,OneFlow的執行基本單元就是Actor,Actor之間通過消息機制通信,靜態執行圖上的節點就是Actor,節點之間的連邊是Register,Register存儲了Actor之間的生產者消費者信息,以及生產者Actor產生的數據。
1、Actor機制實現去中心化調度
OneFlow的運行時去中心化調度就是用Actor機制實現的。在整個由Actor構成的靜態圖中,沒有一個中心的調度節點,每個Actor都只需要關心消費的那些Actor,和消費生產出的數據的那些Actor即可。這樣在超大規模的分布式訓練場景下,完全的去中心化調度可以避免中心調度的單點性能瓶頸問題。
每個Actor內部都有一個狀態機,Actor之間的收發消息和Actor內部執行都會改變狀態。當Actor收到了執行所需要讀取的Register,且當前有空閑的Register可以生產的時候,這個Actor就執行(Act)一次,生產出一個Register。生產完以后,該Actor就向需要消費這個Register的那些消費者Actor們發消息,表示可以來讀取產生的數據了;同時該Actor還需要把消費完的那些Register還給這些Regsiter的生產者Actor們,表示用完了數據,可以回收了。Actor內部的狀態機如圖4所示。
2、Actor機制實現流水線
上面介紹了Actor的內部狀態機,Actor之間的消息傳遞和數據傳遞是依賴Register實現的。一個Actor是否能執行,只跟兩個條件相關:1)消費的那些Register是否可讀;2)生產的那些Register是否有空閑塊可寫。對于一個Register,如果運行時分配多個空閑塊,那么相鄰的兩個Actor就可以同時工作,重疊起來,這樣就實現了各個Actor之間的流水線。理想狀態下整個靜態執行圖的執行時間就是整個系統中是性能瓶頸的那個Actor運行的總時間,其余Actor的執行時間就被流水線掩蓋起來了。
舉一個例子來解釋Actor機制下的流水線是如何運轉起來的。圖2是一個由3個Actor(a,b,c)組成的靜態圖的時序圖。其中深綠色的Regst方塊表示正在被使用的Register塊,白色的Regst方塊表示同一個Register的備用空閑塊。
1)在Time0時刻,Actor a產出了一個Regst_a_0,Actor b 和 Actor c由于沒有可讀的Register,所以處在等待狀態。假設每個Actor的執行時間都是單位時間。
2)到Time1時刻,Actor a給Actor b發消息說可以來讀產出的Regst_a_0了,Actor b收到了消息,并檢查生產的Register b是否有空閑Regst塊可用,發現有可用的Regst_b_0,于是Time1時刻Actor b執行,讀取Regst_a_0,寫Regst_b_0;同時Actor a還會去看是否有空閑塊可寫,發現有,Time1時刻Actor a也在執行,寫Regst_a_1。(這里需要說明的是,Regst_a_0和Regst_a_1邏輯上是屬于同一個Register,只是空間上分成了不同的空閑塊備份而已。在深度學習訓練任務中,Regst_a_0和Regst_a_1里存放的是同一個operator產出的不同batch的數據。)于是Actor a和Actor b就并行工作起來了。Actor c由于沒有數據可讀,仍在等待。
3)到Time2時刻,Actor b 生產出了Regst_b_0,于是給下游的消費者Actor c發消息說可以來讀生產的Regst_b_0,同時給上游的生產者Actor a發消息說用完了的Regst_a_0。此時Actor a 已經把剛剛生產的Regst_a_1又發給了Actor b,Actor b檢查仍有Regst_b_1空閑,于是Actor b開始讀Regst_a_1,寫Regst_b_1;Actor c收到Regst_b_0,發現有Regst_c_0空閑,于是Actor c開始執行,讀Regst_b_0, 寫Regst_c_0;Actor a收到了Actor b用完還回來的Regst_a_0,檢查Regst_a_0所有的消費者都用完了,于是將Regst_a_0回收,標記為空閑塊,同時Actor a還可以繼續執行,寫Regst_a_2。
在上面的例子中,到了Time2時刻,其實Actor a、b、c都在工作,在深度學習訓練任務中,Time2時刻Regst_b_0、Regst_c_0存放的是Batch 0的數據,Regst_a_1、Regst_b_1存放的是Batch 1的數據,Regst_a_2存放的是Batch 2的數據。通過一個Register有多個空閑塊的設計,Actor機制就實現了流水并行。
3、數據搬運
在多機多卡的分布式環境中,各個機器和各個設備之間的數據傳輸往往是影響系統橫向擴展性的最重要因素,如果傳輸開銷可以被計算開銷掩蓋,那么分布式深度學習訓練就可以達到理想的線性加速比。相較于其框架,OneFlow把數據搬運視為跟數據計算同等地位的操作,提出“數據搬運是一等公民”的思想。
已有框架在編譯期的關注焦點是數據計算,數據搬運是背后隱式發生的,在靜態分析計算圖時略過計算和搬運的重疊編排,OneFlow在計算圖中顯式表達了數據搬運,而且在靜態分析時同等對待數據搬運和數據計算,以最大化重疊搬運和計算。
在OneFlow最終的執行圖中,數據搬運操作也是一個個Actor。除了在設備上做數據計算用的Actor以外,還有計算機內存到GPU顯存之間的數據拷貝Actor,機器之間做網絡通信的網絡Actor,負責數據的切分、合并、復制的Actor,負責讀取磁盤數據的Actor,負責加載保存模型的Actor等等。很多其框架都把數據加載、多卡模型梯度的同步、網絡、模型加載更新等分別做成一個單獨的模塊,而OneFlow的設計是所有的功能都在一張由Actor組成的靜態執行圖里實現了。OneFlow這樣的設計不僅簡潔、優雅,還非常高效。
圖3表示了沒有GPU-Direct的情況下,在OneFlow的Runtime階段,一個設備上的計算節點如果消費了另一個設備的計算節點,數據是如何搬運過去的。
4、盡可能并行
在OneFlow的設計中,所有的出發點都是希望可以盡可能并行,從而達到最優的分布式性能。比如考慮到分布式訓練模型梯度同步時,顯存到內存的傳輸帶寬高于機器之間的網絡傳輸帶寬,OneFlow會做兩級的scatter和gather操作(本機的和各個機器之間的),用于增加locality,提高整體性能;又比如在異步啟動深度學習訓練時,python端用戶的控制邏輯跟OneFlow運行時的執行圖是并行執行的,同時OneFlow有一套互斥臨界區的設計保證執行的高效性和正確性;數據加載部分無論是從磁盤讀數據還是從python端喂數據,OneFlow都能保證盡可能并行,使得計算設備不會因為要等數據而導致性能下降。
已有框架要盡可能重疊數據搬運和計算,一般借助多層回調(callback)函數,當嵌套層次過多時,會遇到所謂的callback hell麻煩,正確性和可讀性都可能下降。但在OneFlow中,以上的這些并行并發特性,都是在這一套簡潔的Actor機制下實現的,解決了令人頭禿的callback hell問題。
此外,在多機的網絡通信部分,OneFlow底層的網絡通信庫原生支持RDMA的高性能通信,也有一套基于epoll的高效通信設計。而相比于最流行的Pytorch,多機還需要通過RPC來做數據同步。
Placement+SBP——OneFlow如何做到分布式最易用
OneFlow是目前分布式場景中支持數據并行、模型并行、流水并行等最易用的深度學習框架。用戶只需要像單卡一樣去搭建網絡模型,并告訴OneFlow有哪些機器哪些卡,OneFlow就會用最高效的方式把這些機器和設備使用起來。
這源于OneFlow的一套獨特的設計:ConsistentView(一致性視角)。對于多機多卡,OneFlow會抽象成一個超級大的設備,稱之為邏輯上的設備,這個邏輯設備的顯存是實際多個物理設備的顯存之和,這個邏輯設備的算力也是實際多個物理設備的算力之和。用戶只需要定義在這個邏輯上的超級設備里,深度學習模型是如何構建的,其余的便不需要用戶來操作,由OneFlow來完成邏輯上的設備到物理上的設備的映射。
這里先明確兩個概念:“邏輯上的”和“物理上的”。“邏輯上的”表示OneFlow把分布式集群抽象成一個超級計算機之后的計算和數據,“物理上的”表示那些真實的部署到各個機器和設備上的計算和數據。深度學習網絡是由Op構成的計算圖,Op之間生產消費Tensor數據。在多機多卡的環境下,一個邏輯上的Op會對應多個真實的物理上的Op,每個物理上的Op實際執行的計算都是這個邏輯Op計算的一部分,一個邏輯上的Tensor也會對應多個物理上的Tensor,每個物理上的Tensor都是邏輯Tensor的一部分。
對于其框架定義的分布式訓練,每張卡是一個“world”,多卡之間根據暴露出來的接口來同步模型梯度;而對于OneFlow而言,多機多卡也都是一個“world”,使用一套Placement+SBP的方式做全局的統籌管理。
1、Placement
在OneFlow的計算圖搭建過程中,每個計算Op都有一個屬性叫做Placement,表示了該邏輯上的Op,是要部署到哪些機器哪些設備上的。對于常見的數據并行,就是所有的Op都部署到所有的設備上。但OneFlow也支持用戶指定Op的Placement,比如當網絡過大單卡根本放不下的時候,在OneFlow可以讓網絡的前一部分在一張卡上,后一部分在另一張卡上,用一種“接力”的方式工作,實現流水并行。
圖4展示了一種可能的Placement例子:
2、SBP
SBP是OneFlow獨有的概念,他是三個單詞的首字母組合:Split、Broadcast、PartialSum(以PartialSum為例,實際上還可以是PartialMin, PartialMax等reduce操作),全稱叫SbpParallel,表示一種邏輯上的Tensor跟物理上的多個Tensor的映射關系。
其中Split表示物理上的Tensor是邏輯Tensor按照某一維度切分后得到的, Split有個參數axis,表示切分的維度,如果把多個物理上的Tensor按照Split的維度進行拼接,就能還原出邏輯Tensor;Broadcast表示物理上的Tensor是跟邏輯上的Tensor完全相同的;PartialSum表示物理上的Tensor雖然跟邏輯上的Tensor形狀一致,但是物理上的Tensor里的值是邏輯Tensor里對應位置的一部分,如果把物理上的多個Tensor按照對應位置相加,即可還原出邏輯上的Tensor。圖5展示了SBP的簡單示例。
SbpSignature是一個SbpParallel的集合,在OneFlow的設計里是Op的屬性,描繪了一個邏輯上的Op被映射成各個設備上的多個物理上的Op以后,這些物理上的Op是如何看待輸入輸出Tensor在邏輯上和物理上的映射關系的。一個Op會有多個合法的SbpSignature,一個最簡單的合法signature就是輸入輸出都是Broadcast,這表示了這個Op需要整個邏輯上的Tensor數據。
當用戶構建的邏輯上的計算圖確定以后,OneFlow在Compiler生成分布式的物理上的執行圖時,會考慮每個Op的Placement和該Op允許的合法SbpSignature列表,在其中選擇一個傳輸開銷最小的SbpSignature作為本次訓練的SbpSignature,用于指導Compiler生成最高效的執行圖。
關于Op的合法SbpSignature的列表,舉一個矩陣乘法(matmul)的Op的例子。定義: Y = matmul(A,B) , A, B, Y都是Tensor,表示Y = AB。那么至少存在兩種合法的SbpSignature:
? Y: Split(0), A:Split(0), B:Broadcast
? Y: Split(1), A:Broadcast, B: Split(1)
兩種合法的signature在兩個設備上的示意圖如圖6所示。假設邏輯上的MatMul的輸入輸出Tensor的形狀是:
A(64, 10) X B(10, 50) -> Y(64, 50)
且該Op分布在兩個設備上。在第一種SbpSignature下,0號設備上的A是邏輯上A的前一半,1號設備上的A是邏輯A的后一半(按照第0維切分),而兩個設備上的B跟邏輯上的B完全一致,兩個設備輸出的Y分別是邏輯上的Y的前一半和后一半。同樣可以分析第二種SbpSignature。
值得一提的是,當A是數據,B是模型的時候,第一種SbpSignature就是數據并行,第二種SbpSignature就是模型并行。如果兩個相鄰的MatMul op,前一個使用第一種SbpSignature,后一個使用第二種SbpSignature,整個網絡就實現了混合并行。圖7是一個混合并行的示例,定義了Y0 = MatMul_0(A0, B0) , Y1 = MatMul_1(Y0, B1)這樣一個由兩個op組成的計算圖,其中A0, Y0, Y1是數據Tensor,B0, B1是模型Tensor。
在圖7中MatMul_0產出的Y0被MatMul_1消費,但是這兩個op對同一個Tensor的SBP看待方式是不一樣的,MatMul_0認為Y0是Split(axis=0)切分,但是MatMul_1需要一個Broadcast的Y0輸入。這時候OneFlow會自動插入一個“萬能”的Boxing Op做必要的數據裁剪、拼接、搬運和求和等等操作,使得所有的Op都可以在分布式環境下拿到想要的數據。
另外在數據并行的時候,訓練的前向模型Tensor的是Broadcast,對應反向傳播的梯度就是PartialSum,當Optimizer需要全部的梯度來更新模型時,就會觸發OneFlow的Boxing機制進行高效的梯度同步工作。
3、最易用的分布式并行框架
OneFlow的這套Placement + SBP + Boxing的機制,可以使得用戶定義的計算圖中的Op、Tensor以任意的方式分布在各個機器和各個設備上,無論是數據并行、模型并行還是流水并行,對于OneFlow而言,都只是一個特定Placement下的特定SbpSignature的組合而已,用戶可以方便的配置,也可以交給OneFlow來做自動的處理。
另外,早在微軟推出ZeRO-2框架之前,OneFlow就已經支持了類似的功能,多機多卡情況下,每個模型Tensor都只保存在其中一個設備上,降低梯度計算中的內存占用。
綜上,在編譯期,OneFlow通過在一套數學上嚴謹的形式系統來表示所有合法的并行模式,并支持編譯器較方便地自動搜索最優并行方案;在運行期,則通過Actor系統最優地、靈活地支持并行、并發執行。OneFlow的內核具有簡潔、高效和高擴展性的優點。
淺談人工智能和深度學習的未來
OneFlow區別于其它深度學習框架的核心就是對于分布式訓練的思考和設計,OneFlow致力于分布式訓練最快。
隨著深度學習的發展,模型會越來越大,訓練深度學習模型所需的算力會越來越高,同時模型增大的速度要大于GPU單卡顯存擴容的速度;訓練對算力的增長要求要大于GPU單卡算力增長的速度。所以分布式深度學習訓練會越來越常見、越來越重要。BERT、GPT-3等模型的出現印證了判斷。單個芯片的能力總是受到物理約束,相信算力增長的解決之道在于橫向擴展,而這必須從系統軟件的角度去求解,一旦解決,即使是把性能“一般”的芯片互聯起來,只要系統軟件可以讓協同工作好,就可以滿足任何算力需求。
在分布式深度學習訓練中,性能至關重要。性能和橫向擴展性是深度學習框架的核心競爭力。
如果深度學習的未來是單機單卡就能搞得定的規模的話,那么深度學習也就只能做做目前已知的一些簡單任務,深度學習的浪潮也會褪去。
開源,OneFlow還有很長的路要走
OneFlow團隊是一個小團隊,OneFlow作為一個開源框架還有很多不足。框架的易用性和完善性不如Pytorch和Tensorflow。團隊還在努力補上OneFlow的短板,同時也非常需要開源社區的支持。
總結
以上是生活随笔為你收集整理的AI推理与Compiler的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: ALD技术产品形态
- 下一篇: 自动驾驶传感器比较:激光雷达(LiDAR