今天要介紹的文章是來自阿裡雲機器學習平台PAI與香港大學吳川教授團隊等人在 INFOCOM 2022 發表的論文，論文名是"Efficient Pipeline Planning for Expedited Distributed DNN Training"。論文提出了一個支援任意網絡拓撲的同步流水線并行訓練算法，有效減少大規模神經網絡的訓練時間。
原有技術問題
- 同步流水線并行（pipeline parallelism）在訓練輪間插入同步牆，用于累計梯度和更新模型保證模型收斂性。但同步牆會阻塞流水線，導緻同步流水線并行的 GPU 使用率低，訓練時間長。
- 異步流水線并行去掉同步牆，允許提前往流水線中加入未來訓練輪的 microbatch 來提高 GPU 使用率，導緻 microbatch 在過時的模型參數上訓練，會影響模型的收斂速度甚至導緻模型無法收斂。
新技術創新點
本文提出了一個支援任意網絡拓撲的同步流水線并行訓練算法，保證模型的收斂性能不改變，同時相比之前的同步流水線并行政策，有效減少大規模神經網絡的訓練時間。論文主要的貢獻包括：
- 一個基于遞歸最小割的GPU排序算法，通過分析GPU間網絡拓撲确定GPU的模型部署順序，保證最大化利用 GPU 間帶寬；
- 一個基于動态規劃的模型切分部署算法，高效率找到最優的模型分割與部署方案，平衡模型在每個 GPU 上的運算時間與模型切片間的通信時間；
- 一個近似最優的清單排序算法，決策每個 microbatch 在各個 GPU 上的執行順序，最小化模型的訓練時間；
新技術主要架構以及關鍵技術點
作為分布式機器學習的一種主流訓練方式，流水線并行通過同時進行神經網絡計算與中間資料通信，減少訓練時間。我們考慮混合的同步流水線并行，通過結合資料并行（data parallelism）進一步減少訓練時間。一個同步流水線并行方案包含模型切分裝置部署與 microbatch 執行排程兩個部分。上圖給出了一個 6 層神經網絡模型在 4 塊 GPU 上進行同步流水線并行訓練的示例。由圖表（a）所示，模型被切分成三個 stage，其中第二個 stage 由于其計算量較大，被複制到兩個 GPU 上通過資料并行（data parallelism）的方式訓練。圖表（b）表示模型的三個 microbatch 的具體訓練過程，包括每一個 microbatch 在每一個 stage 上的前向和後向計算。其中，由于第二個片段以資料并行方法在 GPU2 和 GPU3 上訓練，在全部 microbatch 訓練完成後通過 AllReduce 算子同步模型片段參數。
microbatch 執行排程
先考慮在給定了模型切分裝置部署方案下的執行排程問題，這個問題屬于著名的 NP-hard 的 job shop 問題的一個特殊情況。針對此，我們設計了一個基于清單排序的近似最優的排程算法，包括1）一個執行順序排序算法，決策每一個 stage 上每一個 microbatch 的前向和後向計算的執行順序；和2）一個基于上述排序的排程算法。
我們定義每個 microbatch 在一個 stage 上的一個前向或者後向計算為一個 computation block，定義每個 microbatch 在兩個連續的 stage 間的前向通信（傳遞 activations）或者後向通信（傳遞 gradients）為一個 communication block。由于一個 microbatch 在最後一個 stage 上的後向計算可以緊接在其前向計算之後，我們将這兩個計算合并成一個 computation block。假設模型被切分成了 個stage，每個 microbatch 在每輪訓練中包括 個 computation block 和 個 communication block。我們将這 個 block 按照執行順序依次排列，組織成一個有序表 。對每個 block ，我們維護一個隊列 ，用來存儲所有的其前序 block 已經被處理了而目前 block 尚未被處理的 microbatch index。換言之， 存儲全部目前準備執行 block 的 microbatch index。同時，對每個 stage ，我們維護一個執行順序隊列 ，用來表示每個 microbatch 在目前 stage 的前向和後向計算的順序，每個元素用 (microbatch index, block number)表示。
我們的排序算法具體如下。從 1 到 依次周遊 中，我們從每一個非空的 ， 隊列中取出一個 microbatch m，将其加入到 中。如果 j 是一個 computation block，我們就将 加入到 中。在周遊 J 一輪的過程中，對于每個 block，我們最多将一個 microbatch 加入到該 block 對應的 中。我們重複上述周遊過程，直到全部的 變為空，表示我們已經将全部的 microbatch 的 computation block 都加入到了執行順序隊列 中。當排序算法完成後，我們嚴格按照每個 stage 上的執行順序隊列 指定的順序對每個 microbatch 的 computation block（即該 microbatch 的前向和後向計算）排程。對于存在資料并行的 stage（即該 stage 被複制到了多個 GPU 上），我們在該 stage 全部 microbatch 計算結束完成後加入 AllReduce 算子同步模式參數。
通過理論分析，證明了該排程算法下的單輪訓練時間不會超過最優單輪訓練時間的 倍，其中 為一輪訓練中的 microbatch 數量，而 為模型在單一 stage 上的包括 AllReduce 的全部計算時間，和兩個連續 stage 間的全部通訊時間的最大值。是以，在模型切分與部署階段，我們需要最小化 來取得最優的流水線訓練效率。
模型切分與部署
遞歸 GPU 排序
實際系統裡 GPU 間的網絡拓撲複雜，本文設計了一個遞歸 GPU 排序算法（recursive device ordering，RDO）确定GPU的模型部署順序，保證最大化利用 GPU 間帶寬。GPU 間的網絡拓撲可以模組化為一個圖 ，其中圖節點 表示每個 GPU ，邊 的權重表示連結兩個 GPU 間最小帶寬大小。排序算法 RDO 将拓撲 作為輸入，利用最小割算法 [1] 找到圖 G 上權重和最小的割，并基于此圖将一分為二為 和 ，繼續作為輸入調用 RDO 算法遞歸切割，直到我們可以對全部的 GPU 排序，記為 。
基于動态規劃的模型切分與部署
基于 ，本文采用動态規劃對模型進行切分和部署。用 表示當我們将模型的前 層切分成 個 stage 部署在 GPU 到 上，且最後一個 stage 被複制到 個 GPU 上時，最優的（即最小的）單個 stage 全部計算時間和任意連續 stage 間全部通信時間的最大值。我們可以按如下公式遞歸計算出 ：
第一個部分表示将前 層切分成 個 stage 部署在 GPU 到 上，且最後一個 stage 被複制到 個 GPU 上時的最優單一 stage 計算或者任意連續 stage 間全部通信時間的最大值。第二個部分表示前 層和後 層間通信的時間，而第三個部分表示 層的包括 AllReduce 的全部計算時間。
對于一個 L 層的深度學習模型，我們有 。是以，給定模型切分的 stage 數量 ，我們可以基于動态規劃算法得到最優的 以及對應的切分方案。周遊全部可能的模型切分 stage 數量 ，我們可以找到最優的模型切分和部署方案。
實驗
我們基于 Tensorflow 1.14.1 實作了我們的同步流水線規劃算法（synchronous pipeline planning，SPP），并在兩個測試平台上驗證了我們的算法。一個測試平台由四個 GPU 伺服器構成，每個伺服器陪伴有兩塊 1080Ti GPU，通過一個 50Gbps 交換機相連。第二個測試平台為一個單機四卡伺服器，配備有四塊 Tesla V100 GPU，通過 NVLink 相連。我們選取了七個具有代表性的深度學習模型用于測試，模型具體資訊如下表：
本文将 SPP 與四個具有代表性的分布式機器學習訓練方式進行比較驗證，包括：1）資料并行；2）GPipe[2]，一個同步流水線并行架構；3）PipeDream[3]，一個異步流水線并行架構；和 4）HetPipe[4]，一個針對異構叢集的異步流水線并行架構。針對異步流水線算法 PipeDream 和 HetPipe，作者通過插入同步牆的方式将他們改為同步流水線并行算法用于比較。具體結果如下，其中 speed-up 的計算方式為
可以看到 SPP 在每個模型上都取得了超越基線算法的表現。進一步，作者将 SPP 與原版的 PipeDream 進行比較，在 V100 GPU server 上訓練 VGG19 模型，采用 ImageNet 訓練集，将 90% 驗證集準确率作為最終的訓練目标。實驗結果如下：
可以看到，盡管 SPP 與使用我們實作的同步流水線模式 PipeDream 相比，每個訓練輪訓練時間隻有很小的加速比，但在端到端的訓練過程中，SPP 比 PipeDream（采用原始的異步流水線設計）訓練速度高出了9.05%。這是因為 PipeDream 的異步流水線訓練減緩了模型收斂的進度，進一步延長了模型訓練時間。
結論
本文提出了一個支援任意網絡拓撲的同步流水線并行訓練算法，有效減少大規模神經網絡的訓練時間。算法對一個給定的深度學習模型進行模型切分與部署，充分利用 GPU 間通信帶寬，并設計一個高效的排程器，在不同的 GPU上的 stage 進行 microbatch 執行排程，最大限度地減少訓練時間。我們在兩個 GPU 測試平台上進行的比較實驗證明，我們的設計優于最先進的方法。
參考文獻：
[1] M. Stoer and F. Wagner, “A Simple Min-Cut Algorithm,” Journal of the ACM, vol. 44, no. 4, pp. 585–591, 1997.
[2] Y. Huang, Y. Cheng, A. Bapna, O. Firat, D. Chen, M. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu et al., “GPipe: Efficient Training of Giant Neural Networks Using Pipeline Parallelism,” in Proc. of NeurIPS, 2019.
[3] D. Narayanan, A. Harlap, A. Phanishayee, V. Seshadri, N. R. Devanur, G. R. Ganger, P. B. Gibbons, and M. Zaharia, “PipeDream: Generalized Pipeline Parallelism for DNN Training,” in Proc. of ACM SOSP, 2019.