技術博客神經網絡分佈式訓練中參數優先傳播方法
作者：倪昊
這篇論文來自 2019 年 SysML 會議中 Parallel & Distributed Learning 的部分。
數據並行訓練（Data parallel training） 已經廣泛地運用在在深度神經網絡的分佈式計算中，但是，分佈式計算帶來的性能提升經常受限於參數同步性能的瓶頸。作者等人提出了一種新的參數同步機制：Priority-based Parameter Propagation （P3），提高了模型的訓練集羣對網絡帶寬的利用率，加快了模型的訓練速度。 我們首先來回顧一下神經網絡的分佈式訓練。
深度神經網絡（DNN）的分佈式訓練
神經網絡的分佈式訓練通常有兩種策略。一種是模型並行，指的是分佈式系統中的不同機器分別負責在單個網絡的不同部分計算，即把模型的不同層（layer）放到不同的工作節點上，常用與模型過大的情況，它的計算效率不高。另一種是數據並行，不同的機器有着整個模型的完全拷貝，每個機器只獲得整個數據的不同部分，即把訓練數據分成許多塊，分給不同的 worker 節點，各個節點單獨進行計算，但是共享一個模型，計算的結果通過某些方法結合起來。當然，這兩種策略並不衝突，完全可以混合使用。
數據並行常用的架構是參數服務器（Parameter Server），即 PS 架構。
數據並行有兩種不同的並行模式：同步訓練和異步訓練。同步訓練是指所有 worker 節點的梯度都計算完成，統一進行梯度更新。而異步訓練是指 worker 節點計算是獨立地從 PS 節點獲取模型參數，獨立計算梯度並更新。採用同步隨機梯度下降算法（SGD）算法的數據並行訓練是一種非常流行的方法。
每個節點的工作大致可以分成三步：
- 節點從 PS 節點獲取最新的模型參數，對各自負責的訓練數據進行前向傳播。
- 各個節點獨立地對不同部分的訓練數據進行反向傳播，得到每個模型參數的梯度。
- 對各個 worker 節點的梯度計算結果進行同步，利用更新模型的參數
然後不斷的循環迭代。
節點的每一次迭代都需要通過網絡進行海量參數的同步，這對網絡帶寬提出了非常高的要求。解決這個問題的一個方法是增加網絡的帶寬，但這樣做意味着極高的成本，同時隨着模型體量的增大，成本也無法估量。因此，我們應該尋求在有限帶寬下的解決方案。
近年來一種流行的方案是梯度壓縮，但是這種方案的缺點在於丟失了一部分信息，可能會影響到模型的準確度。
還有一種方法，就是提高網絡帶寬的利用率。在訓練的過程中，如果我們在每次迭代完成後纔去同步梯度，就會導致網絡流量爆炸性地增高，而在 worker 節點計算的過程中，網絡卻基本處於空閒狀態。因此，我們可以將已經計算出來的梯度在 worker 節點計算其他層梯度的時候發給 PS 節點，即讓 worker 節點與 PS 節點之間的通信和 worker 節點的反向傳播同時進行，這樣，網絡帶寬就可以得到更有效的利用。一些深度學習框架已經實現了這一點，比如 TensorFlow、MXNet 和 Caffe2。
侷限性與優化
作者等人在工作中發現了以上方案還存在一定的侷限性，通過優化，模型對網絡帶寬的利用率和計算速度可以進一步提高。
反向傳播中參數傳遞的優先級
在神經網絡模型的訓練中，我們通常是：反向傳播——用反向傳播所得的參數進行前向傳播，更細緻一點，從最後一層開始進行反向傳播的計算，再從第一層開始前向傳播，如下圖所示。
在第 N 次迭代中，我們從輸出層開始根據 Loss 計算出梯度並更新 L4 層的參數，然後遞歸地求出輸入層即 L1 層的梯度同時更新整個模型的參數，然後我們從輸入層開始，進行 N+1 次的迭代，利用更新後模型的參數一層一層地計算，直到獲得輸出層的結果，然後進行再進行反向傳播。
我們可以觀察到一點，在迭代的過程中，最先計算出來的參數（L4層的參數）總是最後被使用，而反向傳播最後計算出來的參數卻最先被前向傳播所使用。我們發現在神經網絡的計算過程中，各個層的參數從獲得到使用中間的間隔是不同的，層數靠近輸出層，這一間隔也越大。在以往的分佈式訓練中，參數往往在某一層的反向傳播結束時就開始同步，這就有可能導致 L2 層的參數還沒有同步完，但是 L1 層已經反向傳播結束，不得不等待 L2 層完成同步，L1 層完成參數同步然後再接受 PS 節點的參數更新，才能開始進行 L1 層的前向傳播。這樣，就會導致前向傳播和反向傳播之間的間隔過大。
所以，我們認爲，在參數的同步中，層數低的層相比層數高的層，應該具有更高的優先級。
上圖中（a）代表以往的同步機制，我們對其進行改進，優先同步較低的層的梯度，如（b）所示，當 L1 層的梯度計算完成，但是更高層的參數還沒有完成同步時，優先進行 L1 層的參數同步，這樣就縮短了反向傳播和正向傳播之間的間隔，但是並沒有增加網絡的負載。
參數同步中的粒度選取
參數同步所需的通信時間主要由三部分組成：
- 梯度從 worker 節點傳輸到 PS 節點的時間
- PS 節點利用梯度更新模型參數所需的時間
- PS 節點發送更新完的參數到各個 worker 節點
就如之前描述的一樣，以往我們以層爲粒度去進行參數同步，如圖（a）所示，通過實現數據的傳輸和計算並行執行來儘可能地利用空閒的網絡帶寬，縮短通信所需要的時間。
但是在模型中，如果有一層的參數非常多，比如圖（a）中的 L2 層，它完成參數同步所需的時間每一步都是 L1 和 L3 的三倍，我們會發現他會拖慢參數同步，或者說「阻塞」了參數的傳遞。比如在時間爲 4 時，只有梯度的傳輸，卻沒有進行任何參數更新的計算。這是因爲我們的參數同步是以層爲單位，每一層中參數的更新必須要等到該層所有的節點的梯度都傳輸完成，但是我們並不需要等到接收完所有節點的梯度後再開始更新參數。
所以，我們可以採取更小的參數同步粒度，比如，將 L2 層的參數分成 3 份，單獨地進行參數同步，如（b）所示。通過採取更小的粒度，我們可以儘可能地提升帶寬的利用率，直觀來講就是增大圖中不同層之間的重合度。
總結
以上兩點就是 P3 採取的優化方式，作者在 MXNet 的基礎上實現了這一機制，有興趣的同學可以訪問 GitHub 閱讀相關源碼。根據作者的測試，使用 P3，ResNet-50、Sockeye、VGG-19 等模型的訓練效率分別提升了25%、38%、66%。這種同步機制的優點在於它與其他的優化方式並不衝突，且不會影響模型的精度，但是它對分佈式模型訓練效率的提升也有限，且並不是所有的模型都有明顯的提升，特別是模型較小或者網絡帶寬十分受限時。