作者：Amar Phanishayee
譯者：夏夜
發佈時間：2019 年 11 月 15 日 10:00
本文介紹了微軟新提出的 PipeDream，旨在使深度學習網絡訓練並行化水平更高，進而提高訓練效率。
深度神經網絡（DNNs：Deep Neural Networks）已經在大量應用中取得了巨大進展，這些應用包括圖像分類、翻譯、語言建模以及視頻字幕等。但 DNN 訓練極其耗時，需要多個加速器高效地並行化。
在論文“ PipeDream: Greneralized Pipeline Parallelism for DNN Training ”（該論文發表於第 27 屆 ACM 操作系統原理研討會： SOSP 2019 ）中，微軟系統研究小組的研究員，與來自卡內基梅隆大學和斯坦福大學的同事和學生們一起提出了一種 DNN 並行化訓練的新方法。正如論文裏在大量模型上展示的那樣，這套被稱爲 PipeDream 的系統比傳統的訓練方法快了最多 4.3 倍。
DNN 訓練是在前向和後向通道計算中迭代進行的。在每輪迭代中，訓練過程循環處理輸入數據的一個 minibatch，並且更新模型參數。最常見的 DNN 並行化訓練方法是一種被稱爲數據並行化的方法（見下圖 1），這種方法把輸入數據分散到各個 workers（加速器）中運行。
不幸的是，儘管在數據並行化加速方面有一些性能優化的進展，但若要放在雲基礎設施上訓練就會產生很大的通信開銷。而且，GPU 計算速度的飛快提升，更進一步地把所有模型訓練的耗時瓶頸轉向了通信環節。
不那麼常見的並行化訓練形式是模型的並行化（見下圖 1），是把算子分散到各個 worker 上計算的，這在以往通常用於訓練大型 DNN 模型。模型並行化也遇到了挑戰：它不能高效地利用硬件資源，並且需要程序員決定怎樣按照給定的硬件資源來分割特定的模型，這其實無形中加重了程序員的負擔。
PipeDream 是微軟研究院 Fiddle 項目開發的一個系統，它引入了流水線並行化，這是一種 DNN 並行化訓練的新方法，結合了傳統的 batch 內並行化（模型並行化和數據並行化）和 batch 間並行化（流水線）。
圖 1：傳統的 batch 內並行化訓練方式（如數據並行化和模型並行化）對硬件利用率太低。左圖中，數據並行化中的單個 worker 在交換梯度數據時不得不進行通信等待。右圖是模型並行化，worker 之間只能同時處理一個 minibatch，這大大限制了硬件利用率。
使用流水線並行化訓練來解決 batch 內的並行化限制
PipeDream 重新審視了模型的並行化，希望以此來優化性能，這與以往的動機不同，以前模型並行化是因爲訓練大型模型時，訓練過程受限於數據集的大小。PipeDream 使用多個輸入數據的流水線作業，來克服模型並行訓練時硬件效率限制的問題。典型的流水線並行化設置涉及到不同 stage 之間 layer 的分割，每個 stage 都會被複制並且並行運行數據。流水線會被注入多個 batch，以使流水線滿負荷運行在穩定狀態。在大部分情況下，流水線並行化訓練比數據並行化訓練需要通信的數據要少很多，因爲它只需要在兩個 stage 邊界之間傳輸 activation 和梯度。在穩定狀態下，所有的 workers 時刻都在運轉，不像模型並行化訓練中會有停下來等待的時候（如下圖所示）。
圖 2：左圖中，我們展示了一個流水線並行化的例子，8 塊 GPU 被分配到 4 個 stage 中。通信只用於 stage 邊界的 actiavtion 和梯度上。stage 1,2 和 3 通過對其各自的 stage 進行復制來保證流水線的負載平衡。右圖展示了具有 4 個 worker 的一個流水線，展示了啓動階段和穩定階段。在這個例子中，後向處理花費的時間是前向處理的兩倍。
在 PipeDream 中克服流水線並行化訓練的挑戰
爲了獲得流水線並行化訓練的潛在收益，PipeDream 必須克服三個主要挑戰：
- 首先，PipeDream 必須在不同的輸入數據間，協調雙向流水線的工作。
- 然後，PipeDream 必須管理後向通道里的權重版本，從而在數值上能夠正確計算梯度，並且在後向通道里使用的權重版本必須和前向通道里使用的相同。
- 最後，PipeDream 需要流水線裏的所有 stage 都花費大致相同的計算時間，這是爲了使流水線得到最大的通量（因爲最慢的 stage 會成爲流水線通量的瓶頸）。
圖 3 展示了 PipeDream 工作流程的頂層概覽。給定一個模型和硬件部署方式，PipeDream 在單個 GPU 上進行短暫的運行時性能分析後，可以自動決定怎樣分割這些 DNN 算子，如何平衡不同 stage 之間的計算負載，而同時儘可能減少目標平臺上的通信量。PipeDream 即使是在不同的模型（不同點體現在計算和通信方面）和不同的平臺上（不同點體現在互聯的網絡拓撲和分層帶寬上）也能夠有效的均衡負載。
由於 DNN 並不總是在可用的 worker 間進行均等分割，所以 PipeDream 可能在某些 stage 上使用數據並行化——多個 worker 會被分配到給定的 stage 上，並行化地處理不同的 minibatch。PipeDream 使用稱作 1F1B 的調度算法來使硬件保持滿負荷運轉，同時還能達到類似數據並行化的語義。
在 1F1B 的穩定狀態下，每個 worker 爲它所在的 stage 嚴格地切換前向和後向通道，保證資源的高利用率（可忽略的流水線暫停，沒有流水線 flush），即使在常見情況下，後向通道花費的時間多於前向通道時也是如此。上面的圖 2 已經通過例子展示了這一點。與數據並行化不同，1F1B 還使用不同版本的權重來維持統計有效性。最後，PipeDream 擴展了 1F1B，在數據並行 stage 上引入了循環調度策略，保證了後向通道計算的梯度被引流至前向通道對應的 worker 上。
PipeDream 是基於 PyTorch（PipeDream 的早期版本使用 Caffe）構建出來的。我們的評估圍繞着 DNN 模型、數據集和硬件配置的不同組合進行，證實了 PipeDream 流水線並行化所帶來的訓練時間上的收益。相比於數據並行化訓練，PipeDream 在多 GPU 機器上達到了很高的精確度，性能方面，在圖像分類上有 5.3 倍提升，機器翻譯上 3.1 倍提升，語言建模任務有 4.3 倍提升，視頻字幕模型則有 3 倍提升。PipeDream 比模型並行化有 2.6 至 15 倍的性能提升，相比混合並行化有 1.9 倍提升。
如果你對 PipeDream 更多細節感興趣，可以在 GitHub 上找到源碼。
原文鏈接：