全面支援 PyTorch 2.0：BladeDISC 5 月~11 月新功能發布

作者：BladeDISC研發團隊

BladeDISC 上一次更新主要發布了 GPU AStitch 優化，方法來源於我們發表在 ASPLOS 2022上的論文AStitch。這一次，我們發布了 0.3.0 版本。

本次更新中 BladeDISC 社區全面支援了 PyTorch 2.0 編譯，推進了和 Torch-MLIR 社區的合作；增加了 CPU 量化編譯和倚天新硬體支援；在編譯優化方面 BladeDISC 社區增加了一系列特性，包括改進了 GPU 訪存密集計算的性能，完成了 Shape Constaints IR 功能設計和支援。

本文描述 BladeDISC v0.3.0 版本對於 v0.2.0 的主要更新內容。

PyTorch 2.0 和動態編譯支援

BladeDISC 近半年來高度關注 PyTorch 2.0 相關特性的新動態，積极參与社區的協作，在此期間完成了 TorchBlade 編譯架構的調整，更好地 PyTorch 動態編譯和訓練支援。

TorchDynamo 優化

現在使用 PyTorch nightly 版本和 BladeDISC，只需要額外的兩行程式碼改動即可完成 BladeDISC 的編譯加速:

import torch_blade # one more extra line
model = ...
compiled_model = torch.compile(model, backend='disc')

TorchBenchmark

BladeDISC 持續關注深度學習模型的通用優化，我們將 TorchBenchmark 作為優化的指南針評估和持續提升 BladeDISC 在不同模型上的魯棒性和優化效果。

image.png

TorchMLIR(MHLO) 和動態性貢獻

BladeDISC 一直是 MLIR/MHLO 動態性的主要貢獻和推動者。這個版本中我們與位元組 AML 團隊合作往 Torch-MLIR 貢獻了 Torch-to-MHLO 的模組，特別是對動態性的支援，請參考 [RFC] Adding Torch to MHLO conversion #999，特別感謝位元組 AML 的同事的共同推動!

image.png

並且我們對 BladeDISC 中的 PyTorch 編譯流程做了架構調整，Torch-MLIR 成為了 BladeDISC 的基礎模組。下圖中藍色線框表示了 Torch-to-MHLO 工作在 BladeDISC 中相對的位置。

image.png

我們號召社區對此模組感興趣的朋友們一起合作推動 Torch-MLIR(MHLO) 模組的迭代與演進，減少基礎工作的碎片化。

PyTorch 訓練支援

BladeDISC 正在逐步支援 PyTorch 模型的訓練優化，目前已經成功支援 BERT 等模型的編譯優化。BladeDISC 支援採用 PyTorch 2.0 的 TorchDynamo 和 Lazy Tensor Core 兩種方式開啟訓練優化(當前此模組的介面沒有完全確定，鑒於目前 PyTorch 社區的活躍度，我們會更優先支援 TorchDynamo 的使用方式)。目前此模組仍在持續更新和迭代，更多狀態請關注我們的開發動態。

EasyCV/NLP 推理加速編譯支援

  • BEVFormer: 是一種純視覺的自動駕駛感知演算法，藉助 BladeDISC 優化，取得了 1.42 倍的端到端性能提升。
  • PAI-Diffusion Model: 在此版本中 BladeDISC 社區也對最近大熱的 AIGC Diffusion 相關模型提供了通用優化支援。BladeDISC 為 PAI-Diffusion 提供了接近 3 倍的端到端的性能提升。

更多資訊請查看 PAI EasyCV/NLP 相關係列文章和開源地址:

BladeDISC 量化 (Experimental)

在這半年中我們完成了編譯+量化結合的一系列初步摸索，在包括X86、ARM多個不同硬體平台上完成了早期的方案及性能驗證，下表中展示了我們一些初步的成果。

模型 輸入 設備 優化前PyTorch/FP32 優化後Int8量化+編譯
bert-mini 8*64 g6r / Ampere Altra / 單核 135.9 ms 39.6 ms
bert-mini 8*64 g8m / YiTian /單核 127.8 ms 31.1 ms
bert-mini 8*64 hfg7 / Cooper Lake 8369 /單核 37.5 ms 21.5 ms

近期我們將支援更多硬體平台（例如CUDA），並提供如何量化PyTorch/TensorFlow模型的示例。此外，我們也將繼續優化提升量化模型的推理性能。

BladeDISC 編譯優化改進

新硬體支援：AArch64（倚天）

在這半年中我們進一步完善了對AArch64類硬體（尤其是倚天）的支援，進行了一系列針對性改進：

  • 增加對BF16/int8 GEMM/Conv的支援，充分利用倚天硬體的能力；
  • 對Arm Compute Library進行了一系列的訂製和改進，解決其在dynamic shape及高並發場景下的各種可用性問題；
  • 訪存密集型運算元CodeGen品質改進，包括Stitch-CPU對reshape類運算元的支援從而支援更大粒度的kernel，以及op duplication策略的引入進一步減少訪存量；

GPU上訪存密集計算codegen性能增強

針對GPU上訪存密集計運算元圖程式碼生成提供了一系列的深度優化，單個LayerNorm子圖在常規推理shape下可帶來最高2X的性能提升。主要的優化功能包括：

  • Fusion中若干獨立控制流的合併。比如一個stitch fusion中包含多個獨立的且shape相同的row reduce計算，則將這幾個獨立row reduce計算的控制流合併為同一個，一方面減少不必要的計算，另一方面增大ILP；
  • Row-reduce的schedule選擇邏輯優化。針對不同的shape，選擇更加合適的row reduce計算的的schedule；
  • 優化element-wise fusion的向量化優化。通過指令交叉來支援element-wise fusion的數據讀和計算的向量化；
  • Loop相關優化。添加loop unroll及instruction interleave優化，增大ILP；添加循環不變數外移優化，減少不必要的計算；
  • 消除kernel的無效argument ，減小kernel launch開銷。

以上功能可以通過設置變數DISC_MEM_INTENSIVE_OPT_EXPERIMENTAL=true來打開。

Shape Constraint IR

在這半年中我們完成Shape Constraint IR的設計和開發，通過將shape constraint作為第一等公民引入到IR中，可以方便我們充分挖掘計算圖中蘊含的結構化約束，並以此來輔助完成一系列動態shape語意下的優化，進一步縮小與static shape compiler在優化空間上的差異。感興趣的讀者可以在這裡了解我們的設計文檔，也可以閱讀我們在知乎上分享的技術文章 (link1, link2)。

對二次訂製開發支援的增強

我們基於MLIR社區PDL的工作重構了BladeDISC中接入一個custom library call流程，極大的簡化了相關的開發工作量。在新的方式下，用戶只需要提供一個PDL的pattern描述文件，以及一個符合BladeDISC runtime介面的kernel，便可以在不重新編譯BladeDISC的情況下，實現pattern的替換及對應kernel的調用，明顯改進了BladeDISC對二次訂製開發的支援。我們在量化這個場景下檢驗了新的基礎設施的可用性及工程效率，感興趣的同學可以參考這裡和這裡的例子。後續我們還將藉助PDL和transform dialect進行進一步拓展，使得不僅僅是對custom kernel，也能對特定pattern的CodeGen策略進行訂製。

Runtime Abstraction Layer 改進

  • 大模型權重的支援：我們重構了常量（比如模型的權重參數）編譯結果的存儲格式，從基於protobuf改成了自定義的格式， 從而去除了對const上限的約束，方便我們支援大模型的優化。
  • 並發性能改進：針對高並發場景（比如同時服務數百路並發的推理請求）進行了一系列的優化，進一步縮小了kernel 調度以及共享資源鎖同步的開銷，在某語音識別200+並發的場景中進一步取得20%+的性能改進。

Ongoing Work

CUTLASS Gemm CodeGen

在GPU上接入了CUTLASS進行計算密集型運算元的運算元融合與程式碼生成，自動化地將GEMM及後續的element-wise（如GELU等激活函數子圖、Tranpose運算元等）進行計算融合與程式碼生成。目前GEMM + GELU及GEMM + Transpose的通路已經走通，且在BERT模型上取得了加速效果，魯棒性正在提升中，可使用DISC_ENABLE_COMPUTE_INTENSIVE_FUSE=true設置來嘗試使用。

MLIR Transform Dialect Based CodeGen

我們目前正在基於MLIR社區的transform dialect進行計算密集型相關pattern在dynamic shape語意下的程式碼生成，目前第一個目標是在AArch64平台上GEMM相關pattern達到與ACL相當的性能，以期通過白盒的方式徹底解決ACL在dynamic shape以及服務端多路請求並發的場景下的可用性問題，相關的工作的最新進展可以參見這裡。

推薦類稀疏模型

我們針對在業界中廣泛應用的Tensorflow推薦類模型的性能熱點部分–FeatureColumn中的稀疏運算元進行了初步的優化支援。目前已經完成推理場景中常見的稀疏運算元的cpu codegen支援，以及初步的運算元融合支援，基於此目前已經可以在部分模型上獲得了一定的收益。後續我們將支援更多種類，更大粒度的運算元融合以及使用CPU AVX等指令集優化稀疏部分運算元的計算性能，相關進展參見這裡。

以上為本次 release 的主要內容。

項目開源地址：//github.com/alibaba/BladeDISC

Tags:

VirMach 便宜 VPS

QNews

QNews