彈性分散式深度學習系統

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 的分散式訓練 生態，早日實現便捷的端到端的人工智慧開發套件。

　　瀋陽看性病去哪家醫院：