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