近年來，深度學(xué)習(xí)引領(lǐng)了圖像處理、語音識(shí)別和預(yù)測(cè)等方面的巨大進(jìn)步。在 Uber，我們將深度學(xué)習(xí)應(yīng)用到了公司業(yè)務(wù)中，從自動(dòng)駕駛搜索路線到防御欺詐，深度學(xué)習(xí)讓我們的數(shù)據(jù)科學(xué)家和工程師們能夠?yàn)橛脩籼峁└玫捏w驗(yàn)。
TensorFlow 已經(jīng)成為了 Uber 首選的深度學(xué)習(xí)庫。因?yàn)檫@個(gè)框架是目前使用最為廣泛的開源深度學(xué)習(xí)框架，對(duì)于新的開發(fā)者而言非常友好。它結(jié)合了高性能與低級(jí)模型細(xì)節(jié)調(diào)試能力——例如，我們可以使用 Keras 這種高級(jí) API，同時(shí)使用自己定制的 Nvidia CUDA 工具。此外，TensorFlow 還為各種深度學(xué)習(xí)用例提供了端到端支持，從進(jìn)行實(shí)驗(yàn)性探索到將生產(chǎn)級(jí)模型部署到云服務(wù)器、移動(dòng)端 APP、甚至自動(dòng)駕駛汽車上。
上個(gè)月 Uber 工程部門推出了 Michelangelo——一個(gè)內(nèi)部機(jī)器學(xué)習(xí)服務(wù)平臺(tái)，可以讓機(jī)器學(xué)習(xí)輕松部署到大規(guī)模系統(tǒng)中。在本文中 Uber 介紹了 Michelangelo 深度學(xué)習(xí)工具包的重要開源組件 Horovod，它可以讓分布式 TensorFlow 深度學(xué)習(xí)項(xiàng)目更加輕松地實(shí)現(xiàn)。
面向分布式
隨著 Uber 在 TensorFlow 上訓(xùn)練越來越多的機(jī)器學(xué)習(xí)模型，項(xiàng)目的數(shù)據(jù)和計(jì)算能力需求正在急劇增加。在大部分情況下，模型是可以在單個(gè)或多 GPU 平臺(tái)的服務(wù)器上運(yùn)行的，但隨著數(shù)據(jù)集的增大和訓(xùn)練時(shí)間的增長，有些時(shí)候訓(xùn)練需要一周甚至更長時(shí)間。因此，Uber 的工程師們不得不尋求分布式訓(xùn)練的方法。
Uber 開始嘗試部署標(biāo)準(zhǔn)分布式 TensorFlow 技術(shù)，在試驗(yàn)了一些方法之后，開發(fā)者意識(shí)到原有方法需要進(jìn)行一些調(diào)整：首先，在遵循文檔和代碼示例之后，我們并不總是清楚哪些功能對(duì)應(yīng)著哪些模型訓(xùn)練代碼的分布式計(jì)算。標(biāo)準(zhǔn)分布式 TensorFlow 引入了很多新的概念：工作線程、參數(shù)服務(wù)器、tf.Server（）、tf.ClusterSpec（）、 tf.train.SyncReplicasOptimizer（） 以及 tf.train.replicas_device_setter（） 等等。它們?cè)谀承┣闆r下能起到優(yōu)化作用，但也讓我們難以診斷拖慢訓(xùn)練速度的 bug。
第二個(gè)問題有關(guān) Uber 規(guī)模的計(jì)算性能。在進(jìn)行了一些基準(zhǔn)測(cè)試之后，我們發(fā)現(xiàn)標(biāo)準(zhǔn)的分布式 TensorFlow 機(jī)制無法滿足需求。例如，在使用 128 個(gè) GPU 進(jìn)行訓(xùn)練時(shí)，我們因?yàn)榈托蕮p失了一半的計(jì)算資源。
圖 1. 標(biāo)準(zhǔn) TensorFlow 基準(zhǔn)套件，使用英偉達(dá) Pascal GPU（從 1 塊到 128 塊）運(yùn)行 Inception V3 和 ResNet-101 模型，與理想狀態(tài)下的分布式計(jì)算（單 GPU 算力簡(jiǎn)單疊加）每秒處理的圖像數(shù)量對(duì)比。從中我們發(fā)現(xiàn)標(biāo)準(zhǔn)方法很難釋放出硬件的全部潛能。
當(dāng)我們使用標(biāo)準(zhǔn) TensorFlow 基準(zhǔn)測(cè)試套件在 128 塊英偉達(dá) Pascal GPU 上進(jìn)行測(cè)試時(shí)（如圖 1 所示），無論是 Inception V3 還是 ResNet-101 都浪費(fèi)了將近一半 GPU 算力。
充分利用 GPU 資源是目前大規(guī)模訓(xùn)練的一大課題，此前 Facebook 的一小時(shí)訓(xùn)練 ImageNet 論文《Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour》介紹了使用 256 塊 GPU 進(jìn)行 ResNet-50 網(wǎng)絡(luò)「數(shù)據(jù)并行」訓(xùn)練的方法，引起人們的廣泛關(guān)注，這也證明了大規(guī)模分布式訓(xùn)練可以顯著提高生產(chǎn)力。
利用不同類型的算法
圖 2.「數(shù)據(jù)并行」方法在分布式訓(xùn)練上包含在多節(jié)點(diǎn)上并行分割數(shù)據(jù)和訓(xùn)練。在同步情況下，不同批次數(shù)據(jù)的梯度將在不同節(jié)點(diǎn)上分別進(jìn)行計(jì)算，但在節(jié)點(diǎn)之間進(jìn)行互相平均，以對(duì)每個(gè)節(jié)點(diǎn)中的模型副本應(yīng)用一致化更新。
在 Facebook 的研究之后，Uber 的研究人員開始尋找更好的分布式 TensorFlow 模型訓(xùn)練方法。由于我們的模型小到可以在單個(gè) GPU 或多 GPU 的單服務(wù)器上運(yùn)行，我們開始嘗試使用 Facebook 的數(shù)據(jù)并行方法。
在概念上，數(shù)據(jù)并行的分布式訓(xùn)練方法非常直接：
1. 運(yùn)行訓(xùn)練腳本的多個(gè)副本，每個(gè)副本：
a）讀取數(shù)據(jù)塊
b）將其輸入模型
c）計(jì)算模型更新（梯度）
2. 計(jì)算這些副本梯度的均值
3. 更新模型
4. 重復(fù) 1a 步驟
標(biāo)準(zhǔn)分布式 TensorFlow 包使用參數(shù)服務(wù)器的方法來平均梯度。在這種方法之下，每個(gè)進(jìn)程都有一到兩個(gè)角色：工作線程或參數(shù)服務(wù)器。工作線程處理訓(xùn)練數(shù)據(jù)，計(jì)算梯度，并把它們傳遞到參數(shù)服務(wù)器上進(jìn)行平均。
圖 3. 分布式訓(xùn)練中的參數(shù)服務(wù)器可以按照不同比例的參數(shù)服務(wù)器和工作線程進(jìn)行配置，每一個(gè)都有著不同的配置數(shù)據(jù)。
盡管這種方法可以提升表現(xiàn)，但我們?nèi)匀幻媾R兩大挑戰(zhàn)：
- 確定工作線程與參數(shù)服務(wù)器的正確比例：一旦使用參數(shù)服務(wù)器，它就可能變成網(wǎng)絡(luò)或計(jì)算的瓶頸。如果使用多個(gè)參數(shù)服務(wù)器，通信模式就會(huì)變成「all-to-all」的狀態(tài)，網(wǎng)絡(luò)可能會(huì)很快飽和。
- 應(yīng)對(duì)不斷增加的 TensorFlow 程序復(fù)雜性：在測(cè)試中我們發(fā)現(xiàn)，每個(gè)使用分布式 TensorFlow 的案例都需要指定初始工作線程和參數(shù)服務(wù)器，傳遞服務(wù)發(fā)現(xiàn)信息，如所有工作線程和參數(shù)服務(wù)器的主機(jī)和端口，并使用合適的 tf.ClusterSpec（） 構(gòu)建 tf.Server（），進(jìn)而調(diào)整訓(xùn)練程序。此外，用戶必須保證所有的操作都正確地使用 tf.train.device_replica_setter（），并使用 towers 讓代碼符合服務(wù)器中多 GPU 的設(shè)置。這通常導(dǎo)致陡峭的學(xué)習(xí)曲線和大量的代碼重構(gòu)，壓縮了實(shí)際建模的時(shí)間。
在 2017 年上半年，百度發(fā)表了研究《Bringing HPC Techniques to Deep Learning》（參見百度將 HPC 引入深度學(xué)習(xí)：高效實(shí)現(xiàn)模型的大規(guī)模擴(kuò)展），提出使用不同的算法來平均梯度，并讓這些梯度在所有節(jié)點(diǎn)之間交流，這被稱為 ring-allreduce，他們使用 TensorFlow 也實(shí)現(xiàn)了這種算法（https://github.com/baidu-research/tensorflow-allreduce）。該算法的思路基于 2009 年 Patarasuk 與 Xin Yuan 的論文《Bandwidth Optimal All-reduce Algorithms for Clusters of Workstations》。
圖 4. ring-allreduce 算法允許工作線程節(jié)點(diǎn)平均梯度并將其分散到所有節(jié)點(diǎn)——無需參數(shù)服務(wù)器。
在 ring-allreduce 算法中，每個(gè) N 節(jié)點(diǎn)與其他兩個(gè)節(jié)點(diǎn)進(jìn)行 2*（N-1） 次通信。在這個(gè)通信過程中，一個(gè)節(jié)點(diǎn)發(fā)送并接收數(shù)據(jù)緩沖區(qū)傳來的塊。在第一個(gè) N-1 迭代中，接收的值被添加到節(jié)點(diǎn)緩沖區(qū)中的值。在第二次 N-1 迭代中，接收的值代替節(jié)點(diǎn)緩沖區(qū)中保存的值。百度的文章證明了這種算法是帶寬上最優(yōu)的，這意味著如果緩沖區(qū)足夠大，它將最大化地利用可用的網(wǎng)絡(luò)。
除了網(wǎng)絡(luò)最優(yōu)化，allreduce 方法也易于理解和應(yīng)用。用戶可以利用消息傳遞接口（Message Passing Interface，MPI）實(shí)現(xiàn)，如 Open MPI，來啟動(dòng) TensorFlow 程序的所有副本。MPI 明確地建立了在分布式條件下工作線程互相通信的范式。用戶需要使用 allreduce（） 來調(diào)整自己的程序以平均梯度。
Horovod 簡(jiǎn)介
意識(shí)到 ring-allreduce 方法能夠改善易用性和性能，這激勵(lì)我們繼續(xù)研究適合我們的實(shí)現(xiàn)，以滿足 UberTensorFlow 的需求。我們采用了百度的 TensorFlow ring-allreduce 算法實(shí)現(xiàn)，并在此基礎(chǔ)上進(jìn)行構(gòu)建。流程如下：
1. 我們將代碼轉(zhuǎn)換成獨(dú)立的 Python 包 Horovod，它的名字來自于俄國傳統(tǒng)民間舞蹈，舞者手牽手圍成一個(gè)圈跳舞，與分布式 TensorFlow 流程使用 Horovod 互相通信的場(chǎng)景很像。Uber 的不同團(tuán)隊(duì)可能使用不同版本的 TensorFlow。我們希望所有團(tuán)隊(duì)無須更新到 TensorFlow 最新版，就可以利用 ring-allreduce 算法，使用補(bǔ)丁，甚至構(gòu)建框架。擁有獨(dú)立的 Python 包使安裝 Horovod 的時(shí)間從一個(gè)小時(shí)縮減至幾分鐘，時(shí)間長短取決于硬件條件。
2. 我們用 NCCL 替換百度的 ring-allreduce 實(shí)現(xiàn)。NCCL 是英偉達(dá)的集合通信庫，提供高度優(yōu)化的 ring-allreduce 版本。NCCL 2 允許在多個(gè)機(jī)器之間運(yùn)行 ring-allreduc，這使得我們利用其多種性能提升優(yōu)化。
3. 我們支持模型適應(yīng)單個(gè)服務(wù)器和多個(gè) GPU，原始版本只支持單個(gè) GPU 模型。
4. 最后，我們根據(jù)大量初始用戶的反饋對(duì) API 進(jìn)行了多處改進(jìn)。特別是，我們實(shí)現(xiàn)了廣播操作，使模型在所有工作線程中實(shí)現(xiàn)一致性初始化。新的 API 允許我們將用戶在單個(gè) GPU 項(xiàng)目中的運(yùn)算量減少到 4。
接下來，我們將討論如何在團(tuán)隊(duì)中使用 Horovod 進(jìn)行機(jī)器學(xué)習(xí)。
使用 Horovod 分配訓(xùn)練任務(wù)
分布式 TensorFlow 的參數(shù)服務(wù)器模型（parameter server paradigm）通常需要對(duì)大量樣板代碼進(jìn)行認(rèn)真的實(shí)現(xiàn)。但是 Horovod 僅需要幾行。下面是一個(gè)分布式 TensorFlow 項(xiàng)目使用 Horovod 的示例：
在該示例中，粗體文字指進(jìn)行單個(gè) GPU 分布式項(xiàng)目時(shí)必須做的改變：
hvd.init（） 初始化 Horovod。
config.gpu_options.visible_device_list = str（hvd.local_rank（）） 向每個(gè) TensorFlow 流程分配一個(gè) GPU。
opt=hvd.DistributedOptimizer（opt） 使用 Horovod 優(yōu)化器包裹每一個(gè)常規(guī) TensorFlow 優(yōu)化器，Horovod 優(yōu)化器使用 ring-allreduce 平均梯度。
hvd.BroadcastGlobalVariablesHook（0） 將變量從第一個(gè)流程向其他流程傳播，以實(shí)現(xiàn)一致性初始化。如果該項(xiàng)目無法使用 MonitoredTrainingSession，則用戶可以運(yùn)行 hvd.broadcast_global_variables（0）。
之后，用戶可以使用 mpirun 命令使該項(xiàng)目的多個(gè)拷貝在多個(gè)服務(wù)器中運(yùn)行：
mpirun 命令向四個(gè)節(jié)點(diǎn)分布 train.py，然后在每個(gè)節(jié)點(diǎn)的四個(gè) GPU 上運(yùn)行 train.py。
Horovod 還通過同樣的步驟分布 Keras 項(xiàng)目。（TensorFlow 和 Keras 的腳本示例地址：https://github.com/uber/horovod/blob/master/examples/）
Horovod 的易用性、調(diào)試效率和速度使之成為對(duì)單 GPU 或單服務(wù)器項(xiàng)目感興趣的工程師和數(shù)據(jù)科學(xué)家的好搭檔。下面，我們將介紹 Horovod Timeline，它在分布式訓(xùn)練工作中提供對(duì)工作線程節(jié)點(diǎn)狀態(tài)的高度理解。
Horovod Timeline
我們?cè)谠试S用戶使用 Horovod 時(shí)，就意識(shí)到需要向用戶提供一種能夠輕松識(shí)別代碼中 bug 的方式，這也是處理復(fù)雜分布式系統(tǒng)時(shí)常常面臨的問題。尤其是，由于用戶需要收集和交叉引用不同服務(wù)器上的文件，用戶很難使用原始的 TensorFlow timeline 或 CUDA 分析器。
我們希望用 Horovod 創(chuàng)造一種方式，提供節(jié)點(diǎn)之間操作 timeline 的高度理解。因此，我們構(gòu)建了 Horovod Timeline。用戶可以使用 Horovod Timeline 清晰看到每個(gè)節(jié)點(diǎn)在訓(xùn)練過程的每個(gè)時(shí)間步的狀態(tài)。這有助于識(shí)別 bug，解決性能問題。用戶可通過設(shè)置單個(gè)環(huán)境變量啟用 timeline，通過 chrome://tracing 在瀏覽器中查看分析結(jié)果。
圖 5：Horovod Timeline 在 Chrome 的事件追蹤性能分析工具（trace event profiling tool）中描述分布式訓(xùn)練過程中的高級(jí)別 timeline。
Tensor Fusion
我們分析了多個(gè)模型的 timeline 之后，發(fā)現(xiàn)具有大量張量的模型，如 ResNet-101，有很多小的 allreduce 操作。之前我們注意到，ring-allreduce 在張量足夠多的情況下可以最大化利用網(wǎng)絡(luò)，但工作效率和速度都不如張量少的情況。于是問題來了：如果在張量上執(zhí)行 ring-allreduce 之前，先融合多個(gè)小張量，會(huì)發(fā)生什么呢？
答案就是：Tensor Fusion，一種在執(zhí)行 Horovod 的 ring-allreduce 之前先融合張量的算法。我們使用該方法進(jìn)行實(shí)驗(yàn)，發(fā)現(xiàn)在未優(yōu)化的傳輸控制協(xié)議（TCP）網(wǎng)絡(luò)上運(yùn)行的多層模型性能提升了 65%。我們簡(jiǎn)要介紹了 Tensor Fusion 的使用方法：
- 確定要減少哪些向量。首先選擇幾個(gè)在緩沖區(qū)（buffer）中適用且具備同樣的數(shù)據(jù)類型的張量。
- 為未分配的張量分配融合緩沖區(qū)（fusion buffer）。默認(rèn)的融合緩沖區(qū)大小是 64 MB。
- 將所選張量的數(shù)據(jù)復(fù)制到融合緩沖區(qū)。
- 在融合緩沖區(qū)上執(zhí)行 allreduce 操作。
- 將融緩沖區(qū)中的數(shù)據(jù)復(fù)制到輸出張量中。
- 重復(fù)直到該循環(huán)中沒有需要減少的張量。
我們使用 Horovod、Tensor Fusion 和在 Michelangelo 平臺(tái)上構(gòu)建的其他特征，提高模型在我們的機(jī)器學(xué)習(xí)系統(tǒng)中的效率、速度和易用性。下一部分，我們將分享現(xiàn)實(shí)世界的基準(zhǔn)，來展示 Horovod 的性能。
Horovod 基準(zhǔn)
圖 6：Inception V3 和 ResNet-101 TensorFlow 模型在 25GbE TCP 上使用不同數(shù)量的 NVIDIA Pascal GPU 時(shí)，使用標(biāo)準(zhǔn)分布式 TensorFlow 和 Horovod 運(yùn)行分布式訓(xùn)練工作每秒處理的圖像數(shù)量對(duì)比。
我們重新運(yùn)行調(diào)整后適合 Horovod 的官方 TensorFlow 基準(zhǔn)，并與常規(guī)的分布式 TensorFlow 的性能進(jìn)行對(duì)比。如圖 6 所示，Horovod 的能力有大幅改進(jìn)，我們不再浪費(fèi)一半的 GPU 資源。事實(shí)上，使用 Inception V3 和 ResNet-101 模型進(jìn)行縮放可以達(dá)到 88% 的計(jì)算效率。也就是說，訓(xùn)練速度是標(biāo)準(zhǔn)分布式 TensorFlow 的兩倍。
圖 7：Horovod 在 25GbE TCP 和 25GbE RDMA 網(wǎng)絡(luò)上每秒處理的圖像對(duì)比。它們?cè)诓煌瑪?shù)量的 NVIDIA Pascal GPU 上為 Inception V3、ResNet-101 和 VGG-16 運(yùn)行分布式訓(xùn)練工作。
由于 MPI 和 NCCL 都支持遠(yuǎn)程直接內(nèi)存訪問（RDMA）網(wǎng)絡(luò)，我們使用 RDMA 網(wǎng)卡運(yùn)行額外的基準(zhǔn)測(cè)試，來確定它們提升的效率是否能夠超過 TCP 網(wǎng)絡(luò)。
我們發(fā)現(xiàn) RDMA 沒有明顯提升 Inception V3 和 ResNet-101 模型上的性能，僅比 TCP 網(wǎng)絡(luò)提高了三四個(gè)百分點(diǎn)。但是，RDMA 幫助 Horovod 在兩個(gè)模型上實(shí)現(xiàn)了超過 90% 的縮放效率（scaling efficiency）。
與此同時(shí)，VGG-16 模型在使用 RDMA 網(wǎng)絡(luò)時(shí)速度提升了 30%。這可以用 VGG-16 的大量模型參數(shù)來解釋，全連接層和少量層的結(jié)合引起大量模型參數(shù)。這些特征改變了從 GPU 計(jì)算到通信的關(guān)鍵路徑，造成了網(wǎng)絡(luò)瓶頸。
這些基準(zhǔn)說明 Horovod 在 TCP 和 RDMA 網(wǎng)絡(luò)上的縮放效果很好，盡管使用 RDMA 網(wǎng)絡(luò)的用戶能夠在使用大量模型參數(shù)的模型如 VGG-16 時(shí)才能獲取最優(yōu)性能和顯著效率提升。
我們使用 Horovod 探索深度學(xué)習(xí)中的性能優(yōu)化還只是開始。未來，我們將持續(xù)利用開源社區(qū)使用我們的機(jī)器學(xué)習(xí)系統(tǒng)和框架實(shí)現(xiàn)性能提升。
下一步
今年早些時(shí)候，Uber 開源了 Horovod，讓這一可擴(kuò)展機(jī)器學(xué)習(xí)模型走向整個(gè)社區(qū)。目前 Horovod 還在發(fā)展之中，我們正在向以下幾個(gè)方向繼續(xù)推進(jìn)：
1. 讓 MPI 更易安裝：雖然在工作站上安裝 MPI 比較容易，但是在集群上安裝 MPI 仍然需要一些努力；例如，有很多工作負(fù)載管理器，我們需要根據(jù)不同的硬件進(jìn)行相應(yīng)的調(diào)整。我們正在開發(fā)為集群運(yùn)行 Horovod 的參考設(shè)計(jì)，為此，我們希望與 MPI 社區(qū)和網(wǎng)絡(luò)硬件供應(yīng)商合作，開發(fā)安裝 MPI 和相關(guān)驅(qū)動(dòng)程序的說明。
2. 收集和分享調(diào)整分布式深度學(xué)習(xí)模型參數(shù)的心得：Facebook 的「一小時(shí)訓(xùn)練 ImageNet 論文」描述了與在單 GPU 上訓(xùn)練模型相比，分布式訓(xùn)練任務(wù)需要超參數(shù)調(diào)整以達(dá)到甚至超越前者的準(zhǔn)確性。Facebook 證明了在 256 塊 GPU 上訓(xùn)練 TensorFlow 模型的可行性。
3. 加入超大模型示例：Horovod 目前支持適用于單 GPU，同時(shí)也支持多 GPU 服務(wù)器的模型。我們希望在更多形式的硬件上應(yīng)用更大的模型。
我們希望 Horovod 的簡(jiǎn)潔性可以使大家采用分布式訓(xùn)練，更好地利用計(jì)算資源用于深度學(xué)習(xí)。
