阿里雲機器學習平臺PAI與香港大學合作論文入選INFOCOM 2022

近日，阿里雲機器學習平臺PAI與香港大學吳川教授團隊合作的論文”Efficient Pipeline Planning for Expedited Distributed DNN Training”入選INFOCOM（IEEE International Conference on Computer Communications） 2022，論文提出了一個支援任意網路拓撲的同步流水線並行訓練演算法，有效減少大規模神經網路的訓練時間。

作為分散式機器學習的一種主流訓練方式，流水線並行透過同時進行神經網路計算與中間資料通訊，減少訓練時間。一個典型的同步流水線並行方案包含模型切分裝置部署與微批次（micro-batch）執行排程兩個部分。

以下的兩個圖給出了一個6層神經網路模型在4塊GPU上進行同步流水線並行訓練的示例。由圖表1所示，模型被切分成三個片段，其中第二個片段由於其計算量較大，被複制到兩個GPU上透過資料並行的方式訓練。圖表2表示模型的三個微批次的具體訓練過程，其中，由於第二個片段以資料並行方法在GPU2和GPU3上訓練，在全部微批次訓練完成後透過AllReduce運算元同步模型片段引數。

圖表1模型切分裝置部署

圖表2微批次執行排程

然而，設計高效的流水線並行演算法方案仍然存在諸多挑戰，例如深度學習模型各異，每層的訓練時間也不相同，因此難以找到最優的模型切分部署方案；當前的流水線並行演算法侷限於同質化的GPU間網路拓撲，而現實機器學習叢集具有複雜的混合GPU間網路拓撲（例如，單個機器上的GPU可以透過PCIe或者NVLink連線，跨機通訊可以基於TCP或者RDMA），導致現有方案無法使用等，以上問題導致實際訓練中的GPU使用效率低。

針對以上難點，團隊提出了一個近似最優的同步流水線並行訓練演算法。演算法由三個主要模組構成：

1） 一個基於遞迴最小割的GPU排序演算法，透過分析GPU間網路拓撲確定GPU的模型部署順序，保證最大化利用GPU間頻寬；

2） 一個基於動態規劃的模型切分部署演算法，高效率找到最優的模型分割與部署方案，平衡模型在每個GPU上的運算時間與模型切片間的通訊時間；

3） 一個近似最優的列表排序演算法，決策每個微批次在各個GPU上的執行順序，最小化模型的訓練時間。

從理論上對演算法做出詳盡分析，給出了演算法的最壞情況保證。同時，在測試叢集中實驗證明團隊的演算法相對PipeDream，可以取得最高157%的訓練加速比。

INFOCOM是計算機網路三大頂級國際會議之一，涉及計算機網路領域的各個方面，在國際上享有盛譽且有廣泛的學術影響力。此次入選意味著阿里雲機器學習平臺PAI在分散式深度學習模型訓練最佳化領域的工作獲得國際學界的廣泛認可，進一步彰顯了中國在分散式機器學習系統領域有著核心競爭力。

阿里雲機器學習PAI是面向企業及開發者，提供輕量化、高價效比的雲原生機器學習平臺，一站式的機器學習解決方案，全面提升機器學習工程效率。