標籤：

快速打造分散式深度學習訓練平台

02-12

最近幾年，深度學習摧朽拉枯之勢席捲IT的各個角落，改變了各領域演算法研究和軟體開發的模式，也給IT基礎設施建設和平台工具研髮帶來了新的要求。今天，我們分享一下如何在開源技術的幫助下，快速搭建起一個分散式的深度學習訓練平台，加速深度神經網路的訓練，提高公司的競爭力。

接下來我們從四個角度來談一下這個深度學習訓練平台：

  • 1. 為什麼需要分散式訓練平台
  • 2. 搭建這個平台需要使用什麼關鍵技術和框架
  • 3. 當前的系統架構和實現
  • 4. 當前架構的不足和優化

一、為什麼需要分散式訓練平台

深度學習框架的多樣性和趨同性

當前的深度學習框架非常之多，耳熟能詳的就有Caffe/Caffe2、MxNet、TensorFlow、Torch、CNTK、Theano等等。這些框架開發語言不一，介面設計各異，這對於AI公司，尤其是中小型創業團隊的框架選擇、技術積累和快速研髮帶來很多困難。不過，雖然深度學習框架眾多，框架本身提供的功能、實現思路和核心技術有一定的趨同性。因此我們也看到有一些項目在這些框架之上做一層封裝從而實現邏輯和介面上的統一，比如Kera和DIGIST。不過想短時間內「一統江山」難度會比較大。

因此，我們可以考慮從另一個層面來「統一」多個深度學習的框架，即通過工程的手段開發一個訓練平台，將深度網路的訓練、測試過程抽象道統一的界面中。同時，配合數據管理（網路存儲和對象存儲），環境管理（容器），CI和自動部署等，形成一個深度學習的「循環」，這能極高的提升迭代的速度。

深度學習的迭代性和周期性

深度學習訓練的一個特點是具有很強的迭代性，即網路結構確定後，可以周期性地通過增加訓練數據而提高模型的泛化能力。這個迭代周期越短，模型更新的速度越快，就能獲得更好的效果和更快的產品更新。這個特點使得訓練過程非常適合自動化。

網路和演算法的複雜度快速增長，帶來計算資源的稀缺和浪費

2012年AlexNet開啟這一輪深度學習熱潮的時候，網路只有8層，但就算只有8層也需要大量的計算。更何況2016年153層的ResNet，最近甚至出現了上千層的網路結構，這讓計算資源變的非常稀缺。另一個方面，隨著訓練數據的快速增長，分散式訓練變得很迫切和必需，這進一步增加了計算資源的稀缺性。

目前大多數的訓練過程使用Nvidia GPU，而Nvidia GPU虛擬化的欠缺，使得GPU資源變的不容易管理。當數據中心有成百上千塊GPU的時候，很容易出現一個矛盾的現象：GPU計算資源稀缺，同時又有部分GPU資源被閑置浪費，整體的利用率比較低。

其實這個現象在雲計算和虛擬化出現之前就存在，只是那個時候浪費的是CPU、內存、IO和存儲。現在同樣的問題出現在GPU上而且顯得更嚴重和迫切。當然，無論是國內還是國外的公有雲、私有雲廠商都在積極的解決這個問題。在沒有成熟和穩定的方案出來之前，我們需要先自力更生，研發一個私有「GPU雲」來應對這個問題。

國內外的互聯網和IT公司都在非常積極的解決上訴問題，有更深層次和更廣泛的平台級解決方案，比如Google TensorFlow、FaceBook FBLearner、阿里雲PAI、騰訊DX-I等等。你可能已經發現，目前能提供這些平台服務的廠商也基本上是公有雲的提供商。其實這是自然而然的事情，在Iaas，Saas之後，Aaas(AI as a service)已經是業界的共識和趨勢。但是這些公有雲平台的現狀還在快速發展且不太完善，比較難以滿足我們實際的需要。例如，這種通用平台難以快速響應最新的、實驗性的深度學習技術；在細節實現和靈活性上難以滿足中小客戶的需求；另外一個很重要的原因是，目前公有雲上的GPU或HPC實例的價格比較高，從成本角度，中小型公司還沒有足夠的動力把深度學習的業務全部遷移到公有雲。

基於上面談到的幾點原因，我們決定研發一個深度學習的訓練平台，這個平台要具備的功能包括：

  • 管理多台訓練伺服器，尤其是帶有GPU的高性能計算伺服器，可把訓練任務分到到分散式的計算節點上執行計算
  • 集成多種訓練框架，抽象訓練過程，提供Web界面，上傳和指定相關數據和參數，即可啟動訓練任務並監控和分析訓練過程；
  • 池化計算資源，尤其是GPU資源，做成「GPU雲」。啟動訓練任務時，平台會自動把訓練任務分配到合適的GPU上；
  • 打通數據中心，可以直接把數據存儲平台中的數據導入到訓練節點上中；
  • 隔離計算節點中的資源和環境，兼容不同型號的GPU、不同版本的CUDA/CuDNN和不同的深度學習框架

二、關鍵技術和框架

Mesos+Marathon

Mesos是Apache下的開源分散式資源管理框架，它被稱為是分散式系統的內核，又被稱作是數據中心的操作系統。簡單的講，Mesos實現了一個資源管理的框架，它在數據中心層次上管理集群資源（CPU、GPU、RAM等），提供資源分配和任務調度的能力。為了進一步把資源和任務隔離，Mesos把具體的任務調度能力抽象出來交給具體的Framework來實現，比如Hadoop，Spark，MPI和Marathon等。

圖2-1

圖2-1是Mesos框架的架構圖。Mesos本身分為master節點和agent(slave)節點。多個master節點通過zookeeper進行選舉實現高可用。Agent節點部署在每一台伺服器實例中並與master相連，周期性地向master節點更新自身狀態。Hadoop或MPI等Framework通過實現mesos定義的scheduler來與mesos master節點通信，調度使用集群資源運行相應任務，任務在Agent上通過Framework executor來實際執行。最典型的executor是bash和docker，如果任務可以被設計成bash命令或者docker容器，那麼無需額外的Framework和Framework executor，僅通過mesos agent就可以運行任務。

作為比較老牌的分散式框架，Mesos支持的Framework非常之多，列表在這裡並且在不斷增加中。最近豆瓣就開源了TFMesos，提供對TensorFlow的支持。不過，最知名和通用的Framework之一就是Marathon。正如其名，Marathon被設計用來運行長時間的任務，比如Web服務，同時支持運行bash腳本和docker container。提供了很多任務調度的策略和約束以提高靈活性，以及任務組、動態埠、持久化、Pods（類似kubenetes pods）、健康檢查等高級特性。Marathon會監控任務狀態並在任務退出重新調度運行，同時還提供負載均衡等分散式特性，保證任務不間斷運行。

不過，Marathon本身並不太適合做機器學習或深度學習任務。這類任務的運行時間雖然很長，但是多數是one-off的，也就是訓練完一個模型後此次任務就已經完成；有些這類任務還需要在運行過程中根據模型收斂的情況，暫停任務，修改參數，重新運行或者修改訓練集，重新運行。甚至有些任務是一些實驗性質的任務，會被頻繁的啟動和停止。所以Marathon從設計上不太符合需要，我們需要的是一個支持深度學習的Framework。很可惜，這樣的通用Framework目前並不存在，且考慮到文章開始提到的一些現狀，很有可能難以實現。幸好，我們目前的需求還沒有那麼複雜，可通過一些trick和流程設計來實現，Marathon暫且可以拿來一用。

Docker

Docker本身無需多言，這裡從深度學習的角度談一下為什麼要使用Docker。

環境隔離

絕大多數的深度學習依賴Nvidia GPU和相應的運算庫CUDA、CuDNN等。Nvidia GPU的硬體型號和驅動繁煩複雜且存在一定的兼容性問題。CUDA等基礎庫眾多並更新頻繁。同時不同的深度學習框架對底層庫的依賴也有所不同。如果直接在裸機上運行深度學習任務，尤其是新手，會花費大量的時間在環境的安裝、配置和解決各種不兼容的問題，令人抓狂。我們把一系列的基礎庫、運算庫和深度學習庫打包成Docker鏡像，充分利用Docker的分層機制在不同層面進行共享和組合。Docker隔離了系統環境和執行環境，即隔離同一台伺服器上不同訓練任務的環境。我們可以把同一個任務分發到帶有不同型號GPU卡的伺服器上，也可以在同一台伺服器上同時運行不同CUDA版本、不同深度學習框架的多個任務。比較完美的解決了上述問題。

資源隔離

深度學習是典型的計算密集、IO密集和數據密集三者都滿足的任務。這裡的計算密集除了CPU密集外還增加了GPU密集。這對單機和集群的硬體資源管理和隔離提出了更高的要求。Mesos完成了對集群硬體資源的管理和有效調度，但是我們還需要進一步隔離硬體資源，在滿足任務要求的同時，還要避免多任務對資源的惡性和無序競爭。Docker就是解決這個問題的最優選擇。Docker在資源隔離方面的優異表現已經被業界充分驗證且成功運用在大大小小的雲裡面。

但是，Docker或者說Linux cgroup對GPU資源的管理沒有像CPU那麼完善和成熟。不過我們仍然可以通過設備映射的方式，在GPU卡這個粒度上對GPU資源進行隔離，雖然粒度比較粗，但是考慮到深度學習任務一般會把GPU資源吃滿，這樣的粗粒度尚可接受。其實，不使用Docker，Mesos本身就提供資源隔離的特性。不過術業有專攻，Mesos跟Docker還有一定的距離。

代碼共享

雖然目前有很多的深度學習框架和經典的神經網路，為了提高模型的性能和其他的工程因素，我們需要對不同的深度學習框架和計算庫作不同層面的修改，作為公司內部的私有版本來使用。這些代碼通常是滿足特定網路和用途因此在代碼管理、共享和使用上非常麻煩。通過Github + CI + Docker，我們把不同repo和branch的代碼打包成完成不同任務的Docker鏡像從而實現了更靈活和細粒度的共享。

Caffe/MxNet

Caffe和MxNet是在計算機視覺領域比較常用的兩個深度學習框架，也是我們第一步支持的深度學習框架。對二者的支持沒有本質的區別，主要是把框架的訓練過程抽象和統一到訓練平台層面，提供統一的介面啟動訓練任務。

三、系統架構和實現

圖3-1

針對上面提到的需求和關鍵技術，整個平台的架構如圖3-1所示。左邊是存儲所有數據的數據中心，使用分散式網路存儲，提供基於網路的對象存儲和文件存儲服務。右邊的Docker倉庫存放訓練所需要的深度學習框架鏡像，比如Caffe或者MxNet。中間包括兩部分，底層搭建了Mesos+Marathon的集群，上層是業務模塊，用於提供Web API、準備環境和數據，調用下層的集群API啟停和監控訓練任務等。

圖3-2

把中層模塊細化後的架構如圖3-2所示。集群中部署三個mesos master節點，這三個節點通過zookeeper集群完成選舉和高可用，同一時刻只有一個節點在工作，其他兩個節點處於standby狀態。同樣，marathon也按照這樣的結構來部署。Mesos-slave根據需要部署，有多少個計算節點就部署多少個mesos slave。所有的mesos slave都與mesos master相連。

圖3-3

上層的DgTrainer主要封裝了文件處理、集群操作、Docker相關操作、任務管理和日誌管理等。是對數據中心、Docker倉庫和Mesos集群的功能性封裝。在這些模塊上提供RESTFul的介面定義，實現新建任務、查看任務狀態、停止任務等介面。以新建Caffe訓練任務為例，見圖3-3，提供一系列的介面參數，輸入網路定義文件train.prototxt，訓練鏡像caffe:cuda8v5等。即可創新一個訓練任務。在任務被創建後，整個系統的數據流向如圖3-4所示。

  1. DgTrainer接受訓練參數
  2. 在數據中心準備數據
  3. Docker倉庫中準備鏡像
  4. 調用marathon介面，啟動訓練任務
  5. mesos調度任務到合適的slave節點上
  6. slave節點從數據中心讀取（預讀取）數據
  7. slave節點從Docker倉庫下載對應鏡像，資源映射並執行任務
  8. 記錄日誌和訓練結果，並反饋到DgTrainer

圖3-4

接下來討論幾個具體的技術實現問題。

運行複雜任務

我們通過調用marathon的介面來創建任務，marathon提供兩個任務啟動模式，命令行和docker。這對於比較簡單的任務足夠了，但是對於訓練任務而言就不太現實。比如，訓練任務在執行前需要先把GPU映射到容器中，而映射過程本身會比較繁瑣（原因下面細說）。因此，比較好的方式是我們實現一個比較全面的腳本（比如shell或者python），任務啟動時，先讓marathon把這個腳本下載到對應的節點上，通過運行這個腳本來執行真正的任務。這個腳本可根據實際的需要定義多個參數來靈活的定義任務的行為。

Marathon本身就提供了這樣的機制。在marathon的啟動任務介面中有個fetch參數，可以定義任務所依賴的 「資源」。Marathon會自動下載指定的資源，解壓縮並執行。

"fetch": [ { "uri": "https://foo.com/setup.py" }, { "uri": "https://foo.com/archive.zip", "executable": false, "extract": true, "cache": true, "destPath": "newname.zip" } ]

容器GPU映射

如果要在Docker內使用GPU硬體，需要做GPU的設備映射。這可以通過--device參數完成。其中nvidia0表示映射第一個GPU卡, nvidia3表示映射第四個GPU卡。

--device /dev/nvidiactl:/dev/nvidiactl /dev/nvidia-uvm:/dev/nvidia-uvm /dev/nvidia0:/dev/nvidia0 /dev/nvidia3:/dev/nvidia1

不過由於GPU的特殊性，在container內部還需要有對應的驅動才能調用GPU硬體資源。這裡有幾種解決辦法。最直接的辦法就是在鏡像內部直接安裝好驅動。這樣做的壞處是驅動版本被固定了，如果一台host的驅動比鏡像內的驅動版本老，極有可能出現無法使用的情況。

第二個辦法是使用Nvidia-docker。這是Nvidia公司針對GPU的特殊性在Docker上封裝了一層，使用方法和原生Docker幾乎一樣。Nvidia-docker會自動處理好硬體映射和驅動映射的問題。這樣的壞處就是必須使用Nvidia-docker來啟動容器，原生Docker的一些工具和API可能會受到限制。

第三個做法就是模擬Nvidia-docker的做法，把host上的驅動做成一個Docker volume並在啟動容器的時候，同時把驅動volume和設備映射到容器內。硬體映射採用開頭提到的方式。

GPU映射順序

在Mesos管理下，每一個任務所使用的GPU順序是不固定的。如果我們的訓練任務需要同時使用兩個GPU卡，對於一台四卡的伺服器，當前被使用的可能是第二和第三塊GPU卡。那麼再啟動一個任務的時候，需要使用的就是第一和第四塊卡，如開頭的示例。這裡我們需要在任務運行之前就確定好要被映射的GPU卡。由於Mesos本身就有資源隔離，因此在沙箱中看到就只有兩個GPU並且序號是0和1，而實際上，這兩個GPU在操作系統中的設備文件其實是/dev/nvidia1和/dev/nvidia2。

因此，我們需要查出兩個GPU卡的「真實身份」。這裡不能使用GPU卡的順序號而應該通過順序號確定MinorNumber，MinorNumber才是/dev中的真實設備文件號。MinorNumber可以調用nvidia nvml庫獲得，這個庫有多種語言的實現。以python為例：

device = nvmlDeviceGetHandleByIndex(i) minorNum = nvmlDeviceGetMinorNumber(device)

四、不足和優化

基於上面提到的技術方案，我們在幾周時間內就搭建起一個分散式的深度學習訓練平台。整體上而言，藉助成熟的開源方案Mesos，Marathon和Docker，這個平台隔離了數據、計算和環境，基本實現了一個分散式系統所需要具備的特性，比如計算資源池化、任務調度、故障遷移、環境隔離、高可用、數據集中等。當然，這個平台還有一些問題尚需解決。

分散式訓練

目前的深度學習框架大都支持分散式訓練。分散式訓練分為模型並行和數據並行。多數情況下並行採用後者，即把相同的模型分布在不同的計算幾點上，通過提高batch size從而加速訓練過程。因此，一個訓練平台需要支持分散式訓練的能力。在技術實現上，一是需要依賴具體的深度學習框架，結合不同的分散式技術，比如MPI、Spark、ParameterServer等；二是需要定義更複雜的資源調度策略，避免出現「木桶效應」；三是支持更複雜的調試、調優等能力。整體來講，分散式訓練+分散式平台是趨勢和必然，也有更多的技術挑戰。

定時任務

多數深度學習任務是周期性的批量任務，因此需要給一個任務增加定時執行的能力。Mesos chronos框架是解決這個需求的一個方案。

自動部署和測試

在模型訓練過程中和訓練完畢後，需要把模型做自動化部署並在測試集上進行測試，以此評測對比模型的效果。

數據中心調優

深度學習對數據訪問有極高的要求，因此可以在平台層面調優數據中心的性能，同時還需要調優數據中心和訓練平台的數據傳輸性能。目前有專門的硬體設備，比如InfiniteBand可以有效提高數據傳輸性能

推薦閱讀：

TAG:深度學習DeepLearning | 分散式系統 | 機器學習 |