ElasticDL: Kubernetes-native 彈性分佈式深度學習系統

9月11日，螞蟻金服在 Google Developer Day Shanghai 2019 上宣佈開源了基於 TensorFlow 2.0 eager execution 的分佈式深度學習系統 ElasticDL。基於 TensorFlow 的支持彈性調度的深度學習系統，據咱們所知，ElasticDL 是第一 個。項目負責人王益和咱們分享了 ElasticDL 項目的設計意圖和現狀，尤爲是 ElasticDL 與 TensorFlow 2.0 以及 Kubernetes 的技術關聯。

分佈式深度學習的技術思路

基於 TensorFlow 的分佈式訓練系統大體能夠分爲如下四類：

 

其中，ElasticDL 位於田字格的右上角。之因此選擇這條技術思路，是爲了利用 Kubernetes 實現容錯和彈性調度。

高性能計算和雲計算

在深度學習技術研發的早期，涉及的人員相對少，共用一個計算集羣的人相對少， 計算做業之間的協調能夠經過口頭交流實現。你們更關心縮短運行時間，也就是 從做業啓動到結束的這段時間。高性能計算技術（HPC）是解決這個問題的有效 途徑，好比 NVIDIA 的 cuBLAS 和 cuDNN 優化高性能數學計算、NCCL 優化 GPU 之間的通訊效率。

隨着深度學習技術的大規模使用，不少工程師和研究員共用一個集羣，經過商量 來協調調度顯然不可行了，你們開始使用集羣管理系統調度分佈式做業。這其中， Kubernetes 近年來一枝獨秀，已經在各大公有云中普遍使用。

雲計算和彈性調度

在 Kubernetes 上啓動分佈式 TensorFlow 做業的經常使用方式是使用 Google Cloud 開源的 Kubeflow。Kubeflow 是 Kubernetes 的一個」插件「，它詢問 Kubernetes 計劃分配哪幾臺機器來運行一個分佈式做業中的各個進程，隨後告 知每一個進程，全部其餘進程的 IP 地址和 port。從而保證一個做業裏各個進程 之間互相知道對方。

爲何須要讓全部進程互相知道對方呢？這是 TensorFlow ps-based distribution 方式（上述表格中的左上）要求的。TensorFlow 1.x 原生的分佈 式訓練功能讓一個做業中全部進程都執行 TensorFlow 1.x runtime 程序。這些 進程互相通訊，互相協調成爲一個「分佈式 runtime「，來解釋執行表示深度學習 計算過程的計算圖（graph）。在開始分佈式訓練之初，graph 被 TensorFlow runtime 拆解成若干子圖；每一個進程負責執行一個子圖 —— 任何一個進程失敗 （多是被更高優先級做業搶佔），則整個大圖的執行就失敗了。因此 TensorFlow 原生的分佈式訓練能力不是容錯的（fault-tolerant）。不過， 它是能夠從錯誤恢復（fault-recoverable）—— TensorFlow API 提供 checkpoint 的能力；若是一個做業失敗了，能夠重啓做業，從最近的 checkpoint 開始繼續執行。

Kubeflow 能夠在 Kubernetes 上啓動基於 TensorFlow 原生的分佈式計算能力的做業。可是 由於後者並不能容錯，因此 Kubeflow 並不能無中生有。不能容錯，也意味着不 能彈性調度。

對彈性調度的訴求

在不少人共用計算集羣的狀況下，支持彈性調度意味着極大提高團隊效率和集羣 的整體利用率。前者支持快速迭代以保持技術領先；後者決定企業成本和雲計算 業務的盈利能力。

一個展現彈性調度效果的例子以下。假設一個集羣裏有 N 個 GPU，一個做業包 括一個進程，佔用了 N/2 個 GPU。第二個做業須要 N/2+1 個 GPU；可是此時機 羣裏空閒 GPU 只有 N/2 個。若是沒有彈性調度能力，那麼第二個做業被迫等待， 直到第一個做業結束釋放資源。這個等待時間極可能和第二個做業的運行時間同 量級。此時，集羣的利用率很低，是 50%。若是有彈性調度，那麼第二個做業可 以立刻啓動，用 N/2 個 GPU 作計算。往後若是有更多空閒資源了，調度系統可 以增長其進程數量，充分利用資源。

另外一個例子是，假設有一個做業已經在執行了，此時一個新的更高優先級的做業 須要資源，因此調度系統殺掉了（preempt）了第一個做業的幾個進程來騰出資 源啓動第二個做業。若是沒有彈性調度和容錯，那麼第一個做業會失敗，全部進 程都結束。直到有足夠資源重啓它，而且沿着最近的 checkpoint 繼續。若是有 彈性調度，則第一個做業的剩下的進程能夠繼續執行，只是由於可用的進程 （GPU）少了，因此速度慢一些而已。

以上兩個例子都展現了彈性調度對集羣利用率的提高，以及對團隊工做效率的保 障。須要注意的是：容錯和彈性調度互爲因果。容錯的意思是，做業不受其 中進程數量變化影響。彈性調度時，做業裏的進程數量會隨集羣 workload 狀況 增減，因此做業必須是容錯的，才能和調度系統配合，實現彈性調度。也由於如 此，彈性調度依賴 分佈式編程框架和調度系統配合。

今天，不少分佈式編程框架均可以和 Kubernetes 配合實現容錯和彈性調度。比 如 用於離線數據處理的 Spark、用於在線數據處理的 Storm、在線 流數據引擎 Flink、分佈式存儲系統 Redis 和 HBase。其中適合深度學習的框 架有 Paddle EDL。基於 TensorFlow 的支持彈性調度的深度學習系統，據咱們 所知，ElasticDL 是第一個。

Kubernetes-native 的彈性調度

ElasticDL 經過實現一個 Kubernetes-native 的框架，調用 TensorFlow 2.0， 來實現彈性深度學習。

所謂 Kubernetes-native 指的是一個程序調用 Kubernetes API 來起止進程。 Google MapReduce 是一個 Borg-native 的分佈式計算框架。用戶經過運行一個 Borg 的客戶端程度啓動一個 MapReduce 做業。Borg 客戶端調用 Borg API 提 交做業，而且啓動一個 master 進程。這個 master 調用 Borg API 啓動其餘 workers 進程。ElasticDL 也相似，用戶調用 ElasticDL 的命令行客戶端程序 啓動做業。這個客戶端程序調用 Kubernetes API，啓動 master 進程。master 進程繼續調用 Kubernetes API 啓動其餘進程。master 進程也能夠調用 Kubernetes API 監控其餘進程。

若是 worker 掛了，按照分佈式深度學習訓練算法的數學特性，能夠不用處理， 便可確保訓練過程繼續。若是一個 parameter server 進程掛了，master 會選 擇一個 worker 進程，讓它轉換角色替補上掛掉的 parameter server 進程。在 以上兩種狀況下，master 都會調用 Kubernetes API，請它再啓動一個額外的 worker 進程。若是啓動成功，master 要帶它入門，加入到與其餘進程的協做中。 master 進程的狀態（主要是三個 task queues：todo、doing、done）能夠保留 在 Kubernetes 集羣的 etcd 存儲系統中。這樣，萬一 master 掛了，重啓的 master 進程能夠從 etcd 繼承前世的狀態。

以上是一個簡化的描述。 ElasticDL 實現了多種分佈式計算模式，每種模式實 現 fault-tolerance 的方式略有不一樣。咱們會在後續文章中詳細介紹。

Kubernetes-native 架構使得 master 進程有機會與 Kubernetes 協做實現容錯 和彈性調度。不過，由於 ElasticDL 調用 Kubernetes API，也就意味着 ElasticDL 只能運行在 Kubernetes 上。

TensorFlow 原生的分佈式計算能力不是 Kubernetes-native 的。因此 TensorFlow 不是綁定在 Kubernetes 這個平臺上的。這是你們若是要用現有技 術在 Kubernetes 運行 TensorFlow 做業的話，須要依賴 Kubernetes 的擴展 Kubeflow 的緣由。

理論上，不調用 Kubernetes API 也是能夠實現必定程度的容錯的。即便沒有 Kubernetes 的通知，master 能夠經過檢查其餘繼承的心跳（heartbeat）或者 檢查 TCP 連接狀態，判斷其餘進程的生死存亡。可是，不調用 Kubernetes API （或者其餘調度系統的 API），master 沒法通知調度系統重啓進程，也沒法得 知新啓動的進程的信息，而且幫助它加入做業。這種「非 Kubernetes-native」的 容錯方式頗爲被動，只能接受資源緊張時一些進程被搶佔而掛掉的事實，而不能 在其餘做業釋放資源後增長進程充分利用空閒資源。

TensorFlow 2.0

如上文解釋，爲了保證 TensorFlow 最核心的 runtime 是平臺無關的，咱們沒 法經過修改 runtime 實現完備的主動的容錯和彈性調度。因此如文首的田字格 所示，ElasticDL 和 Uber Horovod 都是在 TensorFlow 的 API 上包一 層。

Horovod 基於 TensorFlow 1.x。 一個 Horovod 做業的每一個進程調用單機版 TensorFlow 作本地計算，而後收集 gradients，而且經過 AllReduce 調用匯聚 gradients 而且更新模型。Horovod 也是平臺無關的，因此它提供的 AllReduce 操做不支持容錯和彈性調度。這一點和 ElasticDL 不同。

和 ElasticDL 同樣的是，Horovod 須要從 TensorFlow 偷偷「截獲」 gradients， 在 TensorFlow 1.x 中，深度學習計算是表示成一個計算圖（graph），而且由 TensorFlow runtime 解釋執行，因此 Horovod 爲了得到每一個進程算的 gradients 而且 AllReduce 它們，就得 hack 進入圖執行的過程。爲此， Horovod 要求使用者使用特定的 optimizer 代替 TensorFlow 提供的 optimizer，從而能夠在優化模型階段透露出 gradients。

一個調用 Horovod 的用戶程序的結構以下。其中標記爲 () 和 (*) 的部 分是 Horovod 要求用戶寫的，幫助 Horovod 截獲 TensorFlow 計算獲得的 gradients 的代碼。若是用戶不慎忘記寫了，那麼程序執行結果就不對了。

 

ElasticDL 沒有這些問題，由於它依賴的是 TensorFlow 2.0。TensorFlow 2.0 主推的 eager execution mode 採用和解釋執行圖徹底不一樣的深度學習計算方式。 相似 PyTorch 的作法，前向計算過程把對基本計算單元（operator）的調用記 錄在一個內存數據結構 tape 裏，隨後，反向計算過程（計算 gradients 的） 能夠回溯這個 tape，以此調用 operator 對應的 gradient operator。這個 tape 提供一個操做讓用戶能夠獲取每一個參數的 gradient。

ElasticDL 經過調用 TensorFlow 2.0 API 能夠很直接地獲取 gradients：

 

並且上面這段代碼不是須要用戶寫的，而是 ElasticDL 的一部分。ElasticDL 用戶須要寫的代碼對應上述 Horovod 代碼範例中的一行 —— 定義模型。

極簡的 API 和使用方式

訓練一個模型不僅須要上述模型定義，還須要指定數據、優化目標（cost）、和 優化算法（optimizer）。用戶老是但願能以儘可能精簡的方式指定這些信息，以 儘可能少的代碼描述訓練做業。

ElasticDL 和 TensorFlow 其餘的 high-level API，例如 Keras 和 TensorFlow Estimator 同樣， 幾乎調用一個 API 函數就能夠執行一個分佈式訓練做業。下 面這個程序使用 Keras。Keras 使用 TensorFlow 原生分佈式訓練能力，不支持容 錯和彈性調度。

 

ElasticDL 的 API 相對更加精簡一些。上述範例程序對應的 ElasticDL 版本以下：

 

主要的區別在於：在 Keras 程序裏用戶要選擇分佈式執行策略；而在 ElasticDL 程序裏則不須要。這是由於 ElasticDL 自動選擇分佈式訓練算法和 策略。

簡單的說，對於有很大參數（須要 model parallelism）的模型，ElasticDL 使 用 asynchrnous SGD。這個方法配合 delayed model update 能把網絡通訊量減 少一個數量級。不少 NLP、搜索、推薦、廣告的模型都符合這一類。 Asynchronous SGD 對於這類模型的表現比較穩定。對於圖像識別和語音識別這 一類參數不太大的模型，ElasticDL 團隊在開發一個 Kubernetes-native 的 AllReduce。和 Horovod 使用的 AllReduce 同樣，ElasticDL AllReduce 把進 程間通訊的拓撲組織成一個環，從而實現高性能的模型更新。與之不一樣的是， ElasticDL AllReduce 是容錯的 —— 在有進程失敗致使 AllReduce 調用失敗的 狀況下，master 組織剩下的活着的進程構造一個新的環。

ElasticDL 項目但願經過這樣的分而治之的策略，提供高性能而且易用的深度學習系統。

ElasticDL 和 SQLFlow 的關係

今年早些時候，王益團隊 開源了 SQLFlow。用戶能夠 用擴展後的 SQL 語法，很是精煉地描述整個數據流和 AI 流程。

好比，若是咱們要爲一個電子商務網站構造一個推薦系統，須要開發日誌收集、 在線數據清洗、特徵工程、模型訓練，驗證和預測等模塊。每一個模塊可能須要投 入一個團隊數軸甚至數月的時間。

最近幾年裏，不少互聯網服務開始把數據直接上傳到通用數據庫中，好比螞蟻金 服的不少數據是在 ODPS（也就是阿里雲上的 MaxCompute 服務）以及新一代的 智能數據系統 。這促使咱們考慮把數據清洗和預處理放在數據庫中作，而特徵工程、自動機器 學習、和訓練過程在 ElasticDL 這樣的 AI 引擎裏作。SQLFlow 把擴展語法的 SQL 程序翻譯成一個 Python 程序，把兩部分連接起來。

在這樣的場景中，若是 AI 須要不少參數，則用戶也就須要在 SQL 程序中提供 這些參數。好比下面 SQL 語句從數據庫中提取用戶的年齡、工做部門、和工做 地點，來預測其收入。

 

其中，TRAIN 從句指定要訓練的模型；COLUMN 從句指定如何把數據映射成 特徵；LABEL 指定要預測的值；WITH 指定訓練過程當中的各類參數，其中 dist_strategy 是調用 Keras/TensorFlow 作訓練是須要指定的分佈式策略， gpus 指定須要的資源。而這些，在 SQLFlow 調用 ElasticDL 的時候都是不 須要的，由於 ElasticDL 自動選擇分佈式策略和算法。

從這個例子能夠看出，若是要讓用戶能提供儘可能少的參數，人工智能引擎還須要 更加智能，提供包括 AutoML 和 自動特徵工程 的功能。 ElasticDL 項目任重道遠。咱們期待把上述 SQL 程序簡化爲以下形式：

 

ElasticDL 項目的現狀

ElasticDL 項目處於早期探索階段。API 還在演化過程當中。此次開源的版本，尚 不包括自動選擇分佈策略和算法的代碼。相比在 TensorFlow runtime 中實現分 布式計算，基於 TensorFlow 2.0 eager mode 的 Python API 實現的分佈式訓 練性能差距還很大。ElasticDL 團隊在和 Google Brain 團隊合做，開發上述 asynchronous SGD + delayed model update 能力、以及 Kubernetes-native AllReduce。但願在下一個版本中能夠提供給你們使用。

目前 ElasticDL 實現的基於 parameter server 的分佈式SGD 訓練方法驗證了 容錯和彈性調度。而且在 Google Cloud 上的 Kubernetes 1.12 集羣和阿里 Sigma 3.1（一個 Kubernetes 的高性能實現）上均可以運行。而且，ElasticDL 團隊開發了 SQLFlow 生成 ElasticDL 程序的 code generator。

咱們但願儘早開源 ElasticDL 和儘早分享其設計意圖，能匯聚來自不一樣公司和 社區的力量，一塊兒探索 Google TensorFlow 2.0 和 Kubernetes 的分佈式訓練 生態，早日實現便捷的端到端的人工智能開發套件。

原文連接

本文爲雲棲社區原創內容，未經容許不得轉載。