9 月 11 日，螞蟻金服在2019谷歌開發者大會上海站上開源了 ElasticDL 項目，這是業界首個基於 TensorFlow 實現彈性深度學習的開源系統。
開源地址爲：elasticdl.org
開源中國採訪了ElasticDL項目負責人王益，對該深度學習系統的技術細節進行了全面介紹。
一、基於 TensorFlow 2.0 和 Kubernetes 實現彈性深度學習
這個基於 Eager Execution 模式的開源項目名爲“ElasticDL”，它是一個 Kubernetes 原生深度學習框架，根據介紹，ElasticDL 主要有四大特點：
- 容錯性
- 彈性調度
- 易用性
- 高效
其中又以容錯與彈性調度特性最具特色。
ElasticDL 實現了容錯和彈性調度的分佈式深度學習，可以極大提升集羣的總體利用率，同時顯著減少用戶提交作業之後等待作業啓動的時間（pending time）。
王益介紹：“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 和集羣的忙碌情況。
1. 基於 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 和集羣剩餘資源情況。”
2. 基於 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 系統的研發。