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