ElasticDL：首個基于 TensorFlow 實現彈性深度學習的開源系統
9 月 11 日，螞蟻金服開源了 ElasticDL 項目，據悉這是業界首個基于 TensorFlow 實現彈性深度學習的開源系統。
Google Brain 成員 Martin Wicke 此前在公開郵件中透露了 TensorFlow 2.0 的規劃，他提到 Eager Execution 將是 2.0 的核心功能。TensorFlow 2.0 還沒有正式發布，但是在目前推出的新版本中已經添加了該強大特性。
簡單來講，TensorFlow Eager Execution 是一種命令式接口，類比 PyTorch，開發者在調用其進行計算時可以直接直觀地得到結果，這使得基于 TensorFlow 的開發更加簡單明了。
利用 Eager Execution 特性，螞蟻金服完成了業界首個基于 TensorFlow 實現彈性深度學習的系統。在 11 日“谷歌開發者大會 2019”（GDD 2019）上，螞蟻金服研究員王益宣布該系統正式開源。
項目地址為：https://github.com/sql-machine-learning/elasticdl。
基于 TensorFlow 2.0 和 Kubernetes 實現彈性深度學習
這個基于 Eager Execution 模式的開源項目名為“ElasticDL”，它是一個 Kubernetes 原生深度學習框架，根據介紹，ElasticDL 主要有四大特點：
其中又以容錯與彈性調度特性最具特色。
ElasticDL 實現了容錯和彈性調度的分布式深度學習，可以極大提升集群的總體利用率，同時顯著減少用戶提交作業之后等待作業啟動的時間（pending time）。
王益向 OSCHINA（開源中國）介紹：“ElasticDL 是我們知道的第一個基于 TensorFlow 實現彈性深度學習的開源系統。具體地說，ElasticDL 是基于 TensorFlow 2.0 和 Kubernetes 實現彈性深度學習的。”
集群效用從 1/N 到 N/N
在深度學習技術研發的早期，公用一個計算集群的人相對少， 計算作業之間的協調可以通過口頭交流實現。開發者更關心縮短運行時間，也就是從作業啟動到結束的這段時間。高性能計算技術（HPC）是解決這個問題的有效途徑，比如 NVIDIA 的 cuBLAS 和 cuDNN 優化高性能數學計算、NCCL 優化 GPU 之間的通信效率。
隨著深度學習技術的大規模應用，在許多工程師和研究員公用一個集群的情況下，通過商量來協調調度顯然不可行，于是大家開始使用集群管理系統調度分布式作業。
Kubernetes 近年來已經逐漸成為集群管理的重要工具，目前已經在各大公有云中廣泛采用。因此，讓 TensorFlow 能更好地運行在 Kubernetes 集群上，同時提升利用集群進行深度學習的效率和資源利用率（效用），顯得非常具有實際意義。
關于提升集群資源利用率，王益舉了一個比較極端的例子：假設一個集群有 N 個 GPU，而一個任務只使用其中一個，現在有一個任務占用了一個 GPU。當沒有彈性調度機制時，一個要求所有 N 個 GPU 的任務需要等待前一個任務結束才能開始，這個等待時間可能高達數天甚至數周，在等待期間，集群的效用是 1/N；而擁有彈性調度能力之后，新的任務可以在 N-1 個 GPU 上立即運行，并且 Kubernetes 可以在第一個任務完成后將占用的 GPU 賦予這個任務，這種情況下，集群整體效用是 100%。
ElasticDL 在容錯與彈性調度上都有不錯的表現，它的現實意義便是高效解決集群效用問題。
ElasticDL 如何實現？
前邊講到集群資源利用率提高的前提其實就是 ElasticDL 的“彈性調度”特性帶來的，而彈性調度依賴于容錯能力。
容錯是指作業不受其中進程數量變化的影響，在彈性調度過程中，作業里的進程數量會隨集群 workload 情況相應增減，所以作業必須是容錯的，才能配合調度系統，實現彈性調度。
在這個過程中，容錯通常由分布式框架實現，比如 Spark 和 ElasticDL 都可以做到當有進程掛掉，或者新的進程加入時，作業不會暫停或者重啟，而是平滑地繼續。而彈性調度是由分布式框架和分布式操作系統（集群管理系統）一起實現的。比如，當有進程掛掉的時候，分布式框架應該通知集群管理系統新啟進程來補位 —— 至于集群管理系統能不能啟動起來，取決于用戶剩余 quota 和集群的忙碌情況。
基于 Kubernetes-native
通常使用 Keras 的 model-fit API 和 Estimator，開發者只需要調用 API 即可進行分布式訓練或預測，然而 ElasticDL 不依賴于 TensorFlow runtime 實現分布式計算，它的實現在 runtime 之外。
ElasticDL 通過 Kubernetes-native 機制來完成分布式計算，而這也為其帶來了容錯性與彈性調度的能力。
所謂 Kubernetes-native 指的是一個程序調用 Kubernetes API 來起止進程，它與 Google MapReduce 的機制類似。MapReduce 是一個 Borg-native 的分布式計算框架，用戶通過運行一個 Borg 客戶端程度啟動一個 MapReduce 作業；Borg 客戶端調用 Borg API 提交作業，并且啟動一個 master 進程；這個 master 調用 Borg API 啟動其它 workers 進程。
在 ElasticDL 中，用戶調用 ElasticDL 的命令行客戶端程序啟動作業；這個客戶端程序調用 Kubernetes API 啟動 master 進程，master 進程繼續調用 Kubernetes API 啟動其它進程。
“ElasticDL 的整個容錯和彈性調度機制都依賴于 Kubernetes-native 架構”，王益介紹：“如果 worker 掛了，按照分布式深度學習訓練算法的數學特性，可以不用處理， 即可確保訓練過程繼續。如果一個 parameter server 進程掛了，master 會選擇一個 worker 進程，讓它轉換角色替補上掛掉的 parameter server 進程。”
在這兩種情況下，master 都會調用 Kubernetes API，請它再啟動一個額外的 worker 進程。如果啟動成功，master 會帶其加入到與其它進程的協作中。master 進程的狀態（主要是三個 task queues：todo、doing 與 done）可以保留在 Kubernetes 集群的 etcd 存儲系統中。
“這樣，萬一 master 掛了，重啟的 master 進程可以從 etcd 繼承前世的狀態。任何進程掛了，master 都會請 Kubernetes 去啟動一個新的進程代替掛掉的進程。而 Kubernetes 是否能完成使命取決于用戶剩余 quota 和集群剩余資源情況。”
基于 TensorFlow 2.0 Eager Execution
為什么 ElasticDL 又基于 TensorFlow 2.0 呢？王益介紹這是因為 TensorFlow 2.0 帶來了 Eager Execution 特性，正是針對這一特性的嘗試，讓開發團隊實現了 Kubernetes-native 的調度方式，從而讓 ElasticDL 支持容錯和彈性調度。
分布式學習需要了解每個進程根據局部訓練數據計算得到的 gradients，才能匯總這些 gradients 來更新模型。
TensorFlow 1.x 的執行方式被稱為 Graph Mode —— 深度學習計算步驟被表示成一個 graph 數據結構，TensorFlow runtime 會解釋執行這個 graph。其中，gradients 的計算過程是 graph 的一部分，所以為了得到 gradients，分布式深度學習系統需要 hack 進入 graph 的執行過程“偷取”gradients。
這個做法需要用戶寫程序的時候寫一些幫助“偷取”的代碼，增加了程序的復雜度，也增加了對編程者的要求。
TensorFlow 2.0 提供的 Eager Execution Mode 中，通過一個叫 tape 的數據結構，它可以把獲取 gradients 的能力以 API 的方式暴露給開發者，ElasticDL 正是以這樣的方式將其實現。
通過這種對比，其實也反映了業界基于 TensroFlow 進行分布式深度學習的不同設計思路。王益介紹，當前基于 TensorFlow 的分布式訓練系統大致可以分為四類：
其中需要修改 TensorFlow runtime 的工作主要由 Google TensorFlow 團隊完成。因為 TensorFlow runtime 是用 C++ 寫的，把網絡通信和同步功能實現在這個層次里，運行效率很高。而且，理論上 C++ 代碼可以通過感知 TCP/IP 鏈接是否中斷，來判斷進程是否掛掉，從而實現容錯。
“但是 TensorFlow runtime 應該是平臺無關的，所以不應該包含訪問特定集群管理系統，請它重啟掛掉的進程的代碼，所以不易實現彈性調度”，王益指出了二者的區別：“與之相對應的，通過調用 TensorFlow API 實現分布式計算的思路，通信性能往往受到 Python 語言性能以及不能在 runtime 內部實現‘微操作’的限制。但它的好處是可以自由調用集群管理系統 API 來管理進程。”
很明顯，ElasticDL 通過 TensorFlow 2.0 帶來的新特性實現了 TensorFlow runtime 外直接調用集群管理 API 而完成了彈性調度。
ElasticDL 替代 Kubeflow 的使用
Kubernetes 本來是一個用來管理無狀態應用的容器平臺，但是當前越來越多公司用它來運行各種各樣的工作負載，特別是使用它來運行機器學習相關任務。
Kubeflow 基于 Kubernetes，它把模型訓練、超參數訓練與模型部署等機器學習任務類型進行組合并以容器化的方式部署，提供了整個機器學習流程各個系統的高可用與便捷性，使用 Kubeflow 就可以進行各種各樣的機器學習任務。
目前 Kubeflow 是在 Kubernetes 上啟動分布式 TenosrFlow 作業的主流操作，這可能也是開發者更為熟悉的模式。
“具體來講，Kubeflow 會詢問 Kubernetes 計劃分配哪幾臺機器來運行一個分布式作業中的各個進程，隨后告知每個進程所有其它進程的 IP 地址和 port，從而保證一個作業里各個進程之間互相知道對方。”
為什么需要讓所有進程互相知道對方呢？這是 TensorFlow ps-based distribution 方式要求的。（也就是前邊提到的對比基于 TensorFlow 的分布式訓練系統表格中左上角的類型）
王益解釋：“TenosrFlow 1.x 原生的分布式訓練功能讓一個作業中所有進程都執行 TensorFlow 1.x runtime 程序。這些進程互相通信，互相協調成為一個‘分布式 runtime’來解釋執行表示深度學習計算過程的 graph。在開始分布式訓練之初，graph 被 TensorFlow runtime 拆解成若干子 graph；每個進程負責執行一個子 graph —— 任何一個進程失敗 （可能是被更高優先級作業搶占），則整個大 graph 的執行就失敗了。所以 TensorFlow 原生的分布式訓練能力不是容錯的（fault-tolerant）。”
不過，TensorFlow Python API 提供了 checkpoint 的能力：如果一個作業失敗了，可以重啟作業，從最近的 checkpoint 開始繼續執行。所以它可以從錯誤中恢復（fault-recoverable）。
Kubeflow 可以在 Kubernetes 上發揮 TensorFlow 原生的分布式計算能力，但是因為后者并不能容錯，所以 Kubeflow 并不能無中生有。不能容錯，也意味著不能彈性調度，而這正是 ElasticDL 的特長。
與 SQLFlow 聯動
前邊介紹了 ElasticDL 的實現機制與現實意義，總結起來主要是因為 ElasticDL 通過 TensorFlow 2.0 提供的新特性 Eager Execution 實現了 TensroFlow runtime 外直接調用集群管理 API，從而實現了 Kubernetes-native 機制來完成分布式計算，繼而實現容錯與彈性調度，最終達到極大提升集群總體利用率的目標。
除此之外，ElasticDL 還有一個重要的特性——易用性。ElasticDL 的易用性與另一個工具密不可分。
幾個月前，螞蟻金服開源了一個機器學習工具 SQLFlow，這個工具旨在讓開發者調用 AI 像寫 SQL 一樣簡單。據介紹，SQLFlow 能夠抽象出端到端從數據到模型的研發過程，配合底層的引擎及自動優化，具備基礎 SQL 知識的開發者即可完成大部分機器學習模型訓練及預測任務。
通過與 SQLFlow 聯動，開發者可以用擴展后的 SQL 語法，非常精煉地描述整個數據流和 AI 流程。SQLFlow 把一個 SQL 程序翻譯成一個實現整個 end-to-end machine learning 的程序，這個程序可以調用 xgboost、PyTorch，或者 ElasticDL、TensorFlow 實現訓練任務。
王益舉例，在有 SQLFlow 之前，如果要為一個電子商務網站構造一個推薦系統，需要開發日志收集、在線數據清洗、特征工程、模型訓練、驗證與預測等模塊，每個模塊可能需要投入一個團隊數周甚至數月的時間。
而 SQLFlow 出現之后，這個流程可以用 SQL 語言描述成一個很簡短的程序，SQLFlow 可以把它翻譯成上述數據和 AI 流。
因為 SQL 是一種只描述意圖，不描述過程的語言，所以 SQL 程序通常都很簡短。但是也因為這個原因，SQL 程序里包含的信息量有限。比如，用戶不會通過 SQL 指定分布式調度和訓練算法。“這些部分需要 ElasticDL 根據模型特點自主決定”，王益補充：“這也是為什么說 ElasticDL 也可以反過來為 SQLFlow 提供易用性。”
ElasticDL 開源的下一步計劃
關于開源后接下來的發展，王益表示，ElasticDL 項目目前處于早期探索階段，API 還在演化過程中。“此次開源的版本，尚不包括自動選擇分布策略和算法的代碼，相比在 TensorFlow runtime 中實現分布式計算，基于 TensorFlow 2.0 Eager Mode 的 Python API 實現的分布式訓練性能差距還很大”，他介紹：“ElasticDL 團隊在和 Google Brain 團隊合作，開發上述 asynchronous SGD + delayed model update 能力、以及 Kubernetes-native AllReduce，希望在下一個版本中可以提供給大家使用。”
隨后王益又具體介紹，上述兩種分布式訓練策略，一種會用于模型中有較大的參數的情況，比如分布式 embedding table，另一種用于模型參數較小的情況。而這也是 ElasticDL 自動決斷分布式訓練算法的一個例子。
另一方面，在 SQLFlow 中，如果要讓用戶能提供盡量少的參數，AI 引擎還需要更加智能，提供包括 AutoML 等功能。
王益感嘆：“ElasticDL 項目任重道遠。”
受訪嘉賓
王益，目前在螞蟻金服負責 AI 基礎架構工作。他于 2007 年從清華大學計算機系博士畢業，先后在 Google（中國）、騰訊、LinkedIn（美國總部）與百度硅谷研究院工作，期間在硅谷和北京各有一次創業經歷。參加工作以來，王益一直專注于 AI 基礎架構工作，參與和領導了多個核心 AI 系統的研發。