天天看點

OneFlow 并行特色

OneFlow 并行特色

在 Consistent 與 Mirrored 視角中,已經知道 OneFlow 提供了 mirrored 與 consistent 兩種看待分布式系統的視角,并且提前知道了 OneFlow 的 consistent 視角頗具特色。

因為在 consistent_view 下,OneFlow 提供了邏輯上統一的視角,分布式訓練時,使用者可以自由選擇資料并行、模型并行還是是混合并行。

本文繼續深入介紹 OneFlow 獨具特色的 consistent 視角,包括:

  • OneFlow 在 consistent_view 下純資料并行流程示意
  • OneFlow 在 consistent_view 下混合并行流程示意
  • 混合并行的優勢及适用場景
  • OneFlow 混合并行執行個體

網絡模型訓練的邏輯圖

先設定一個簡單的多層網絡,作為讨論并行方式的載體,其結構如下圖所示:

OneFlow 并行特色

 各層中,有 樣本 (灰色矩形)、 模型 (藍色矩形),以及作用在兩者之上的 op (圓形),為了簡化讨論,也可将樣本與模型限定為 矩陣 ,作用在它們之上的op為 矩陣乘法 。

對照上圖,很容易梳理出該網絡模型的邏輯:

  • 第0層的輸入為 Data 0 矩陣與 Model 0 矩陣,它們進行 op (矩陣乘法)運算後,輸出 Data 1
  • 第1層的輸入為 Data 1 矩陣與 Model 1 矩陣,它們進行 op 運算後,輸出 output
  • 第2層為 output 層,Data 2 作為整個網絡的輸出;當然,在更深的網絡中,它也可以作為下一層的輸入繼續參與訓練

consistent 視角下支援資料并行、模型并行與混合并行,将依次進行介紹,其中混合并行是重點。

Consistent 視角下的并行特色

純資料并行

已經知道,consistent 視角下,預設的并行方式是資料并行;而如果選擇 mirrored 視角,則隻能采用資料并行;若在調用作業函數時直接傳遞 numpy 資料(而不是使用 OneFlow 的 DataLoader 及相關算子),兩者的差別在于:

  • mirrored 視角下,采用純資料并行,需要自己根據參與訓練的卡數對資料進行切分、重組,使用 list 傳遞和接收資料;
  • 而 consistent 視角下提供了邏輯上的統一看待,資料的切分和重組交給了 OneFlow 架構完成。

下圖是 consistent 視角下,采用純資料并行的方式,實作原邏輯網絡模型的流程示意圖:

OneFlow 并行特色

 在純資料并行中,采用了2張顯示卡進行并行訓練,因為采用了 純資料并行 ,可以看到,對于原邏輯模型中的每一層,樣本資料都被平均配置設定到了各個卡上,每張卡上都擁有 完整的模型,與切分的資料進行 op 運算,最後組合各個卡上的樣本,得到完整的輸出。

純模型并行

在 consistent 視角下,也可以通過選擇純模型并行(設定方式在下文執行個體中會介紹),其流程示意圖為:

OneFlow 并行特色

 在純模型并行中,同樣是2張顯示卡進行并行訓練,原邏輯模型中的每一層中,都是 部分模型 與 完整的資料 進行 op 運算,最後組合得到完整的輸出。

值得一提的是,從上圖可以看出,各個卡上第0層的輸出,并 不能 直接作為第1層的輸入:因為模型并行中,為完成 op 操作,需要部分的模型與 完整的 資料;為了解決這個問題,OneFlow 中使用了 boxing 機制。

boxing 機制會統籌分布式訓練中各個節點的資料,并合理切分、合并到對應的卡上,除了模型并行過程中的資料重組問題外,資料并行中的反向梯度同步,也使用 boxing 機制解決。

boxing 的内部機制雖然複雜,但是對于使用者而言是透明的,僅僅是防止讀者産生迷惑才加入了 boxing 的圖示,對于本文而言,隻需要了解:OneFlow 會自動協調好分布式中資料的同步問題。

選擇最優的并行方式

資料并行與模型并行的優劣并不是一成不變的,樣本規模、模型規模及模型結構決定了分布式訓練中的綜合表現,需要具體情況具體分析。

概括而言:

  • 資料并行情況下,需要同步的資訊是反向傳播過程的 梯度,是以應該確定各個訓練節點之間的資訊同步速度要比節點内部的計算速度快,比如說 卷積層 的參數較少,但是計算量大,就比較适合使用資料并行;
  • 模型并行情況下,因為可以将邏輯上作為整體的模型 切分到各個實體卡 上,能夠解決“模型太大,一張卡裝不下”的問題,是以,對于參數量大的神經網絡層(如最後的全連接配接層),可以考慮使用模型并行。

實際上,也可以使用 混合并行,在同一個分布式訓練的不同部分,組合使用資料并行、模型并行。比如,對于神經網絡中靠前的參數較少、計算量大的層,采用資料并行;在最終的參數衆多的全連接配接層,則采用模型并行,以下是針對本文最開始的網絡模型邏輯圖的 混合并行 實作方案的示意圖:

OneFlow 并行特色

 目前,其它的主流架構對于混合并行或者不支援,或者需要深度定制,而 OneFlow 中可以通過簡單的設定,配置混合并行的分布式訓練,還可以用自由度超高的流水并行,深度優化分布式系統。

混合并行執行個體

代碼

以下腳本,在 consistent 視角下,對 MLP 模型采用了混合并行方案:輸入層與隐藏層采用(預設的)資料并行;輸出層采用模型并行并進行列切分。

代碼:hybrid_parallelism_mlp.py

更具體的解析在後文“代碼解讀”可見。

代碼解讀

以上腳本修改自3分鐘快速上手中的示例代碼,比較兩份代碼,也可以體會到在 OneFlow 的 consistent_view 下進行各種并行方案的配置是多麼的簡單,隻需要在單機的程式上稍加修改即可。

以上程式的關鍵部分有:

  • 通過 oneflow.config.gpu_device_num 接口設定參與訓練的GPU數目:
  •   flow.config.gpu_device_num(2)
  • reshape 及 hidden 采用預設的資料并行,不需要修改;輸出層通過設定 model_distribute 為 flow.distribute.split(axis=0) 變為模型并行:
  • def mlp(data):
  •     initializer = flow.truncated_normal(0.1)
  •     reshape = flow.reshape(data, [data.shape[0], -1])
  •     hidden = flow.layers.dense(
  •         reshape,
  •         512,
  •         activation=flow.nn.relu,
  •         kernel_initializer=initializer,
  •         name="dense1",
  •     )
  •     return flow.layers.dense(
  •         hidden,
  •         10,
  •         # dense為列存儲,進行split(0)切分
  •         model_distribute=flow.distribute.split(axis=0),
  •         name="dense2",

有讀者可能好奇為什麼split(axis=0)是列切分?需要說明的是,OneFlow 中的 dense 内部采用列存儲,是以以上代碼的flow.distribute.split(axis=0)确實是在做列切分。

此外,flow.layers.dense 使用 model_distribute 形參設定并行方式,其内部調用了底層更通用的 get_variable 接口建立 blob, get_variable 接口設定并行方式的形參名為 distribute。

可以看到,通過極少量的修改,就能将單機訓練程式改為分布式、混合并行的程式,這是 OneFlow 差別于其它架構的一大特色。

流水并行執行個體

在模型并行之外,OneFlow 還提供了一種靈活度更高的“流水并行”的并行方式,可以讓使用者使用 scope.placement 接口顯式指定用來運作邏輯 op的 實體硬體。

在流水并行中,整個神經網絡有的層次在一組實體裝置上,另外一些層次在另外一組實體裝置上,接力協同工作,分多個階段,在裝置之間流水執行。

在以下示例中,對 Consistent 與 Mirrored 視角中的“在 OneFlow 中使用 consistent 視角”代碼進行簡單修改,展示了流水并行模式。

完整代碼:hybrid_parallelism_lenet.py

更詳細的讨論可見後文的“代碼解讀”。

以上關鍵的代碼隻有2行,且他們的本質作用是類似的:

  • 通過 oneflow.scope.placement ,指定 hidden 層的 op 計算運作在0号 GPU 上
  •   with flow.scope.placement("gpu", "0:0"):
  •         hidden = flow.layers.dense(
  •             reshape,
  •             512,
  •             activation=flow.nn.relu,
  •             kernel_initializer=initializer,
  •             name="hidden",
  •         )
  • 通過 oneflow.scope.placement ,指定 output 層的op計算運作在第0号主機的1号 GPU 上
  •   with flow.scope.placement("gpu", "0:1"):
  •         output = flow.layers.dense(
  •             hidden, 10, kernel_initializer=initializer, name="outlayer"

scope.placement 的具體使用,可參閱 API 文檔。

人工智能晶片與自動駕駛

繼續閱讀