ElasticDL 是一個基於 TensorFlow 2.x 和 Kubernetes 的開源的分佈式深度學習編程框架。2019 年秋天的 Google Developer Day 活動中來自螞蟻金服的 ElasticDL 團隊展示了 ElasticDL 的第一個開源版本。本文更新這大半年來 ElasticDL 項目的進展，尤其是性能優化和業務落地。
ElasticDL 的首要設計意圖是簡化分佈式編程。 它允許用戶只提供用 TensorFlow 2.0 API 描述的模型，而不需要用戶寫分佈式訓練過程代碼。用戶的模型定義只要能在本地調通，即可在分佈式環境下用大規模數據訓練模型，從而提升研發效率。
同時，ElasticDL 提供的彈性調度的能力在實踐中可以讓集羣的利用高達 90%。當集羣資源不足時，一個訓練作業裏的進程減少；當其他作業結束釋放資源後，進程數量隨之增加。這樣的做法比 TensorFlow Distribution Strategy 專注容錯（進程減少的情況下作業不失敗，但不會增加進程數量）更進一步。並且，因爲 ElasticDL 作業容忍變化的 worker 數量，所以每個作業的啓動都不必等待集羣有足夠的資源，而是可以見縫插針的儘早開始訓練，從而縮短等待作業啓動的時間，讓研發人員可以儘快看到第一個迭代的結果，萬一分佈式訓練有問題，也能儘早發現，從而進一步提升了研發效率。
簡化分佈式深度學習編程
爲了從海量數據中學習規律，我們需要編寫分佈式深度學習程序來完成訓練任務。這在工業場景中尤爲常見。
可分佈式深度學習程序的編寫很難——編程者既要了解深度學習，也要了解分佈式系統開發。在一個分佈式深度學習系統中，需要啓動和監控若干個 workers。因爲既要拆分訓練數據給 workers，還要綜合各個 worker 算出的 gradients 來更新模型，所以涉及通信 (Communication) 和 同步 (Synchronization)。此外，當 worker 數目很多時，作業在執行過程中有 worker 掛掉的概率也會變得很大。如果一個 worker 掛掉，則整個作業重啓或者恢復到最近的 checkpoint (Fault Recovery)，那麼重啓之後可能又會有 worker 掛掉導致重啓，於是作業不斷陷入重啓和恢復，永遠也無法完成。這進一步要求編程者具備設計容錯 (Fault Tolerance) 系統的能力。其實不僅分佈式深度學習，其他分佈式機器學習程序、分佈式離線和在線數據處理程序等各種分佈式程序的寫作，都對編程者有類似上述要求。
一個常見的解決思路是爲特定類型的作業提供分佈式編程框架，讓用戶只需要完形填空一樣補上業務邏輯，而分佈式計算（包括通信、同步、和容錯）都由框架的代碼來完成。一個典型的例子是離線數據處理程序用 MapReduce 框架來寫。不管是 Google MapReduce 還是 Hadoop MapReduce，用戶基本都只需填寫 map 和 reduce 兩個函數的實現即可。類似的，在線數據流系統基於 Storm 和 Flink 來寫，用戶只需提供 bolts 和 nuts 這樣的業務邏輯定義。
在 ElasticDL 之前，螞蟻金服的同事們使用過多種框架和類似框架的高層 API。這些方案大都基於 TensorFlow 和 Kubernetes。
- TensorFlow Estimator 作爲構建在 TensorFlow 之上的一層 API，允許用戶只需定義模 型，而訓練過程封裝在一個函數調用裏。利用 Kubeflow 提供的 TF operator，我們可以將該訓練過程以分佈式作業的方式啓動在Kubernetes 上。這個方案的侷限是：它僅支持 TensorFlow 的 graph mode，不支持 eager execution；而 eager execution 可以大幅簡化調試，尤其方便跟蹤網絡各層輸出。
- Keras API 支持 TensorFlow 2.x 和 eager execution。目前 TensorFlow 2.x Keras API 還暫不支持 ParameterServer 分佈式策略，對 AllReduce 分佈式策略提供了實驗性的支持。
- Horovod 對用戶代碼有侵入性，用戶除了必須熟悉 TensorFlow API 之外，還需學習 Horovod API。
以上三個方案的共同侷限是，雖然具備一定的容錯能力，不過不支持彈性調度。而且它們都依賴部署 Kubernetes operator，瞭解 Kubernetes 對 AI 專家來說頗有挑戰。
| 方案 | 模型定義方式 | 分佈式執行機制 |
|---|---|---|
| Estimator | TensorFlow Estimator API | Kubeflow TF-operator |
| Keras | TensorFlow Keras API | Kubeflow TF-operator |
| Horovod | Horovod with TensorFlow | Kubeflow MPI-operator |
| ElasticDL | TensorFlow Keras API | ElasticDL master process per job |
針對這些侷限，我們設計和開發了 ElasticDL 分佈式計算框架。用戶定義可以用 TensorFlow 2.x 的 Keras API 來定義模型。並且，分佈式執行不要求 Kubernetes 集羣有 任何特殊配置，而是利用每個作業裏的 master 進程來協調訓練數據分配、通信、同步和容錯——這也是 ElasticDL 除了容錯，支持彈性調度的原因。
基於 ElasticDL 框架的編程
就像 MapReduce 框架中只需要用戶完形填空兩個函數：map 和 reduce，ElasticDL需要用戶填寫 forward、loss、optimizer、feed 函數。其中 forward 定義深度學習的前向計算過程 (Forward Pass)，ElasticDL 會調用 TensorFlow eager execution 的 GradientTape 機制來自動推導對應的後向計算過程 (Backward Pass)；loss 函數返回模 型訓練時使用的損失函數；optimizer 函數返回模型訓練時使用的優化器；feed 定製化訓練數據到 TensorFlow 模型輸入 (tensors) 的轉換過程。
所有這些函數的編程只需要瞭解 TensorFlow API，不需要對分佈式訓練有任何背景知識。寫完之後，用戶可以在單機上用小數據做調試驗證。如果通過，可以不做任何代碼修改就提交到 Kubernetes 集羣上做分佈式的容錯的大規模訓練。
不同於 Kubeflow/TF-operator 給每個集羣部署一個 Kubernetes Operator 的方式， ElasticDL 爲每個作業引入一個 master 進程。通過調用 Kubernetes API，master 進程瞭解集羣情況；同時，作爲作業的一部分，master 還了解深度學習作業的特點——包括利用 Python inspection 機制瞭解上述各個函數的特點，其中調用的 API 函數等。所以， master 有非常充分的信息來做更優的調度。比如 master 可以請 Kubernetes 把兩個 worker 啓動在同一臺物理機上，共用一個 GPU——當一個 worker 讀數據的時候，請另外一個 worker 來做計算，從而始終保持較高的 GPU 利用率。
一個例子
我們用一個 MNIST 手寫數字識別的例子來說明。
def forward():
inputs = tf.keras.Input(shape=(28, 28), name="image")
x = tf.keras.layers.Reshape((28, 28, 1))(inputs)
x = tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu")(x)
x = tf.keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu")(x)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
x = tf.keras.layers.Dropout(0.25)(x)
x = tf.keras.layers.Flatten()(x)
outputs = tf.keras.layers.Dense(10)(x)
return tf.keras.Model(inputs=inputs, outputs=outputs, name="mnist_model")
除了模型定義之外，用戶還需要指定 feed, loss, optimizer 函數。
def loss(labels, predictions):
labels = tf.reshape(labels, [-1])
return tf.reduce_mean(
input_tensor=tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=predictions, labels=labels
)
)
def optimizer(lr=0.1):
return tf.optimizers.SGD(lr)
def feed(dataset, mode, _):
def _parse_data(record):
if mode == Mode.PREDICTION:
feature_description = {
"image": tf.io.FixedLenFeature([28, 28], tf.float32)
}
else:
feature_description = {
"image": tf.io.FixedLenFeature([28, 28], tf.float32),
"label": tf.io.FixedLenFeature([1], tf.int64),
}
r = tf.io.parse_single_example(record, feature_description)
features = {
"image": tf.math.divide(tf.cast(r["image"], tf.float32), 255.0)
}
if mode == Mode.PREDICTION:
return features
else:
return features, tf.cast(r["label"], tf.int32)
dataset = dataset.map(_parse_data)
if mode == Mode.TRAINING:
dataset = dataset.shuffle(buffer_size=1024)
return dataset
上述每個函數都很容易做單獨測試 (unit test)。而且，利用 TensorFlow 2.x eager execution，上述函數很容易 log 每一層的輸出。基於個特點，ElasticDL worker 在調用 forward 函數的時候，可以打印中間結果，便於調試和復現問題。
ElasticDL 的彈性訓練過程
給定上述模型定義，ElasticDL 的 master 進程按照 asynchronous 或者 synchronous SGD 方法，協調 workers 來做訓練。當使用 asynchronous SGD 方法時，master 會啓動一個高性能的 parameter server，供各個 workers 使用。當使用 synchronous SGD 時，ElasticDL 使用和才雲科技合作研發的一個 Kubernetes-native 的 fault-tolerable AllReduce 實現 FTlib。
Master 負責動態數據劃分
彈性訓練過程的一個容易被忽略的前提是動態數據劃分 (Dynamic data partitioning)。在用 MPI 寫分佈式程序的時候，因爲作業中進程數量是恆定的，所以經常採用靜態數據 劃分的做法 —— 在訓練之前把訓練數據預先分成 N 個文件，對應作業中的 N 個 worker 進程。這個做法在彈性調度的時候就失效了 —— 因爲彈性調度時，作業中的進程數量是可變的。爲此，需要實現動態數據劃分。
ElasticDL 的動態數據劃分是基於索引的。ElasticDL 要求訓練數據是一個或者多個 RecordIO 格式的文件，或者是 MaxCompute 數據庫系統中的表 (table)。這兩種數據源都允許 master 進程在開始訓練之前，在基本存儲單元 (block) 間快速跳躍着掃描數據，把數據分成小段，稱之爲任務 (task)。每個 task 包括的內容如下：
-
文件名或者表名，
-
第一條記錄相對於文件（或者表）開始處的偏移 (offset)，
-
這個 task 裏的總記錄數。
-
RecordIO
-
MaxCompute
掃描結果是很多 tasks，master 把這些 tasks 放進一個 TODO 隊列裏。這個隊列不一定需要是 master 進程裏的數據結構，可以是放在 etcd 裏的 —— 因爲 etcd 是不死的，所以 master 即使被高優先級作業搶佔了，這個信息也不會丟失；可以通過在資源富餘時重啓 master 進程來恢復作業狀態。
掃描和劃分數據的同時，master 開始請 Kubernetes 啓動 workers，總數不超過用戶指定的數量 N（最大併發度）。每當一個 worker 啓動起來了，master 會收到 Kubernetes 發來的通知；master 在一個 etcd 數據結構裏記錄“活着”的 workers。
掃描和劃分數據結束之後，master 就依次從 TODO 隊列裏取出 task，通過 gRPC 發給某一個活着的 worker，同時 master 把這個 task 挪進 DOING 隊列裏。接收到 task 的 worker 負責打開文件（或者表），並且從指定的 offset 開始依次讀取記錄，並且更新本地模型。根據用戶選擇的 asynchronous 或者 synchronous 算法，workers 會通過調用 parameter server 或者 AllReduce 來協調更新全局模型。
當一個 worker 處理完了接收到的 task，它通過 gRPC 返回一個表示成功的標記；master 就把這個 task 從 DOING 隊列挪到 DONE 隊列了。當所有 task 都從 TODO 挪進了 DONE， 則說明一個 epoch 完成了。
如果一個 worker 失敗了（比如被更高優先級作業搶佔了），則 master 的 gRPC call 會 timeout；此時，master 把對應的 task 從 DOING 隊列挪回 TODO 隊列了。下一次有 worker 完成 task 時，master 會把這個 task 再發出去。這裏有一個細節：有的 task 可能被某個 worker 使用了一部分，也因此影響到了模型更新；此時 worker 被搶佔，那麼這部分已經被處理的數據會因爲 task 的下一次分發，被重複使用。不過這個並不影響機器學習訓練要求數據統計一致性的假設。而且其他動態數據劃分方法造成的數據複用情況可能更嚴重。
Worker 調用 TensorFlow Eager Execution
ElasticDL worker 接收到的一個 task 通常包括多個 minibatches。對於每個 task， worker 打開對應的文件或者表，隨後做如下操作：
- 讀取一個 mini-batch 的訓練數據。
- 用本地模型 (local model) 作爲參數調用用戶定義的 forward 函數以計算 cost。如果 模型很大，則部分參數可能來自於 parameter server。
- 給定 cost，worker 利用 TensorFlow eager execution 的 GradientTape 機制，進行 backward 計算，得到梯度 (gradient)。
- 如果是 synchronous SGD，此時 worker 調用 AllReduce 實現 FTlib 來同步 gradients 並且更新模型。如果是 asynchronous SGD，worker 不定時的向 parameter server 上傳 gradients，也不定時地從 parameter server 獲取全局模型參數。
高效訓練的優化
相對於 2019 年秋季 ElasticDL 在 Google Developer Day 上亮相時的狀態，最近幾個月 ElasticDL 項目針對性能優化做了很多工作。當時 ElasticDL 使用 Redis 作爲 parameter server。現在有了自己的用 Go 語言寫的 parameter server。相對於 Redis， ElasticDL parameter server 可以做一些深度學習計算，從而減少 worker 和 parameter server 之間通信的次數。
這個變化和其他優化工作一起讓同樣的訓練作業，總體訓練時間下降了約 13 倍。最近一個基於 DeepFM 模型的試驗展示，用兩個 parameter server 進程和四個 workers 進程來訓練，10 個 epochs 的總體時間從 1350 秒（ElasticDL 的 2019年9月版本）下降到 106 秒（2020年2月版本）。這些優化策略包括：
- 在 parameter server 上惰性初始化 (lazy initialize) embedding vectors —— 在使用到 vector 的時候才初始化。
- 把一個 embedding table 拆分到多個 parameter server 進程裏以均衡存儲與通信負載。
- worker 從 PS 請求 embedding vectors 時，先濾除重複的 embedding ID，只取回不同 ID 的 vectors，從而減少通信量。
- worker 向 PS 發送梯度時，先把相同 ID 的梯度進行合併（調用 TensorFlow 的 embedding vector combination 函數），從而減少通信量。
彈性調度提升集羣利用率
ElasticDL 實現的彈性調度和剛性調度 (Gang Scheduling) 是對應的。剛性調度的簡潔不求甚解的描述是：一個作業裏的 n 個進程，運行時如果有一個進程掛了（比如被更高優先級的作業搶佔了資源），則整個作業掛掉。等資源足夠再啓動所有的 n 個進程了， 則可以重啓（或者從最近的 checkpoint 恢復）。
-
Gang Scheduling
上文提到的幾種分佈式運行 TensorFlow 作業的方式都使用了 Kubeflow 項目提供的 Kubernetes operators，支持在 Kubernetes 上分佈式地運行 TensorFlow 作業。因爲 TensorFlow runtime 目前支持一定程度的容錯，所以作業執行過程中，如果有一些 workers 掛了，剩下的可以繼續。不過不支持因爲日後資源富餘，恢復 workers 數量。XGBoost、MXNet 社區也習慣於複用 Kubeflow 的 Kubernetes operator。用 MPI 寫的程序也可以用 Kubeflow 拉起。
而彈性調度 (Elastic Scheduling) 實現的是訓練作業運行過程中，進程數量的變化不影響作業進行。具體的說，如果一個或者幾個進程被高優先級的作業搶佔，剩下的進程不受影響地繼續進行。如果將來資源豐沛了，系統可以加幾個進程，此時作業仍然不受影響地繼續運行。
上文簡述了 ElasticDL 實現彈性調度的機制，包括動態數據分配以及由 master 來啓動、監控、和管理 workers，而不依賴 Kubernetes operator。本節展示三個 benchmark 試驗，幫助大家直觀地瞭解 ElasticDL 對集羣利用率和研發效率的同時提升。
實驗一：多個AI訓練作業併發
考慮兩個 AI 訓練作業需要的資源總和略超過集羣的情況：如果沒有 elastic scheduling，則兩個作業順序執行。第二個作業的發起人需要等很久——用戶體驗不好。並且任何時刻只有一個作業在運行——集羣資源用不滿。而如果有 elastic scheduling，則兩個作業併發執行，雖然後啓動的作業拿不到期待的全部資源，但是也馬上就開始執行了——用戶體驗好，而且因爲作業併發集羣被用滿。
我們做了一個實驗來驗證上述好處。這個實驗可以在 ASI 集羣和開源 Kubernetes 集羣上覆現。實驗結果如下圖。
上圖對應的實驗裏，我們用 gang scheduling 的方式提交了兩個訓練作業，每個作業都需要 13 個 CPU。而 Google Cloud 上租用的實驗集羣總 CPU 數是 24， 不足同時運行兩個作業，所以依次運行它們。可以看到第一個作業在 395 秒時結束。隨後集羣花了一點時間調度，然後開始運行第二個作業，直到 795 秒時結束。
下圖對應的實驗裏，我們用 ElasticDL 來執行同樣的兩個訓練作業。第一個作業提交之後的 30 秒，我們提交了第二個作業。第二個作業馬上就開始運行，用滿了集羣剩下的資源，而不需要等到第一個作業結束。在 395 秒時，第一個作業結束。隨後，在 580 秒時，第二個作業也結束了。因爲彈性調度，使得兩個作業儘量同時運行，所以總結束時間比也上圖要早。
總結：
- 用戶等待作業啓動時間幾乎是 0。 這對於 AI 工作很重要，因爲用戶最關注的是第一個迭代儘快開始—— 如果第一個迭代失敗了，很可能是用戶程序的 bug。另外，深度學習模型往往需要手動調優，學習率、optimizer、activation 等配置如果不合理，往往在前幾個迭代就能發現；因此第一個迭代能立刻開始，對模型調優的工作效率提高有很大幫助。
- 集羣利用率高。 第二個實驗 (elastic scheduling) 執行期間，有一段時間集羣利用率是 100%；其他時間也不低於第一個實驗 (gang scheduling)。
- 作業完成更快。 第二個試驗裏，兩個作業用了約 580 秒；第一個實驗裏需要約 795 秒。
實驗二：AI作業和在線服務混布
運行各種在線服務的生產集羣，通常需要留出餘量資源，以應付突然增長的用戶請求量。我們希望利用這些“餘量”來做 AI 訓練，從而提升集羣利用率。下面實驗驗證：通過用較低優先級運行 ElasticDL 訓練作業，在用戶請求增加的時候，Kubernetes 自動擴容在線服務 (NGINX)；此時 ElasticDL 作業自動釋放資源，配合在線服務的擴容。當流量高峯過去之後，Kubernetes 自動縮容 NGINX 服務，此時，ElasticDL 自動利用釋放的資源。
圖中紫色曲線是 NGINX 服務使用的 CPU 數量，隨用戶請求數量變化。綠色曲線是 ElasticDL 訓練作業使用的 CPU 數量，隨 NGINX 的資源需求自動變化。藍色曲線是集羣的總體資源利用率——保持在 90% 以上。
實驗三：訓練時更改 worker 數量不影響收斂性
有用戶擔心訓練過程中 worker 的數量發生變化，會導致不收斂。實際情況下從未發生這類問題。用 ElasticDL 和用 gang scheduling 分別訓練 Wide & Deep model 和 xDeepFM model， 收斂曲線如下：
可以看到，採用 gang scheduling 持續用 4 個或者 8 個 workers，和用 ElasticDL 並且 worker 數量在 4 到 8 之間變化，得到的收斂曲線很難分辨。差別在自然誤差範圍之內。
總結
螞蟻金服從事的金融行業涉及支付、微貸、和保險等業務。和搜索、廣告、推薦不同，金融業務的流程要複雜得多——包括對用戶信用的預判以及和其他金融機構的聯動——每一個用戶請求對應很多處理步驟；而搜索、廣告、推薦業務裏針對每個用戶請求的 AI 處理步驟少得多。行業特點導致螞蟻金服要訓練的模型的類型繁多，呈現更長尾的特點。也對工具提升研發效率提出了高要求。ElasticDL 正是針對這些特點設計的。
同時，對集羣的利用率提升是各行各業都關注的。在很多公司和行業，AI 集羣的利用率通常在 30% 以下。當通過全面實現彈性調度，把集羣利用率提升到 90% 左右時，相當於空手套白狼地把集羣規模擴大了爲原來的三倍多。因此節省的硬件投資可能高達數千萬甚至數億元人民幣。
ElasticDL 的設計和實現依託了 TensorFlow 2.x 提供的高效率的模型描述 API。也依賴了 TensorFlow eager execution 提供的 GradientTape 機制 —— 使得 ElasticDL 可以在不改變 TensorFlow runtime 的情況下，結合 Kubernetes 實現徹底的彈性調度（進程數可增也可減），從而實現了減少作業啓動 的等待時間，提升集羣利用率，和提升研發效率的效果。
目前 ElasticDL 在阿里繫結合 PAI 平臺在推廣。PAI 平臺提供的拖拽式編程模式進一步降低了端到端機器學習流程的研發門檻。希望接下來 ElasticDL 團隊可以有更多結合業務實踐的分享。
作者介紹：
螞蟻集團， ElasticDL 主力工程師齊俊
本文轉載自公衆號支付寶技術（ID：Ant-Techfin）。
原文鏈接：