AI for Science系列(三):賽槳PaddleScience底層核心框架技術創新詳解

框架技術創新

飛槳科學計算套件賽槳PaddleScience底層技術依賴飛槳核心框架。爲了有力地支撐科學計算任務高效訓練與推理,飛槳核心框架在自動微分、編譯器、執行器和分佈式等多方面分別進行了技術創新。值得一提的是,自動微分在機制上的革新帶來了框架計算量方面的優化,結合自研編譯器CINN的加持,在典型科學計算任務2D定常Laplace模型上達到業界性能最優。

自動微分

飛槳核心框架在2.3版本中針對自動微分機制和自動微分API提供多種技術創新。

自動微分機制

自動微分在機制方面,通過設計自動微分基礎算子體系,並在其上定義linearize和transpose程序變換規則。飛槳在靜態圖中新增對前向微分的支持,實現前反向不限階的自動微分機制。創新後的自動微分機制具備良好的算子支持擴展性,目前已經支持全連接網絡,並在2D定常Laplace、3D定常圓柱繞流等典型科學計算任務上完成驗證工作。

目前的自動微分機制設計流程爲:用戶完成前向組網後,會首先將前向過程對應的網絡轉化到自動微分基礎算子體系上,然後多次調用linearize和transpose兩種程序變換實現高階前向/反向微分,最後將網絡整體轉化到飛槳原生算子體系上完成組網或者對接編譯器進行更多優化。該機制將計算拆分到更細的粒度,由於引入併合理使用切空間(前向微分空間)信息,相對於原有的自動微分機制有計算量上的優化,在2D定常Laplace模型上實測有最大1.58倍性能提升。

自動微分API

在自動微分API方面,飛槳爲外部用戶新增2個自動微分API,分別是前向自動微分forward_grad與反向自動微分grad兩種接口,均支持高階微分。爲了方便用戶使用,同時提供Jacobian、Hessian的求解,支持按行延遲計算,在複雜偏微分方程組中顯著提升計算性能。另外,也公開了實現此次創新後的自動微分核心機制API jvp與vjp,當前所有API均放在incubate空間下,處於實驗特性,會根據用戶需求及具體應用場景持續演進。上述API列表以及計算示例如下表所示。

分佈式自動並行

爲了支持用戶使用大規模網格點數據進行高效訓練,PaddleScience採用高可擴展的飛槳自動並行技術,它能根據用戶輸入的串行網絡模型和所提供集羣資源信息自動進行分佈式訓練。基於統一分佈式計算圖和統一資源圖設計,飛槳自動並行可支持任意並行策略和各類硬件集羣資源上的分佈式訓練,能將原用戶原計算圖中每個串行tensor和串行operator自動切分成分佈式tensor和分佈式operator,並插入合適高效通信來保證與原串行計算一致,如下圖所示。目前飛槳自動並行支持半自動與全自動兩種模式:

  • 模式一:半自動模式下用戶可以根據自己需要指定某些關鍵tensor和operator的切分方式,而框架會通過數據流分析方式進行自動推導剩餘的tensor和operator;
  • 模式二:全自動模式下不需要用戶標記任何切分方式,所有tensor和operator都由框架基於Cost Model自適應選擇最優的切分策略。基於自動並行所提供的數據並行、梯度累積和數據集並行加載等分佈式能力,在3D圓柱繞流任務上32卡擴展效率可達到93%。結合自動並行,PaddleScience未來將更容易獲得複雜並行策略支持。

編譯器

自動微分機制將科學計算模型中的深度學習算子拆分成若干細粒度的基礎算子,若直接計算,這些細粒度算子將使得模型的大部分時間被浪費在數據交換和執行調度而非實際計算上。這意味着我們並未完全利用到硬件的計算資源,因此模型的計算速度還仍有較大的提升空間。神經網絡編譯器在此類問題上有着顯著的優勢,包括針對計算圖的通用Pass優化、算子融合以及自動代碼生成等。飛槳神經網絡編譯器CINN即是解決此類問題的有力工具,它在消除大量無用數據交換和執行調度開銷的同時,極大地降低算子開發所需的人力成本,相比於未開啓編譯加速的執行方式,開啓CINN在2D定常Laplace模型上實測有最大3.72倍性能提升。

CINN爲科學計算模型提供了一套完善的優化機制,典型的常量摺疊,冗餘算子消除,特定硬件的gemm rewriter庫算子替換等優化Pass在CINN中均已實現。自動算子融合能力也在日趨強化,爲科學計算等領域模型的訓練性能提供更多的加速契機。以2D定常Laplace模型爲例,最大的融合算子集合可包含94個小算子。

圖 左:CINN架構圖 中:gemm rewriter示例 右:5個小算子融合示例

我們以2D定常Laplace模型作爲統一測試模型,設定採樣點數100*100,選取多組網絡層數和隱藏層節點數的全連接網絡訓練2000 epoch,分別在JAX、飛槳靜態圖、飛槳靜態圖+新自動微分、飛槳靜態圖+新自動微分+CINN編譯器幾種模式下進行測試。以飛槳靜態圖模式作爲baseline,計算其他模式相對飛槳靜態圖的加速比。(性能測試硬件平臺:V100 16GB單卡;Intel(R) Xeon(R) Gold 6148 CPU @ 2.40GHz 40-core)

從實驗結果可以看出,受益於新自動微分機制上的革新和自研編譯器CINN的性能優化加持,在典型科學計算任務2D定常Laplace模型上達到業界性能最優。

執行器

飛槳中有兩種數據結構可以表示網絡模型,一種是Program,將模型表示成OP序列,方便在單機單卡下順序執行OP;另一種是Graph,將模型表示成計算圖結構,方便在多卡下並行調度OP,並默認開啓許多圖優化的IR pass以提升訓練性能。由於Program相比Graph修改簡單,可以方便地插入分佈式需要的數據通信等OP,因而飛槳中許多重要的分佈式功能都是基於Program結構進行開發。在科學計算項目中,我們對上述典型的科學計算模型實現Program和Graph兩種表示結構的無損互轉,使得計算圖優化後的Graph可以轉換成Program進行分佈式訓練,打通計算圖優化和分佈式訓練結合的技術路徑,充分利用兩者的優化能力,大幅提升模型訓練性能。

此外,飛槳框架在2.3版本中發佈了全新的靜態圖執行器,其中實現高效的多stream和多線程異步調度組件,使得模型訓練更加性能優越和易擴展,並已在單機單卡場景下默認使用。針對PaddleScience,我們通過對靜態圖新執行器在對接Graph IR體系和支持分佈式訓練等方面的功能適配和擴充,使圖優化後的模型可以利用新執行器進行分佈式訓練。

拓展閱讀

AI+Science系列(一) :飛槳加速CFD(計算流體力學)原理與實踐

AI+Science系列(二):國內首個基於AI框架的CFD工具組件!賽槳v1.0 Beta API介紹以及典型案例分享!
賽槳PaddleScience v1.0 Beta:基於飛槳核心框架的科學計算通用求解器

相關地址

  • 飛槳AI for Science共創計劃

https://www.paddlepaddle.org.cn/science

https://www.paddlepaddle.org.cn/specialgroupdetail?id=9****

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章