彈性分布式深度學(xué)習(xí)系統(tǒng)

9月11日，螞蟻金服在 Google Developer Day Shanghai 2019 上宣布開(kāi)源了基于 TensorFlow 2.0 eager execution 的分布式深度學(xué)習(xí)系統(tǒng) ElasticDL?；?TensorFlow 的支持彈性調(diào)度的深度學(xué)習(xí)系統(tǒng)，據(jù)我們所知，ElasticDL 是第一 個(gè)。項(xiàng)目負(fù)責(zé)人王益和我們分享了 ElasticDL 項(xiàng)目的設(shè)計(jì)意圖和現(xiàn)狀，尤其是 ElasticDL 與 TensorFlow 2.0 以及 Kubernetes 的技術(shù)關(guān)聯(lián)。

　　分布式深度學(xué)習(xí)的技術(shù)思路

　　基于 TensorFlow 的分布式訓(xùn)練系統(tǒng)大致可以分為以下四類(lèi)：

　　其中，ElasticDL 位于田字格的右上角。之所以選擇這條技術(shù)思路，是為了利用 Kubernetes 實(shí)現(xiàn)容錯(cuò)和彈性調(diào)度。

　　高性能計(jì)算和云計(jì)算

　　在深度學(xué)習(xí)技術(shù)研發(fā)的早期，涉及的人員相對(duì)少，共用一個(gè)計(jì)算集群的人相對(duì)少， 計(jì)算作業(yè)之間的協(xié)調(diào)可以通過(guò)口頭交流實(shí)現(xiàn)。大家更關(guān)心縮短運(yùn)行時(shí)間，也就是 從作業(yè)啟動(dòng)到結(jié)束的這段時(shí)間。高性能計(jì)算技術(shù)(HPC)是解決這個(gè)問(wèn)題的有效 途徑，比如 NVIDIA 的 cuBLAS 和 cuDNN 優(yōu)化高性能數(shù)學(xué)計(jì)算、NCCL 優(yōu)化 GPU 之間的通信效率。

　　隨著深度學(xué)習(xí)技術(shù)的大規(guī)模使用，很多工程師和研究員共用一個(gè)集群，通過(guò)商量 來(lái)協(xié)調(diào)調(diào)度顯然不可行了，大家開(kāi)始使用集群管理系統(tǒng)調(diào)度分布式作業(yè)。這其中， Kubernetes 近年來(lái)一枝獨(dú)秀，已經(jīng)在各大公有云中廣泛使用。

　　云計(jì)算和彈性調(diào)度

　　在 Kubernetes 上啟動(dòng)分布式 TensorFlow 作業(yè)的常用方式是使用 Google Cloud 開(kāi)源的 Kubeflow。Kubeflow 是 Kubernetes 的一個(gè)”插件“，它詢問(wèn) Kubernetes 計(jì)劃分配哪幾臺(tái)機(jī)器來(lái)運(yùn)行一個(gè)分布式作業(yè)中的各個(gè)進(jìn)程，隨后告 知每個(gè)進(jìn)程，所有其他進(jìn)程的 IP 地址和 port。從而保證一個(gè)作業(yè)里各個(gè)進(jìn)程 之間互相知道對(duì)方。

　　為什么需要讓所有進(jìn)程互相知道對(duì)方呢?這是 TensorFlow ps-based distribution 方式(上述表格中的左上)要求的。TensorFlow 1.x 原生的分布 式訓(xùn)練功能讓一個(gè)作業(yè)中所有進(jìn)程都執(zhí)行 TensorFlow 1.x runtime 程序。這些 進(jìn)程互相通信，互相協(xié)調(diào)成為一個(gè)“分布式 runtime“，來(lái)解釋執(zhí)行表示深度學(xué)習(xí) 計(jì)算過(guò)程的計(jì)算圖(graph)。在開(kāi)始分布式訓(xùn)練之初，graph 被 TensorFlow runtime 拆解成若干子圖;每個(gè)進(jìn)程負(fù)責(zé)執(zhí)行一個(gè)子圖 —— 任何一個(gè)進(jìn)程失敗 (可能是被更高優(yōu)先級(jí)作業(yè)搶占)，則整個(gè)大圖的執(zhí)行就失敗了。所以 TensorFlow 原生的分布式訓(xùn)練能力不是容錯(cuò)的(fault-tolerant)。不過(guò)， 它是可以從錯(cuò)誤恢復(fù)(fault-recoverable)—— TensorFlow API 提供 checkpoint 的能力;如果一個(gè)作業(yè)失敗了，可以重啟作業(yè)，從最近的 checkpoint 開(kāi)始繼續(xù)執(zhí)行。

　　Kubeflow 可以在 Kubernetes 上啟動(dòng)基于 TensorFlow 原生的分布式計(jì)算能力的作業(yè)。但是 因?yàn)楹笳卟⒉荒苋蒎e(cuò)，所以 Kubeflow 并不能無(wú)中生有。不能容錯(cuò)，也意味著不 能彈性調(diào)度。

　　對(duì)彈性調(diào)度的訴求

　　在很多人共用計(jì)算集群的情況下，支持彈性調(diào)度意味著極大提升團(tuán)隊(duì)效率和集群 的總體利用率。前者支持快速迭代以保持技術(shù)領(lǐng)先;后者決定企業(yè)成本和云計(jì)算 業(yè)務(wù)的盈利能力。

　　一個(gè)展示彈性調(diào)度效果的例子如下。假設(shè)一個(gè)集群里有 N 個(gè) GPU，一個(gè)作業(yè)包 括一個(gè)進(jìn)程，占用了 N/2 個(gè) GPU。第二個(gè)作業(yè)需要 N/2+1 個(gè) GPU;但是此時(shí)機(jī) 群里空閑 GPU 只有 N/2 個(gè)。如果沒(méi)有彈性調(diào)度能力，那么第二個(gè)作業(yè)被迫等待， 直到第一個(gè)作業(yè)結(jié)束釋放資源。這個(gè)等待時(shí)間很可能和第二個(gè)作業(yè)的運(yùn)行時(shí)間同 量級(jí)。此時(shí)，集群的利用率很低，是 50%。如果有彈性調(diào)度，那么第二個(gè)作業(yè)可 以馬上啟動(dòng)，用 N/2 個(gè) GPU 做計(jì)算。日后如果有更多空閑資源了，調(diào)度系統(tǒng)可 以增加其進(jìn)程數(shù)量，充分利用資源。

　　另一個(gè)例子是，假設(shè)有一個(gè)作業(yè)已經(jīng)在執(zhí)行了，此時(shí)一個(gè)新的更高優(yōu)先級(jí)的作業(yè) 需要資源，所以調(diào)度系統(tǒng)殺掉了(preempt)了第一個(gè)作業(yè)的幾個(gè)進(jìn)程來(lái)騰出資 源啟動(dòng)第二個(gè)作業(yè)。如果沒(méi)有彈性調(diào)度和容錯(cuò)，那么第一個(gè)作業(yè)會(huì)失敗，所有進(jìn) 程都結(jié)束。直到有足夠資源重啟它，并且沿著最近的 checkpoint 繼續(xù)。如果有 彈性調(diào)度，則第一個(gè)作業(yè)的剩下的進(jìn)程可以繼續(xù)執(zhí)行，只是因?yàn)榭捎玫倪M(jìn)程 (GPU)少了，所以速度慢一些而已。

　　以上兩個(gè)例子都展示了彈性調(diào)度對(duì)集群利用率的提升，以及對(duì)團(tuán)隊(duì)工作效率的保 障。需要注意的是：容錯(cuò)和彈性調(diào)度互為因果。容錯(cuò)的意思是，作業(yè)不受其 中進(jìn)程數(shù)量變化影響。彈性調(diào)度時(shí)，作業(yè)里的進(jìn)程數(shù)量會(huì)隨集群 workload 情況 增減，所以作業(yè)必須是容錯(cuò)的，才能和調(diào)度系統(tǒng)配合，實(shí)現(xiàn)彈性調(diào)度。也因?yàn)槿? 此，彈性調(diào)度依賴 分布式編程框架和調(diào)度系統(tǒng)配合。

　　今天，很多分布式編程框架都可以和 Kubernetes 配合實(shí)現(xiàn)容錯(cuò)和彈性調(diào)度。比 如 用于離線數(shù)據(jù)處理的 Spark、用于在線數(shù)據(jù)處理的 Storm、在線 流數(shù)據(jù)引擎 Flink、分布式存儲(chǔ)系統(tǒng) Redis 和 HBase。其中適合深度學(xué)習(xí)的框 架有 Paddle EDL。基于 TensorFlow 的支持彈性調(diào)度的深度學(xué)習(xí)系統(tǒng)，據(jù)我們 所知，ElasticDL 是第一個(gè)。

　　Kubernetes-native 的彈性調(diào)度

　　ElasticDL 通過(guò)實(shí)現(xiàn)一個(gè) Kubernetes-native 的框架，調(diào)用 TensorFlow 2.0， 來(lái)實(shí)現(xiàn)彈性深度學(xué)習(xí)。

　　所謂 Kubernetes-native 指的是一個(gè)程序調(diào)用 Kubernetes API 來(lái)起止進(jìn)程。 Google MapReduce 是一個(gè) Borg-native 的分布式計(jì)算框架。用戶通過(guò)運(yùn)行一個(gè) Borg 的客戶端程度啟動(dòng)一個(gè) MapReduce 作業(yè)。Borg 客戶端調(diào)用 Borg API 提 交作業(yè)，并且啟動(dòng)一個(gè) master 進(jìn)程。這個(gè) master 調(diào)用 Borg API 啟動(dòng)其他 workers 進(jìn)程。ElasticDL 也類(lèi)似，用戶調(diào)用 ElasticDL 的命令行客戶端程序 啟動(dòng)作業(yè)。這個(gè)客戶端程序調(diào)用 Kubernetes API，啟動(dòng) master 進(jìn)程。master 進(jìn)程繼續(xù)調(diào)用 Kubernetes API 啟動(dòng)其他進(jìn)程。master 進(jìn)程也可以調(diào)用 Kubernetes API 監(jiān)控其他進(jìn)程。

　　如果 worker 掛了，按照分布式深度學(xué)習(xí)訓(xùn)練算法的數(shù)學(xué)特性，可以不用處理， 即可確保訓(xùn)練過(guò)程繼續(xù)。如果一個(gè) parameter server 進(jìn)程掛了，master 會(huì)選 擇一個(gè) worker 進(jìn)程，讓它轉(zhuǎn)換角色替補(bǔ)上掛掉的 parameter server 進(jìn)程。在 以上兩種情況下，master 都會(huì)調(diào)用 Kubernetes API，請(qǐng)它再啟動(dòng)一個(gè)額外的 worker 進(jìn)程。如果啟動(dòng)成功，master 要帶它入門(mén)，加入到與其他進(jìn)程的協(xié)作中。 master 進(jìn)程的狀態(tài)(主要是三個(gè) task queues：todo、doing、done)可以保留 在 Kubernetes 集群的 etcd 存儲(chǔ)系統(tǒng)中。這樣，萬(wàn)一 master 掛了，重啟的 master 進(jìn)程可以從 etcd 繼承前世的狀態(tài)。

　　以上是一個(gè)簡(jiǎn)化的描述。 ElasticDL 實(shí)現(xiàn)了多種分布式計(jì)算模式，每種模式實(shí) 現(xiàn) fault-tolerance 的方式略有不同。我們會(huì)在后續(xù)文章中詳細(xì)介紹。

　　Kubernetes-native 架構(gòu)使得 master 進(jìn)程有機(jī)會(huì)與 Kubernetes 協(xié)作實(shí)現(xiàn)容錯(cuò) 和彈性調(diào)度。不過(guò)，因?yàn)? ElasticDL 調(diào)用 Kubernetes API，也就意味著 ElasticDL 只能運(yùn)行在 Kubernetes 上。

　　TensorFlow 原生的分布式計(jì)算能力不是 Kubernetes-native 的。所以 TensorFlow 不是綁定在 Kubernetes 這個(gè)平臺(tái)上的。這是大家如果要用現(xiàn)有技 術(shù)在 Kubernetes 運(yùn)行 TensorFlow 作業(yè)的話，需要依賴 Kubernetes 的擴(kuò)展 Kubeflow 的原因。

　　理論上，不調(diào)用 Kubernetes API 也是可以實(shí)現(xiàn)一定程度的容錯(cuò)的。即使沒(méi)有 Kubernetes 的通知，master 可以通過(guò)檢查其他繼承的心跳(heartbeat)或者 檢查 TCP 鏈接狀態(tài)，判斷其他進(jìn)程的生死存亡。但是，不調(diào)用 Kubernetes API (或者其他調(diào)度系統(tǒng)的 API)，master 無(wú)法通知調(diào)度系統(tǒng)重啟進(jìn)程，也無(wú)法得 知新啟動(dòng)的進(jìn)程的信息，并且?guī)椭尤胱鳂I(yè)。這種“非 Kubernetes-native”的 容錯(cuò)方式頗為被動(dòng)，只能接受資源緊張時(shí)一些進(jìn)程被搶占而掛掉的事實(shí)，而不能 在其他作業(yè)釋放資源后增加進(jìn)程充分利用空閑資源。

　　TensorFlow 2.0

　　如上文解釋，為了保證 TensorFlow 最核心的 runtime 是平臺(tái)無(wú)關(guān)的，我們沒(méi) 法通過(guò)修改 runtime 實(shí)現(xiàn)完備的主動(dòng)的容錯(cuò)和彈性調(diào)度。所以如文首的田字格 所示，ElasticDL 和 Uber Horovod 都是在 TensorFlow 的 API 上包一 層。

　　Horovod 基于 TensorFlow 1.x。 一個(gè) Horovod 作業(yè)的每個(gè)進(jìn)程調(diào)用單機(jī)版 TensorFlow 做本地計(jì)算，然后收集 gradients，并且通過(guò) AllReduce 調(diào)用匯聚 gradients 并且更新模型。Horovod 也是平臺(tái)無(wú)關(guān)的，所以它提供的 AllReduce 操作不支持容錯(cuò)和彈性調(diào)度。這一點(diǎn)和 ElasticDL 不一樣。

　　和 ElasticDL 一樣的是，Horovod 需要從 TensorFlow 偷偷“截獲” gradients， 在 TensorFlow 1.x 中，深度學(xué)習(xí)計(jì)算是表示成一個(gè)計(jì)算圖(graph)，并且由 TensorFlow runtime 解釋執(zhí)行，所以 Horovod 為了獲得每個(gè)進(jìn)程算的 gradients 并且 AllReduce 它們，就得 hack 進(jìn)入圖執(zhí)行的過(guò)程。為此， Horovod 要求使用者使用特定的 optimizer 代替 TensorFlow 提供的 optimizer，從而可以在優(yōu)化模型階段透露出 gradients。

　　一個(gè)調(diào)用 Horovod 的用戶程序的結(jié)構(gòu)如下。其中標(biāo)記為 (*) 和 (**) 的部 分是 Horovod 要求用戶寫(xiě)的，幫助 Horovod 截獲 TensorFlow 計(jì)算得到的 gradients 的代碼。如果用戶不慎忘記寫(xiě)了，那么程序執(zhí)行結(jié)果就不對(duì)了。

　　ElasticDL 沒(méi)有這些問(wèn)題，因?yàn)樗蕾嚨氖?TensorFlow 2.0。TensorFlow 2.0 主推的 eager execution mode 采用和解釋執(zhí)行圖完全不同的深度學(xué)習(xí)計(jì)算方式。 類(lèi)似 PyTorch 的做法，前向計(jì)算過(guò)程把對(duì)基本計(jì)算單元(operator)的調(diào)用記 錄在一個(gè)內(nèi)存數(shù)據(jù)結(jié)構(gòu) tape 里，隨后，反向計(jì)算過(guò)程(計(jì)算 gradients 的) 可以回溯這個(gè) tape，以此調(diào)用 operator 對(duì)應(yīng)的 gradient operator。這個(gè) tape 提供一個(gè)操作讓用戶可以獲取每個(gè)參數(shù)的 gradient。

　　ElasticDL 通過(guò)調(diào)用 TensorFlow 2.0 API 可以很直接地獲取 gradients：

　　而且上面這段代碼不是需要用戶寫(xiě)的，而是 ElasticDL 的一部分。ElasticDL 用戶需要寫(xiě)的代碼對(duì)應(yīng)上述 Horovod 代碼范例中的一行 —— 定義模型。

　　極簡(jiǎn)的 API 和使用方式

　　訓(xùn)練一個(gè)模型不只需要上述模型定義，還需要指定數(shù)據(jù)、優(yōu)化目標(biāo)(cost)、和 優(yōu)化算法(optimizer)。用戶總是希望能以盡量精簡(jiǎn)的方式指定這些信息，以 盡量少的代碼描述訓(xùn)練作業(yè)。

　　ElasticDL 和 TensorFlow 其他的 high-level API，例如 Keras 和 TensorFlow Estimator 一樣， 幾乎調(diào)用一個(gè) API 函數(shù)就可以執(zhí)行一個(gè)分布式訓(xùn)練作業(yè)。下 面這個(gè)程序使用 Keras。Keras 使用 TensorFlow 原生分布式訓(xùn)練能力，不支持容 錯(cuò)和彈性調(diào)度。

　　ElasticDL 的 API 相對(duì)更加精簡(jiǎn)一些。上述范例程序?qū)?yīng)的 ElasticDL 版本如下：

　　主要的區(qū)別在于：在 Keras 程序里用戶要選擇分布式執(zhí)行策略;而在 ElasticDL 程序里則不需要。這是因?yàn)?ElasticDL 自動(dòng)選擇分布式訓(xùn)練算法和 策略。

　　簡(jiǎn)單的說(shuō)，對(duì)于有很大參數(shù)(需要 model parallelism)的模型，ElasticDL 使 用 asynchrnous SGD。這個(gè)方法配合 delayed model update 能把網(wǎng)絡(luò)通信量減 少一個(gè)數(shù)量級(jí)。很多 NLP、搜索、推薦、廣告的模型都符合這一類(lèi)。 Asynchronous SGD 對(duì)于這類(lèi)模型的表現(xiàn)比較穩(wěn)定。對(duì)于圖像識(shí)別和語(yǔ)音識(shí)別這 一類(lèi)參數(shù)不太大的模型，ElasticDL 團(tuán)隊(duì)在開(kāi)發(fā)一個(gè) Kubernetes-native 的 AllReduce。和 Horovod 使用的 AllReduce 一樣，ElasticDL AllReduce 把進(jìn) 程間通信的拓?fù)浣M織成一個(gè)環(huán)，從而實(shí)現(xiàn)高性能的模型更新。與之不同的是， ElasticDL AllReduce 是容錯(cuò)的 —— 在有進(jìn)程失敗導(dǎo)致 AllReduce 調(diào)用失敗的 情況下，master 組織剩下的活著的進(jìn)程構(gòu)造一個(gè)新的環(huán)。

　　ElasticDL 項(xiàng)目希望通過(guò)這樣的分而治之的策略，提供高性能并且易用的深度學(xué)習(xí)系統(tǒng)。

　　ElasticDL 和 SQLFlow 的關(guān)系

　　今年早些時(shí)候，王益團(tuán)隊(duì) 開(kāi)源了 SQLFlow。用戶可以 用擴(kuò)展后的 SQL 語(yǔ)法，非常精煉地描述整個(gè)數(shù)據(jù)流和 AI 流程。

　　比如，如果我們要為一個(gè)電子商務(wù)網(wǎng)站構(gòu)造一個(gè)推薦系統(tǒng)，需要開(kāi)發(fā)日志收集、 在線數(shù)據(jù)清洗、特征工程、模型訓(xùn)練，驗(yàn)證和預(yù)測(cè)等模塊。每個(gè)模塊可能需要投 入一個(gè)團(tuán)隊(duì)數(shù)軸甚至數(shù)月的時(shí)間。

　　最近幾年里，很多互聯(lián)網(wǎng)服務(wù)開(kāi)始把數(shù)據(jù)直接上傳到通用數(shù)據(jù)庫(kù)中，比如螞蟻金 服的很多數(shù)據(jù)是在 ODPS(也就是阿里云上的 MaxCompute 服務(wù))以及新一代的 智能數(shù)據(jù)系統(tǒng) 。這促使我們考慮把數(shù)據(jù)清洗和預(yù)處理放在數(shù)據(jù)庫(kù)中做，而特征工程、自動(dòng)機(jī)器 學(xué)習(xí)、和訓(xùn)練過(guò)程在 ElasticDL 這樣的 AI 引擎里做。SQLFlow 把擴(kuò)展語(yǔ)法的 SQL 程序翻譯成一個(gè) Python 程序，把兩部分鏈接起來(lái)。

　　在這樣的場(chǎng)景中，如果 AI 需要很多參數(shù)，則用戶也就需要在 SQL 程序中提供 這些參數(shù)。比如下面 SQL 語(yǔ)句從數(shù)據(jù)庫(kù)中提取用戶的年齡、工作部門(mén)、和工作 地點(diǎn)，來(lái)預(yù)測(cè)其收入。

　　其中，TRAIN 從句指定要訓(xùn)練的模型;COLUMN 從句指定如何把數(shù)據(jù)映射成 特征;LABEL 指定要預(yù)測(cè)的值;WITH 指定訓(xùn)練過(guò)程中的各種參數(shù)，其中 dist_strategy 是調(diào)用 Keras/TensorFlow 做訓(xùn)練是需要指定的分布式策略， gpus 指定需要的資源。而這些，在 SQLFlow 調(diào)用 ElasticDL 的時(shí)候都是不 需要的，因?yàn)?ElasticDL 自動(dòng)選擇分布式策略和算法。

　　從這個(gè)例子可以看出，如果要讓用戶能提供盡量少的參數(shù)，人工智能引擎還需要 更加智能，提供包括 AutoML 和 自動(dòng)特征工程 的功能。 ElasticDL 項(xiàng)目任重道遠(yuǎn)。我們期待把上述 SQL 程序簡(jiǎn)化為如下形式：

　　ElasticDL 項(xiàng)目的現(xiàn)狀

　　ElasticDL 項(xiàng)目處于早期探索階段。API 還在演化過(guò)程中。這次開(kāi)源的版本，尚 不包括自動(dòng)選擇分布策略和算法的代碼。相比在 TensorFlow runtime 中實(shí)現(xiàn)分 布式計(jì)算，基于 TensorFlow 2.0 eager mode 的 Python API 實(shí)現(xiàn)的分布式訓(xùn) 練性能差距還很大。ElasticDL 團(tuán)隊(duì)在和 Google Brain 團(tuán)隊(duì)合作，開(kāi)發(fā)上述 asynchronous SGD + delayed model update 能力、以及 Kubernetes-native AllReduce。希望在下一個(gè)版本中可以提供給大家使用。

　　目前 ElasticDL 實(shí)現(xiàn)的基于 parameter server 的分布式SGD 訓(xùn)練方法驗(yàn)證了 容錯(cuò)和彈性調(diào)度。并且在 Google Cloud 上的 Kubernetes 1.12 集群和阿里 Sigma 3.1(一個(gè) Kubernetes 的高性能實(shí)現(xiàn))上都可以運(yùn)行。并且，ElasticDL 團(tuán)隊(duì)開(kāi)發(fā)了 SQLFlow 生成 ElasticDL 程序的 code generator。

　　我們希望盡早開(kāi)源 ElasticDL 和盡早分享其設(shè)計(jì)意圖，能匯聚來(lái)自不同公司和 社區(qū)的力量，一起探索 Google TensorFlow 2.0 和 Kubernetes 的分布式訓(xùn)練 生態(tài)，早日實(shí)現(xiàn)便捷的端到端的人工智能開(kāi)發(fā)套件。