科普：分散式深度學習系統（二）

01-29

上一篇文章幫助大家梳理了一下分散式深度學習系統面臨的問題以及相關進展。這一篇文章接著上一篇，詳細介紹一下當把深度學習從CPU移植到GPU上後會碰到的幾個明顯問題，以及一些解決方法。

深度學習的計算模型

很多機器學習模型，包括絕大部分的神經網路（Neural Networks）、圖模型（Graphical Models）、矩陣分解（MF）等等，它們的訓練演算法都可以被抽象成一個迭代收斂（iterative-convergent）的過程。雖然針對各個模型，開發者一般需要具體地寫各種不同的計算函數，但是大體上，這些模型的代碼邏輯一般都符合下面這個流程：

$begin{equation} large mbox{for } t = 1 rightarrow T : theta^{(t+1)} = theta^{(t)} + epsilon Delta_{mathcal{L}}(theta^{(t)}, D^{(t)}) end{equation}$

這裡 ${large theta}$ 是這個模型的參數， ${large Delta}$ 是一個更新函數： ${large Delta}$ 函數每次拿當前狀態下的模型參數 ${large theta^{(t)}}$ 和訓練數據 ${large D^{(t)}}$ 當做輸入，然後計算並返回一組更新值（updates），然後把這組值直接加到當前參數 ${large theta^{(t)}}$ 上，從而將參數更新到新的狀態 ${large theta^{(t+1)}}$ ；如此不斷循環，直到達到提前給定的最大迭代次數 ${large T}$ ，或者滿足某些循環終止條件（比如，參數更新值幾乎趨近於0了）。舉個具體的例子，比如在用隨機梯度下降（SGD）訓練神經網路時， ${large theta}$ 就是這個神經網路要訓練的參數，而 ${large mathcal{L}}$ 就是你要優化的目標函數， ${large Delta_{mathcal{L}}}$ 等價於求梯度 ${large nabla_{mathcal{L}}}$ ，其對應的操作就是：在每次迭代中，取一個小batch的訓練數據 ${large D^{(t)}}$ ，計算在這組數據上目標函數 ${large mathcal{L}}$ 對於參數 ${large theta}$ 的梯度；再對原始的梯度附加一系列操作，比如梯度下降取負，乘上learning rate（公式中的 ${large epsilon}$ ），加上momentum，必要時再對梯度進行clip操作等等。隨後，我們把更新值加上參數上，即做梯度下降，並得到一組更新後的參數，並檢查當下神經網路是否已經收斂。

有了這個固定的計算流程，我們可以很容易的通過數據並行（見本專欄第一篇文章），把任何符合這種迭代收斂計算流程的模型訓練程序，並行化到多個計算節點上進行加速。假設我們有 ${large P}$ 個計算節點，按照數據並行的思路，我們把數據均分成 ${large P}$ 份並且分發到這 ${large P}$ 個節點，然後讓每個節點在被分到那份數據上單獨執行 ${large Delta_{mathcal{L}}}$ 計算函數並計算一份更新值。然後我們再想辦法把每個節點上算出來的更新值統一收集起來再更新模型參數，即可。

按照這個思路，我們可以很容易地把上面的計算流程重寫成一個分散式版本：

$begin{equation} large mbox{for } t = 1 rightarrow T : theta^{(t+1)} = theta^{(t)} + epsilon sum_{p=1}^P Delta_{mathcal{L}}(theta^{(t)}, D^{(t)}_p) end{equation}$

和前面的單機版本比較，我們在這裡給數據 ${large D^{(t)}}$ 加上了一個下標 ${large p}$ ，也就是說，這裡我們有 ${large P}$ 個計算節點在不同的數據上同時執行計算函數 ${large Delta_{mathcal{L}}}$ 。當所有節點的計算都執行完畢後，它們的結果被匯總之後（ $sum$ 符號）再用來更新參數 ${large theta}$ 。

一旦我們把這種迭代收斂的計算模式抽象出來，按照這個構造，我可以很容易地把一個單機的機器學習程序改寫成一個數據並行的版本，而不用考慮具體的模型細節。

但是，任何構造其實都隱含一些假設，我們在應用這些構造前需要適當反思一下當前應用場景是否遵從這些假設。那麼，上面這一套用來並行迭代收斂演算法的構造有什麼假設呢？大體上，我們可以總結出下面這三點：

我們希望上述訓練程序的主要計算任務集中在 ${large Delta}$ 函數上，因為真正被並行到多台機器上執行正是 ${large Delta}$ 函數 --- 只有當主要的計算開銷由多個機器共同承擔時，分散式計算才能發揮其優勢；
在每個iteration內，每個節點上的 ${large Delta}$ 運算應是相互獨立的：第一台機器上 ${large Delta}$ 函數的執行不依賴於第二台機器的結果，數據點跟數據點之間沒有依賴性（dependency）；
在上面的構造中，模型參數 ${large theta}$ 在分散式環境下是被所有節點共享的；也就是說，在每個iteration開始之前，每個節點都必須能夠獲得到當前（最新的）模型參數 ${large theta^{(t)}}$ 。在 ${large Delta}$ 函數執行完畢之後，每個節點上算出來的參數更新值 $Delta_{mathcal{L}}(theta^{(t)}, D^{(t)}_p)$ 可以很方便的被收集起來並用來更新模型參數（ ${large theta^{(t)} rightarrow theta^{(t+1)}}$ ）。

回到分散式深度學習上，我們來逐個看看這三個假設是否都滿足。

顯然，第一個和第二個假設都很容易滿足。對應到神經網路訓練上， ${large Delta}$ 函數即對應後向傳播（backpropogation）求梯度的過程，顯然，BP求梯度所涉及到的計算基本上就是神經網路訓練的主要計算負載，所以第一個假設滿足。同時，在絕大部分情況下，我們都會假設訓練數據是i.i.d.的，用隨機梯度下降訓練時，每個計算節點只需要從它所負責的那塊數據上取一個獨立小batch算一個隨機梯度（stochastic gradients）就可以了，這個計算不受限於其他節點上的計算，也沒有數據之間的依賴，所以第二個假設也滿足；再看第三個假設，在每個iteration開始之前，每個計算節點是否能夠「輕易」獲得當前的模型參數 ${large theta^{(t)}}$ ？在每個iteration的計算結束之後，我們能否同樣「輕易」地把所有節點上產生的梯度都收集起來並更新當前參數？顯然，滿足這個假設需要一些額外的條件。

在單機上進行訓練時，我們根本不需要考慮這兩個步驟，因為參數就放在（GPU）內存上，我們只需要從內存上讀或往內存上寫就可以了。但是，當在多機集群上進行訓練時，不管是在計算開始前讀取模型參數，還是在計算結束後收集多個節點上的梯度，都會涉及到網路通訊，如何保證參數共享和梯度同步即為滿足第三個條件要解決的核心問題。

參數伺服器（Parameter Server）

說了這麼多，其實就是為了引出參數伺服器（Parameter Server）。Parameter server（PS）就是為上面這種迭代收斂的計算模型而設計出的一套通訊介面和方法。顧名思義，PS的架構其實和課本上講的client-server（CS）架構差不多。PS主要抽象出了兩個主要概念：伺服器（server, or master）和計算節點（client, or worker）；伺服器裡面放了一些數據，而計算節點則可以向伺服器發數據或者請求伺服器回傳數據。有了這兩個概念，我們可以把分散式機器學習的計算流程和PS的伺服器和計算節點這兩個模塊做如下的映射：PS的伺服器端維護全局共享的模型參數 ${large theta^{(t)}}$ （我們通常稱為parameter state），而客戶端則對應到執行計算任務 ${large Delta}$ 的各個工作節點；同時，伺服器端向客戶端提供兩個主要的API： push和pull。那麼，PS架構如何滿足解決上述的參數共享和梯度同步問題呢？在每個iteration開始的時候，所有的客戶端先調用pull API向伺服器發送一個請求，請求伺服器回傳最新的模型參數 ${large theta^{(t)}}$ 。當每個計算節點收到 ${large theta^{(t)}}$ 後，它就把這份最新的參數拷貝並覆蓋到之前舊的參數（ ${large theta^{(t-1)}}$ ）上（物理上通常這些參數存儲在RAM或GPU內存上），然後執行 ${large Delta}$ 函數計算得到梯度更新值。換句話說，PS的pull API確保了每個計算節點在計算開始前都能獲取一份最新參數的拷貝。

另一方面，梯度更新值計算完畢後，每個計算節點隨後調用push API，把這組更新值發給伺服器。伺服器會收集所有計算節點發來的更新值，並且將它們「加」到當前維護的全局共享參數上（ ${large theta^{(t)}} rightarrow {large theta^{(t+1)}}$ )。在下一個iteration，當伺服器再次收到客戶端的pull請求時，它就會把更新過後的 ${large theta^{(t+1)}}$ 發出去。因此，push API確保了梯度值的收集和模型參數的更新。

Parameter Server的架構示意圖

這裡我要重點說明一點：PS裡面的「伺服器」和「客戶端」其實都是抽象的概念。雖然聽起來像是一個中心化的架構，但是在真正的實現中，並不一定會有一台專門的機器被用來當做中心伺服器；相反地，目前大部分的PS都把伺服器端實現成一個分散式存儲系統，以避免負載不均衡，並減少單機的通信瓶頸，這個我們後面再討論。

將上述過程和單機上的神經網路訓練過程作對比，我們發現一個主要的不同點：PS允許用梯度更新模型參數這一步發生在遠端，而非當前計算節點上。作為一個工具，對於PS的用戶來說，PS的優勢主要體現在它的簡單易用上：如果一個用戶想把訓練程序改成多機版，只需要在每個訓練iteration開始前調一下PS的pull API，在計算結束後調用一下PS的push API，就可以讓這個程序在多機上跑起來並且保證參數同步了，是不是很簡單？其實這就是包括DistBelief等大部分CPU上的分散式機器學習系統的架構。但是，從設計者的角度來講，想要實現這麼一個系統並提供這兩個有效的API，需要考慮的問題就複雜的多了。回到我第一篇專欄裡面介紹關於分散式系統的幾個問題，我們來一一討論：

架構：PS裡面的伺服器端具體應該怎麼設計？物理上，這個伺服器端應該放在哪？是一個單獨的高性能伺服器，還是一個專用的伺服器集群，還是和計算節點直接共享硬體（前面討論過）？計算節點怎麼和伺服器連接？
存儲：伺服器端如何維護模型參數，用什麼數據結構？怎麼保證快速的索引和讀寫？客戶端如何請求伺服器端？客戶端是否需要設計cache或者buffer來加速模型參數和梯度的獲取？
API實現：push和pull API具體應該怎麼實現，阻塞還是非阻塞？是否需要多線程？
同步：伺服器端收集梯度，以及客戶端請求發回參數怎麼保持同步性？完全同步是必須的嗎？完全非同步是否可行？同步和非同步各有什麼優點和缺點？
網路通信：push和pull設計到參數和梯度通信。那麼，如何保證通訊效率？如何有效利用網路帶寬？通信量遠大於可用帶寬的時候應該怎麼辦？
容錯：如果伺服器端宕機了，或者有某幾個計算節點crash了，整個分散式系統應該怎麼保證容錯？
如果計算節點是GPU，或者多GPU節點，問題又會怎麼樣？

上面我列舉了很多問題，其中的某些問題在我下面列舉的這幾篇文章裡面都有陸續討論。如果你對這些問題比較感興趣，我推薦按順序讀一下下面這幾篇文章：

More Effective Distributed ML via a Stale Synchronous Parallel Parameter Server. Qirong Ho, James Cipar, Henggang Cui, Seunghak Lee, Jin Kyu Kim, Phillip B. Gibbons, Garth A. Gibson, Greg Ganger, Eric P. Xing. NIPS 2013.
Exploiting Bounded Staleness to Speed Up Big Data Analytics. Henggang Cui, James Cipar, Qirong Ho, Jin Kyu Kim, Seunghak Lee, Abhimanu Kumar, Jinliang Wei, Wei Dai, Gregory R. Ganger, Phillip B. Gibbons, Garth A. Gibson, and Eric P. Xing. ATC 2014.
Exploiting Iterative-ness for Parallel ML Computations. Henggang Cui, Alexey Tumanov, Jinliang Wei, Lianghong Xu, Wei Dai, Jesse Haber-Kucharsky, Qirong Ho, Gregory R. Ganger, Phillip B. Gibbons, Garth A. Gibson, and Eric P. Xing. SoCC 2014.
Managed Communication and Consistency for Fast Data-Parallel Iterative Analytics. Jinliang Wei, Wei Dai, Aurick Qiao, Qirong Ho, Henggang Cui, Gregory R. Ganger, Phillip B. Gibbons, Garth A. Gibson, and Eric P. Xing. SoCC 2015.
GeePS: Scalable Deep Learning on Distributed GPUs with a GPU-Specialized Parameter Server. Henggang Cui, Hao Zhang, Gregory R. Ganger, Phillip B. Gibbons, and Eric P. Xing. Eurosys 2016.
Poseidon: An Efficient Communication Architecture for Distributed Deep Learning on GPU Clusters. Hao Zhang, Zeyu Zheng, Shizhen Xu, Wei Dai, Qirong Ho, Xiaodan Liang, Zhiting Hu, Jinliang Wei, Pengtao Xie, and Eric P. Xing. ATC 2017.

接下來我們詳細探討其中的某幾個問題。

從CPU到GPU

經過上面的介紹，相信Parameter Server的架構大家心裡有數了。下面我們實踐一下：用PS把一個在單機上運行的神經網路訓練程序並行化到 ${large P}$ 個機器上。按照我前面講的思路，很容易將前面的迭代循環改寫成如下這種形式：

$begin{equation} large begin{aligned} mbox{fo} & mbox{r } t = 1 rightarrow T: & texttt{pull}(large theta^{(t)}) & nabla theta^{(t)} = Delta_{mathcal{L}} (theta^{(t)}, D_p^{(t)} ) & texttt{push}(nabla theta^{(t)}) end{aligned} end{equation}$

上面這段循環會運行在所有的 ${large P}$ 個計算節點上。

如果上述程序的主要計算函數 ${large Delta_{mathcal{L}}}$ 在CPU上完成，那麼上述流程基本概括了絕大部分CPU上的深度學習（甚至更廣泛的機器學習）系統，比如DistBelief，Project Adam的基本原理。

但是，當把核心計算函數 ${large Delta_{mathcal{L}}}$ 挪到GPU上執行時，由於GPU有自己的內存，所以 ${large Delta_{mathcal{L}}}$ 所需要的輸入以及輸出也必須位於GPU內存上（以便GPU能快速讀寫）才能保證運行效率。在分散式環境下， ${large theta^{(t)}}$ 作為 ${large Delta_{mathcal{L}}}$ 的一個輸入，是從遠端pull回來的。目前絕大部分的網路通訊協議只能從一台機器的RAM發到另一台機器的RAM（Nvidia其實有GPUDirect技術，能讓兩個GPU直接通信，不過這個對硬體有一定的要求，我們後面再談），所以其實在pull執行之後和 ${large Delta_{mathcal{L}}}$ 函數執行之前，一般有一個隱藏的memcpy操作，把收到的參數 ${large theta^{(t)}}$ 從RAM移動到GPU內存上；類似地，在 ${large Delta_{mathcal{L}}}$ 算完後，同樣需要有一個memcpy操作把算出來的梯度從GPU內存移動到RAM上，才能調用push進行通信。

總結一下，有了PS的push和pull兩個API，再加上兩個memcpy，我們就設計出了一個可以在GPU集群上跑的深度學習系統了，聽起來是不是特別簡單？其實目前大部分的深度學習框架號稱自己支持分散式GPU集群也大概就是做了上面這點簡單的工作。

如果問題真的這麼簡單，那麼其實這篇專欄文章寫到這裡就算結束了。理解了的同學，只要先到Github找一個PS的實現（比如這個和這個），然後再按照上述思路就可以把任何深度學習框架，或者自己寫的神經網路訓練程序改成一個「能在GPU集群上並行跑的分散式深度學習系統了」。事實上，只要完成了上面這一步，市面上大部分的神經網路確實都可以在GPU集群上跑起來了。

然而，如果你真的動手去實現了這麼一個系統，再去弄一個GPU集群，找幾個流行的神經網路上測試一下自己的系統的性能，你會發現當你用很多台機器並行訓練某個神經網路時，你並不能獲得多大收益。比如，你用8台機器同時訓練一個VGG19，可能只能獲得2-3倍的加速；換句話說，你花8000刀買了8塊顯卡，最後發現你的系統實現讓6000刀直接打了水漂。事實上，如果一個分散式系統用了8台機器實際只2-3台機器的收益，這個系統顯然是不達標的。並且，如果你繼續買更多的機器和GPU，你會發現加速比並不會繼續提高，甚至有可能降低......

GPU快所產生的問題

接下來我們就來詳細討論一下上面這段程序在面對GPU集群的時候會出現什麼樣的問題（當然，如果你自己動手實現了，你能很容易觀測到這些問題）。

如果把上述系統部署到GPU集群上，並且仔細profile一下整個訓練過程，我們就很容易發現幾個很明顯的問題：

內存移動：在GPU內存和RAM之間拷貝數據（memcpy）會產生一定的開銷。這個開銷如果跟CPU上的計算比可能可以忽略不計。但是在高性能計算領域（比如GPU上），完成一步計算和調用一個GPU內核程序（kernel launching）花的時間是差不多的，所以這個memcpy在很多情況下會和實際計算花差不多的時間。在這種接近於1:1的比例下，這個數據拷貝的開銷就沒法忽略不計了。
通訊：機器跟機器之間的遠程通訊需要一定的時間來完成（latency）。更進一步，我們可能還需要保證參數之間的同步性，也就是說，在每個iteration，伺服器端必須等接收到所有計算節點的梯度更新值（ ${large nabla theta^{(t)}_p}$ ），然後把他們加到上個iteration的參數 ${large theta^{(t)}}$ 上，才能響應計算節點的下一輪pull請求 --- 發出新的參數 ${large theta^{(t+1)}}$ ；另一方面，客戶端（每個計算節點）在完成一輪計算並把梯度push出去後，必須等待伺服器回傳新的參數 ${large theta^{(t+1)}}$ ，才能開始下輪計算。考慮到網路狀態一般都不穩定，網路帶寬也不一定充裕；再考慮到每個機器的快慢不一（不同節點完成 ${large Delta_{mathcal{L}}}$ 所需時間有快有慢），那麼想要保證上述同步性，可能就要花相當長的時間純粹用來等待各個節點（包括伺服器）上的通信完成，或者等待最慢的計算節點完成計算任務。
GPU內存：我們都知道，相比較於RAM，GPU內存通常是非常有限的（現在最大的Nvidia GPU也就只有大概12G內存）。但是訓練神經網路的時會產生很多中間狀態值（intermediate states），如各層的activation。在後向傳播的時候，這些中間狀態值會參與梯度的計算，所以必須保存下來，並且很輕易就把GPU內存佔滿了。假設我們訓練的神經網路很大（或很深），GPU內存裝不下了，就無法進行計算了。

上面三個問題裡面，第一個和第三個問題顯然是GPU集群上特有的；第二個問題，也就是如何盡量降低多機同步的時間開銷，幾乎是分散式機器學習（甚至整個分散式系統）這個領域最精經典、最常見、也最棘手的問題，我前面列舉的很多論文就是專門討論這個問題的。為什麼我要在GPU集群這個背景下把這個問題重新拿出來討論呢？考慮下面這個問題：當我們衡量時間開銷的時候，我們到底是在乎絕對時間還是相對時間？回到上面的例子，考慮下面這兩種情況：

情況A：假設每次循環中， ${large Delta_{mathcal{L}}}$ 裡面的計算平均需要5分鐘完成，而push和pull對應的參數同步大概需要10秒完成；
情況B： ${large Delta_{mathcal{L}}}$ 需要0.4秒完成，push和pull對應的通信只需要0.1秒完成。

A和B兩個系統哪個更差？答案顯然是B，雖然從絕對時間上看，花0.1秒在通信上遠比花10秒快，但是其實在評價這個系統的時候，我們並不在意絕對時間，而更看重的是通訊和計算的相對時間（通訊計算時間比）。如果我搭了一個系統有A這種性能，其實這個系統已經足夠好，沒必要再優化了。但是對於B，每個計算節點完成一次計算後，要額外等待1/4的計算時間才能開始下次運算，這就是說有1/5的時間計算資源是閑置的。

這個問題和GPU集群有什麼關係呢？事實上，當把這類分散式系統問題拉到高性能計算（HPC）領域討論的時候，問題是截然不同的（要知道，HPC領域的研究者們通常都是在毫秒和微秒級做優化）。GPU做深度學習的相關運算一般比CPU快了好幾十倍，以前每個iteration ${large Delta_{mathcal{L}}}$ 可能在CPU要個好幾秒才能完成，但是到了GPU上可能只需要幾十毫秒。在分散式系統裡面，網路通訊的快慢受很多因素影響，包括你要通訊的內容的大小，你的通訊程序寫的好不好，你的Ethernet硬體怎麼樣、可用帶寬高不高，等等問題。由於GPU上計算時間大幅減小，通信時間和計算時間的比例會越來越大；隨著這個比例增大，對系統的要求也就越高。

下面我來舉一個簡單的例子，帶大家估算一下GPU上的深度學習對網路通信的要求到底在什麼水平。但提前聲明一點，下面的例子已經被我簡化過了，和實際情況可能有一些出入，但是在幫助理解通信量上是沒有差別的。對更實際的情況感興趣的讀者，可以參考Poseidon論文。

考慮在Nvidia Gefore Titan X (GM200)這塊顯卡上訓練AlexNet這個網路，假設我們把batchsize設成256，大概每0.25秒這塊GPU就可以完成一個iteration的計算。另一方面，AlexNet有大約61.5M個參數，也就是說，在每台機器上，每塊顯卡大概每0.25秒就會產生61.5M個梯度值。假設我們在8台機器上並行訓練AlexNet，每個機器有且僅有一塊顯卡，並且我們把所有的數值都存成單浮點數（float）。這是背景，接下來我們來動手算算具體會產生多少的通信量。

先考慮PS的客戶端。為了保證參數的同步，每個計算節點在每個iteration計算結束後，要先把這61.5M個梯度值發出去，再從PS收回61.5M個值，作為更新後的參數 ${large theta^{(t+1)}}$ 。
再考慮伺服器端：每個iteration內，伺服器端得收61.5M * 8 = 492M個梯度回來，把這些梯度加到 ${large theta^{(t)}}$ 上得到 ${large theta^{(t+1)}}$ ，再把 ${large theta^{(t+1)}}$ 分別發給每個計算節點，也就是再往外發送61.5M * 8 = 492M個參數出去。

我們先假設最理想的情況：我們有辦法將通訊時間和計算時間完全重合起來吧，那麼我們的網路帶寬要到多少才能保證這個通訊任務迅速完成而不阻塞下一輪的運算呢？

我們來算算：對於每個計算節點，我們要在0.25s的計算時間內完成發61.5M和收61.5M這兩步，共需要61.5M × 2 / 0.25s = 492M/s的網路帶寬；而對於伺服器端，我們要完成收492M的梯度再發492M的參數這兩步操作，也就是需要492M * 2 / 0.25s = 3936 M/s的帶寬。如果我們再進一步把這個M換成Ethernet領域常用的Gbps：

對於每個計算節點，所需要的帶寬大概是492M/s * 32 bit = 15744 Mbps = 15.7 Gbps；
對於伺服器端，所需帶寬3936M/s * 32 bit = 125.9 Gbps;

看到這個數據，你大概會驚呆了。你工作用的筆記本或者台式機的標準帶寬大概是1Gbps，最高檔的AWS GPU instance大概可以提供給你20Gbps的帶寬，收費$16每小時；甚至可能市面上壓根兒買不到這麼高帶寬的Ethernet。更何況我們才用了8台機器，隨著機器的增多，通信量顯然會線性增長。而且我前面給的數據是2年前的顯卡的計算速度，現在最新版的Volta可能已經是這個的2 - 4倍了，也就是說計算時間會更小 -- 上面算術的分母會更小……

不過好消息是有一個很簡單的工程方案可以極大的緩解上面這個問題。我前面說過，PS裡面的「伺服器」或「客戶端」只是一個虛擬的概念，我們在實際工程實現中，可以讓集群中的任何一個物理節點既充當計算節點（worker）又是一個伺服器（server）。這樣我們就可以把伺服器端實現成一個分散式的存儲系統 -- 我們要伺服器要維護的模型參數平均分布到多台物理機器，從而可以充分利用每台機器的通訊帶寬了。考慮這種情況，同樣是一個8台機器的集群，讓每台機器既充當伺服器又充當計算節點，每台機器作為伺服器端只需要負責1/8的參數通信。那麼每台機器每秒需要多少帶寬才能完成任務呢？相似的：每台物理節點，作為伺服器需要收一次發一次，而同時作為計算節點需要發一次收一次，所以總體上大概需要在每秒傳輸4 × 61.5M / 0.25s × 7/8 = 840 M/s（為什麼這裡要乘以7/8？讀者可以自己想想），對應的Ethernet帶寬大約是26Gbps。可以看到，這個方法雖然大大的降低了伺服器端的通訊總量，但是26Gbps還是一個太大的值，而且考慮到我們之前還做了大量的假設：我們假設了通信時間和計算時間是完全重疊的，還假設了通信的穩定性（不會有可用帶寬不穩定的情況），而在真實環境下這些假設幾乎都不可能滿足。因此，相比較於CPU環境下，分散式GPU對通信有極其高的要求。

另外，以上對於通信量的估算也詳細解釋了我在知乎上對Facebook那篇訓練ResNet論文的評論：ResNet的計算時間和其他網路比是偏長的（層數多，計算多），而參數量卻很少（都是卷積層，沒多少參數），因此通信計算時間比並不大，所以把ResNet分散式並行起來並獲得不錯的加速相比較於其他網路結構，更簡單。當然，這也從另一個側面反映了ResNet這個網路的優點。

回到這一節剛開始提的第三個問題：GPU內存空間不夠。一個可行的解決方案就是使用模型並行（model parallelism），把內存消耗過大的模型分割成很多份放到不同的計算節點（GPU）上。由於這一系列的專欄主要討論數據並行，在這裡我們就不展開討論模型並行了。而另一個可行的方法就是GeePS這篇文章裡面提出的方法。

接下來我要介紹的我自己的兩個工作，Poseidon和GeePS：

Poseidon: An Efficient Communication Architecture for Distributed Deep Learning on GPU Clusters. Hao Zhang, Zeyu Zheng, Shizhen Xu, Wei Dai, Qirong Ho, Xiaodan Liang, Zhiting Hu, Jinliang Wei, Pengtao Xie, Eric P. Xing. ATC 2017.
GeePS: Scalable Deep Learning on Distributed GPUs with a GPU-Specialized Parameter Server. Henggang Cui, Hao Zhang, Gregory R. Ganger, Phillip B. Gibbons, Eric P. Xing. Eurosys 2016.

主要就是為了解決上面討論的這三個問題的。

Poseidon，海神，深度學習

Poseidon這個系統提出了幾個方案來解決上面提到的前兩個問題。

我們在前面的討論中明確了一點：網路通信無論如何都是有時間開銷的，Ethernet的帶寬高低或者latency大小只會影響這個時間的長短，但並不能把這個時間降到零。以這一點為前提，設想，如果按照前面講的 ${large texttt{pull} rightarrow Delta_{mathcal{L}} rightarrow texttt{push}}$ 三段式流程串列的進行通信和計算，無論這個通信是快是慢，這個時間開銷都會導致在分散式環境下每個iteration的時間比單機版要長，進而導致整個系統無論如何都無法達到線性加速（為什麼？）。所以，把通信和計算重疊（overlap）起來以便「掩蓋」通信時間幾乎是一個必須的步驟。

回到神經網路的訓練過程上，怎麼設計系統來重疊計算和通信？大家回顧一下後向傳播的細節，注意兩點：首先，神經網路的計算是一層接著一層完成的，不管是前向還是後向傳播，算完本層才能算下一層；另一方面，在後向傳播的過程中，一旦後一層拿到前一層的輸入，這一層的計算就不再依賴於前一層了。Poseidon利用了這種層與層之間的依賴性設計了一套pipeline，稱為Wait-free Backpropagation（WFBP），用多線程的方法把計算和通信盡量重疊了起來。具體來說，WFBP試圖並行執行一部分相互獨立的通信和計算，以隱藏參數同步的時間開銷：在後向傳播的時候，當第 ${large i}$ 層的計算完成後，第 ${large i}$ 層參數的同步（push和pull）和它前面所有 ${large i-1}$ 層的梯度計算是獨立的，也就是說二者可以並行；同時，由於參數的更新也是相互獨立的，沒有必要在所有參數都更新完成後再統一pull，取而代之，我們可以在某一層參數更新完成後立刻pull回本地，把上行通信和下行通信也並行起來。這麼做其實還有一個額外的好處：很多常見的神經網路都是前面幾層計算量大，而後面幾層的參數多，比如很多常用的CNN中前面的卷積層計算比較費時，而後面的全鏈接層則參數量大，這樣WFBP正好可以把整個計算流程中計算最耗時和通信最耗時的兩大塊重疊起來，大大隱藏了開銷。

WFBP示意圖：圖中灰色代表計算，綠色代表向PS發送梯度更新（push），淺藍色代表向PS請求最新參數（pull）

我們在普通的Parameter Server架構上加入WFBP後進行了一些實驗，結果證明WFBP的優化效果是十分顯著的。對於Inception-v3和ResNet-152等參數分布比較均勻，且沒有巨大的全連接層的網路，只是使用WFBP就能在32個Titan X節點上達到接近30倍的加速。對WFBP Pipeline具體實現感興趣的同學，可以參考Poseidon論文。

WFBP的主要思想是通過多線程、多stream技術把通信和計算盡量重疊起來，但是僅有這一步仍是遠遠不夠的。我們又在一個稍微大一點、通訊計算時間比更大的網路VGG19上做了一組實驗：僅僅依賴WFBP，在32個節點我們大概最多能達到20x的加速比，雖然相比較於TensorFlow原版的負加速比（是的，TensorFlow自帶的分散式訓練在某些神經網路上可能是負加速，分散式訓練比單機還要慢）已經很不錯了，但是這個結果並不令人滿意：如果你花了10萬刀買了機器，大概有3萬多刀打了水漂。

另外值得強調的是，我們上面的幾組實驗用了40Gbps的Ethernet，這個已經算是奢侈品了。考慮到大部分的雲計算平台（AWS等），以及大多數的機器學習實驗室其實很難有這種硬體配置（畢竟不是人人都是Google或Facebook，感興趣的同學可以去看看Facebook的論文裡面用了多少帶寬的Ethernet），這個對硬體配置的要求顯然並不現實。實際上，我們又把上面測加速比的實驗重新搬到10Gbps和1Gbps的Ethernet上重新跑了幾次，看到結果後心態直接崩了：由於可用帶寬瞬間減小了很多，上述的好幾個實驗在32台機器上大概只剩個位數的加速比了（具體參見Poseidon論文）。我們用的GPU還是2015年的Titan X，而且這幾個網路裡面最大的VGG也僅有120M個參數而已。隨著GPU的飛速發展，以及各種大模型被開發出來，這個問題會變得原來越嚴重。

問題出在哪？與單機上高速的PCIe匯流排相比，多機間的網路傳輸速度是很慢的，而深度神經網路的一大特點就是參數多，通信量大。設想，如果總通信時間超過了計算時間，就算你能夠設計出完美的重疊通信和計算的方法，也無法掩蓋通信時間 -- 你永遠無法達到線性加速，因為在這種環境下，每個iteration所需的時間是由通信時間決定的，而並非計算時間了。其次，突發數據傳輸（burst）也是一大問題，如果網路通信勻速發生，那麼只要可用的帶寬大於所需要的通信速度，就能得到比較好的傳輸性能；如果通訊負載不均勻，在大部分時間網路帶寬處於空閑狀態，但是在某幾個時間點突發很多信息需要傳輸，這就會導致在這一時刻對Ethernet帶寬的要求劇增，並超過上限帶寬。結果就是，在空閑時刻，Ethernet的帶寬利用不充分；在忙碌的時候又負載失衡，在峰值傳輸點達到時網路通信會阻塞。

回到訓練VGG19這個具體問題上。通過一些分析和實驗，我們發現主要問題在於VGG19有兩個巨大的全連接層（FC），同步這兩層參數消耗了大量的時間，使得通信時間超過了計算時間，拖慢了速度。為了解決這個問題，Poseidon採取了一套叫做混合通信（hybrid communication）的機制，試圖直接減小通信量（但不損失精度）。如果通信總量減少了，通信時間自然就變短了。什麼是hybrid communication？大體思想是，對神經網路不同類型的層，使用不同類型的通信模式。舉個例子，在Poseidon裡面，對於一個卷積神經網路裡面的卷積層，Poseidon直接使用PS進行參數同步，因為卷積層有參數共享，參數個數一般都很少；而針對全連接層這種參數特別多的層，Poseidon會根據情況，在必要時選擇另一個叫做Sufficient Factor Broadcasting（SFB）的模式來進行參數同步。

SFB其實是個輕量級的分散式機器學習通信框架，想深究這一塊的同學可以讀讀這篇文章。這套通信模式可以用到很多具備一定性質的機器學習模型上，而很多神經網路正具備這個性質。原理上，我們發現全連接層的參數實際上是個大矩陣，而且在後向傳播時候產生的的梯度矩陣其實是兩個向量的乘積（秩為1）：具體來說，一個參數矩陣大小為 ${N times M}$ 的FC層，在算梯度的過程中，假設batchsize是 $B$ ，那麼它在這個batch上的梯度可以分解為 $B$ 個長度為 $N$ 和 $M$ 的向量的乘積，這些向量被稱為Sufficient Factor（SF）。在實際中， $B$ 往往是遠小於 $N$ 和 $M$ 的，因此我們可以選擇發送SF來減小通信量，而不用再發送原來那個很大的梯度矩陣，在通信結束後再用SF重構原始梯度矩陣即可。

Microsoft發表在OSDI 2014上的Project Adam系統也使用了一個類似的技術進行優化，不過Project Adam只優化了上行通信（push），下行拉取最新參數（pull）時，仍需發送整個大參數矩陣。這一方面沒有解決下行通信量大的問題；另一方面，這種上行發SF下行發參數矩陣的設計導致只能有一個物理節點作為伺服器來接受所有計算節點的SF，再由這個節點把更新過後的參數矩陣返還給所有其他節點。如果你理解了我前面講的PS的設計，你大概會發現這會導致一個嚴重的問題：我們在實現PS的時候通常使用分散式存儲，以便盡量分散利用每個節點的網路傳輸帶寬；但是在這種設計下，某一個物理節點需要承擔很大的通信任務，這個節點很有可能會成為一個通信瓶頸。不同的是，在SFB裡面，每個節點發送自己產生的SF給其他所有節點（broadcast），每個節點收到其他節點發來的SF後在本地重構出梯度，再更新本地參數。每個節點承擔相同的通信量，大大減小了負載不均衡的可能性。但是SFB也有自身的問題：P2P通信的開銷是與集群規模的平方成正比的，當集群規模太大的時候，使用P2P的通信方式不僅不能使通信開銷減小，反而可能增大。

所以這裡Poseidon的hybrid communication就發揮其優勢了。Poseidon並不是在任何情況下對任何FC層矩陣都使用SFB進行通信，而是用一個自適應的方法：用一個不等式去判斷當前情況SFB和PS哪個更適合。不等式很簡單，簡單的推導就能得到，感興趣的同學可以自己動手算一下。有了SFB之後，Poseidon處理大網路的能力得到了增強。對於VGG19，在32個節點以及很小的可用網路帶寬上，也能輕鬆達到30x的加速；對於更加喪心病狂的VGG19-22k（把輸出層從1000改為了21841，ImageNet 22K extreme classification的一個使用場景）也得到了28.5x的加速。相比較一些其它減少通信的方法，Poseidon也具備一定的優勢，具體的可以參加Poseidon論文中的細節。

本篇總結和下篇預告

談完了這些技術細節，我們可以再談談在設計Poseidon這個系統的一些思考。與Caffe, TensorFlow，Torch這種深度學習工具包不同，我們希望Poseidon更多的成為一個平台，而不是框架，或是工具包。在這個平台上，你可以把任何語言寫的深度學習程序 -- 不管是Python寫的TensorFlow或PyTorch程序，還是按照Caffe Prototxt定義的神經網路，抑或是Lua寫torch程序 -- 都輕鬆的部署到Poseidon上面，甚至不需要改一行代碼，就可以在一個GPU集群上跑起來，並且或得滿意的加速。所以Poseidon盡量保存了原有框架的編程語言介面，所以當你把一個TensorFlow程序部署到Poseidon上的時候，幾乎不用改任何代碼。另一方面，我們把Poseidon的底層代碼寫的非常輕量級，用Poseidon把一個現有的非常複雜的深度學習框架（比如TensorFlow）改成分散式版本大概只需要多寫一百行代碼就可以了。在另一方面，我們在部署上花了很大的功夫，讓Poseidon非常易於安裝。當然，這方面會更多的涉及到關於深度學習的產品設計，就不再深入討論了。

這篇文章想要講的內容基本已經涵蓋到了。那麼，Poseidon系統的名字是什麼意思呢？Poseidon，海神，深度學習......這個就留給各位讀者自己聯想了......

我會在第三篇文章，也就是這一系列專欄的最後一篇中談談GeePS這個工作，順便聊聊當下分散式深度學習領域其他比較熱點的話題，比如非同步訓練、動態神經網路、增強學習，等等。

當然，文章的最後不忘打個廣告：Petuum Inc.上個月完成了9300萬美元的B輪融資，歡迎對機器學習和分散式系統感興趣的各位投簡歷。Petuum新一輪融資過後會在灣區建立office，在灣區和匹茲堡都有headcount，歡迎各位諮詢！

* 這篇文章的封面圖來自：https://www.facebook.com/convolutionalmemes/.