雷鋒網(wǎng) AI科技評(píng)論按：本文作者王佐，文章首發(fā)于其知乎主頁(yè)，AI科技評(píng)論獲其授權(quán)轉(zhuǎn)載。

在上一家公司就開(kāi)始實(shí)踐打磨一個(gè)深度優(yōu)化的深度學(xué)習(xí)系統(tǒng)，當(dāng)時(shí)從消除網(wǎng)絡(luò)瓶頸，非凸優(yōu)化，以及具體的深度學(xué)習(xí)算法等方面基于PaddlePaddle做了許多工作。目前公司主要深度學(xué)習(xí)算法都是跑在Tensorflow上，使用配置了GeForce GTX 1080的單機(jī)訓(xùn)練，一次完整的訓(xùn)練至少需要一周的時(shí)間，所以決定從優(yōu)化Tensorflow多機(jī)并行方面提高算力。

• 為什么要優(yōu)化Tensorflow多機(jī)并行

更多的數(shù)據(jù)可以提高預(yù)測(cè)性能[2]，這也意味著更沉重的計(jì)算負(fù)擔(dān)，未來(lái)算力將成為AI發(fā)展的最大瓶頸。在大數(shù)據(jù)時(shí)代，解決存儲(chǔ)和算力的方法是Scale out，在AI時(shí)代，Scale out也一定是發(fā)展趨勢(shì)，并且大數(shù)據(jù)分析任務(wù)和AI/ML任務(wù)會(huì)共享處理設(shè)備（由于AI/ML迭代收斂和容錯(cuò)的特征，這兩種任務(wù)未來(lái)不太可能使用統(tǒng)一平臺(tái)），所以需要在分布式環(huán)境下優(yōu)化資源配置[3]，消除性能瓶頸。雖然現(xiàn)在Tensorflow能支持多機(jī)并行分布式訓(xùn)練，但是針對(duì)復(fù)雜網(wǎng)絡(luò)，其訓(xùn)練速度反而不如單臺(tái)機(jī)器[1]。目前已經(jīng)有IBM[4]和Petuum[1]分別在其深度學(xué)習(xí)系統(tǒng)PowerAI 4.0和Poseidon中實(shí)現(xiàn)多機(jī)并行線性加速，本文介紹我如何通過(guò)消除Tensorflow的網(wǎng)絡(luò)瓶頸，實(shí)現(xiàn)Tensorflow多機(jī)并行線性加速。

• Tensorflow分布式訓(xùn)練的網(wǎng)絡(luò)瓶頸分析

深度學(xué)習(xí)訓(xùn)練需要海量的數(shù)據(jù)，這就需要超大規(guī)模參數(shù)的網(wǎng)絡(luò)模型擬合。如果訓(xùn)練數(shù)據(jù)不足，會(huì)造成欠擬合；如果網(wǎng)絡(luò)模型參數(shù)太少，只會(huì)得到低精度的模型。目前常見(jiàn)網(wǎng)絡(luò)模型參數(shù)已經(jīng)上億，參數(shù)大小達(dá)到數(shù)GB。[10]中給出了訓(xùn)練數(shù)據(jù)和參數(shù)大小一些例子。

訓(xùn)練數(shù)據(jù)和參數(shù)大小（來(lái)自[10]）

目前GPU已經(jīng)成為深度學(xué)習(xí)訓(xùn)練的標(biāo)配。GPU具有數(shù)量眾多計(jì)算單元和超長(zhǎng)流水線，并且具備強(qiáng)大并行計(jì)算能力與浮點(diǎn)計(jì)算能力，可以大幅加速深度學(xué)習(xí)模型的訓(xùn)練速度，相比CPU能提供更快的處理速度、更少的服務(wù)器投入和更低的功耗。這也意味著，GPU集群上訓(xùn)練深度學(xué)習(xí)模型，迭代時(shí)間更短，參數(shù)同步更頻繁。[9]中對(duì)比了主流深度學(xué)習(xí)系統(tǒng)在CPU和GPU上的訓(xùn)練性能，可以看出GPU每次迭代的時(shí)間比CPU少2個(gè)數(shù)量級(jí)。

CPU訓(xùn)練alexnet（來(lái)自[9]）

GPU訓(xùn)練alexnet（來(lái)自[9]）

假設(shè)每0.5秒一個(gè)迭代，每個(gè)worker每秒需要通過(guò)網(wǎng)絡(luò)傳輸?shù)拇笥?GB，即使使用10GbE，參數(shù)同步也會(huì)瞬間把網(wǎng)絡(luò)占滿。考慮到訓(xùn)練數(shù)據(jù)可能通過(guò)NFS或者HDFS加載，也會(huì)占用很多網(wǎng)絡(luò)帶寬。在一個(gè)數(shù)據(jù)分析任務(wù)和AI/ML任務(wù)混合的環(huán)境中，大數(shù)據(jù)分析任務(wù)也會(huì)消耗很多網(wǎng)絡(luò)帶寬（如shuffle操作），網(wǎng)絡(luò)延遲會(huì)更加嚴(yán)重。所以如果想以Scale out的方式提升算力，網(wǎng)絡(luò)將是最大的瓶頸。[1]中通過(guò)實(shí)驗(yàn)證明，在8個(gè)節(jié)點(diǎn)進(jìn)行Tensorflow分布式訓(xùn)練，對(duì)于VGG19網(wǎng)絡(luò)，90%的時(shí)間花在等待網(wǎng)絡(luò)傳輸上面。

網(wǎng)絡(luò)開(kāi)銷（來(lái)自[2]）

• 消除網(wǎng)絡(luò)瓶頸的方法（一）

分布式深度學(xué)習(xí)可以采用BSP和SSP兩種模式。SSP通過(guò)允許faster worker使用staled參數(shù)，從而達(dá)到平衡計(jì)算和網(wǎng)絡(luò)通信開(kāi)銷時(shí)間的效果[8]。SSP每次迭代收斂變慢，但是每次迭代時(shí)間更短，在CPU集群上，SSP總體收斂速度比BSP更快，但是在GPU集群上訓(xùn)練，BSP總體收斂速度比SSP反而快很多[6]。

BSP模型有個(gè)缺點(diǎn)，就是每次迭代結(jié)束，Worker需要發(fā)送梯度更新到PS，每次迭代開(kāi)始，Worker需要從PS接收更新后的參數(shù)，這會(huì)造成瞬間大量的網(wǎng)絡(luò)傳輸。參數(shù)服務(wù)器通過(guò)把參數(shù)切分成block，并且shard到多臺(tái)機(jī)器，比較AllReduce，有效利用網(wǎng)絡(luò)帶寬，降低網(wǎng)絡(luò)延遲。目前主流的深度學(xué)習(xí)系統(tǒng)（Tensorflow，Mxnet，Petuum）都選擇用參數(shù)服務(wù)器做參數(shù)同步。

AllReduce（來(lái)自[5]）

Parameter Server

上圖可以很容易看出，AllReduce拓?fù)渲?，Reducer節(jié)點(diǎn)成為網(wǎng)絡(luò)傳輸?shù)钠款i。PS拓?fù)渲?，通常每臺(tái)機(jī)器啟動(dòng)相同數(shù)量的Worker和Parameter Server，每臺(tái)機(jī)器的網(wǎng)絡(luò)傳輸量基本相同。

ring AllReduce（來(lái)自[5]）

對(duì)于多機(jī)多卡訓(xùn)練，可以把參數(shù)先在本機(jī)聚合，再指定一個(gè)worker跟參數(shù)服務(wù)器交互，可以大量減少網(wǎng)絡(luò)傳輸。可以使用PaddlePaddle提出來(lái)的ring AllReduce，優(yōu)化單機(jī)多卡的本地聚合。

解決瞬間大量的網(wǎng)絡(luò)傳輸問(wèn)題另一個(gè)方法是實(shí)現(xiàn)GPU計(jì)算和網(wǎng)絡(luò)通信的Overlap。在反向傳播的backward階段產(chǎn)生梯度時(shí)，可異步地進(jìn)行梯度更新，并立即計(jì)算下一層網(wǎng)絡(luò)的梯度。梯度更新首先要把新梯度從GPU顯存拷貝到CPU內(nèi)存，這種GPU-CPU的拷貝也可以和GPU計(jì)算做overlap。因?yàn)镻S是跑在CPU上，所以GPU計(jì)算也跟PS參數(shù)更新實(shí)現(xiàn)Overlap。

GPU計(jì)算和網(wǎng)絡(luò)傳輸overlap（來(lái)自[1]）

• 消除網(wǎng)絡(luò)瓶頸的方法（二）

減少網(wǎng)絡(luò)傳輸量也是消除網(wǎng)絡(luò)瓶頸的有效途徑。網(wǎng)絡(luò)模型中90%參數(shù)集中在FC層。很多深度學(xué)習(xí)系統(tǒng)提出了減少FC層參數(shù)大小的方法，比如Adam中的Sufficient Factor，CNTK中的1-bit quantization，Petuum中的Sufficient Factor Broadcasting[7]。

在PS拓?fù)渲?，每個(gè)worker需要發(fā)送梯度  和接收參數(shù)  。SFB通過(guò) 將 轉(zhuǎn)化為兩個(gè)低維度矩陣 和 的傳輸，并采用P2P拓?fù)?，每個(gè)worker本地更新參數(shù)，避免了參數(shù) 的傳輸。SFB和PS比較如下：

PS和SFB（來(lái)自[1]）

（1）PS使用Master-Server架構(gòu)，而SFB使用P2P架構(gòu)，每個(gè)worker將 和 發(fā)送給所有其他worker，每個(gè)worker通過(guò) 和 在本地更新參數(shù) ，從而避免了PS中 的傳輸。

（2）PS每個(gè)worker的傳輸數(shù)據(jù)量是固定的，SFB每個(gè)worker的傳輸數(shù)據(jù)量跟總worker數(shù)有關(guān)，每個(gè)worker需要把 和 發(fā)送給其他worker（發(fā)送numWorkers - 1次）。

（3）SFB傳輸數(shù)據(jù)量還跟batch size有關(guān)。在非凸有限和問(wèn)題中

， 其中 ，

在SGD中， 表示一個(gè)樣本，在mini-batch SGD中， 表示batch size個(gè)樣本。

• 實(shí)現(xiàn)代碼

首先得實(shí)現(xiàn)PS和SFB，可以參照petuum，ps-lite，angel。

Tensorflow 相關(guān)的修改主要有兩個(gè)地方：

tensorflow/core/kernels/http://training_ops.cc中的ApplyXXXOp（ApplyGradientDescentOp，ApplyAdagradOp，ApplyMomentumOp等），將本地的梯度更新修改為 發(fā)送 ->PS端梯度更新->接收

tensorflow/core/kernels/http://matmul_op.cc中的MalMulOp::Compute，這里需要判斷是否使用PS或者SFB，從而將本地更新切換為PS更新或SFB更新。

本地更新

PS更新

SFB更新

目前我們已經(jīng)復(fù)現(xiàn)[1]中的實(shí)驗(yàn)結(jié)果，實(shí)現(xiàn)了Tensorflow多機(jī)并行的線性加速。

參考文獻(xiàn)：

[1]   Hao Zhang, Zeyu Zheng, Shizhen Xu, Wei Dai, Qirong Ho, Xiaodan Liang, Zhiting Hu, Jinliang Wei, Pengtao Xie, Eric P. Xing. Poseidon: An Efficient Communication Architecture for Distributed Deep Learning on GPU Clusters. ATC 2017.

[2]   C. Sun, A. Shrivastava, S. Singh, and A. Gupta. Revisitingunreasonable effectiveness of data in deep learning era. In arXiv:1707.02968, 2017.

[3]   Azalia Mirhoseini, Hieu Pham, Quoc V Le, Benoit Steiner, Rasmus Larsen, Yuefeng Zhou, Naveen Kumar, Mohammad Norouzi, Samy Bengio, and Jeff Dean. 2017. Device placement optimization with reinforcement learning. In International Conference on Machine Learning (ICML).

[4]   PowerAI DDL

[5] allreduce Bringing HPC Techniques to Deep Learning

[6]   H. Cui, H. Zhang, G. R. Ganger, P. B. Gibbons, and E. P. Xing. GeePS: Scalable deeplearning on distributed GPUs with a GPU-specialized parameter server. In Proceedings
of EuroSys, 2016.

[7]   XIE, P., KIM, J. K., ZHOU, Y., HO, Q., KUMAR, A., YU, Y., AND XING, E. Distributed Machine Learning via Sufficient Factor Broadcasting. In arXiv (2015).

[8]   HO, Q., CIPAR, J., CUI, H., KIM, J. K., LEE, S., GIBBONS, P. B., GIBSON, G. A., GANGER, G. R., AND XING, E. P. More Effective Distributed ML via a Stale Synchronous Parallel Parameter Server. In NIPS (2013).

[9]   Benchmarking State-of-the-Art Deep Learning Software Tools

[10]   NanoNets : How to use Deep Learning when you have Limited Data     

雷峰網(wǎng)版權(quán)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見(jiàn)轉(zhuǎn)載須知。