引言
隨著模型規模和資料量的不斷增大，分散式訓練已經成為了工業界主流的 AI 模型訓練方式。基於 Kubernetes 的 Kubeflow 專案，能夠很好地承載分散式訓練的工作負載，業已成為了雲原生 AI 領域的事實標準，在諸多企業內廣泛落地。
儘管 Kubeflow 讓基於 Kubernetes 的大規模分散式訓練變得可行，但是雲原生的極致彈性、降本增效等特性在人工智慧場景下沒有得到很好地釋放。
為了解決目前在雲原生 AI 場景下的成本高，資源利用率低等問題，TKE AI 團隊在 Kubeflow 社群中推動了彈性訓練特性的設計與實現。
在這一文章中，我們主要介紹了資料並行的分散式訓練任務的彈效能力在 Kubernetes 上的設計與實現。並且通過實驗的方式驗證了特定的場景下，在保證訓練精度的同時，這一特效能夠使成本降低 70%。
背景
首先我們簡要回顧一下深度學習的模型訓練。這裡所說的訓練，指的是利用資料通過計算梯度下降的方式迭代地去優化神經網路的引數，最終輸出網路模型的過程。在這個過程中，通常在迭代計算的環節，會藉助 GPU 進行計算的加速。相比於 CPU 而言，可以達到 10-100 倍的加速效果。而分散式的模型訓練，最早是由 Mu Li 在 OSDI'14 上提出的。在傳統的模型訓練中，迭代計算的過程只能利用當前程式所在主機上的所有硬體資源。但是單機的擴充套件性始終是有限的，在資料集規模特別大或者模型特別複雜的時候，單機訓練的速度就會有些捉襟見肘。分散式的訓練可以藉助不同主機上的硬體資源進行訓練加速，大大提高訓練速度。
Horovod 是一款基於 AllReduce 的分散式訓練框架。憑藉其對 TensorFlow、PyTorch 等主流深度學習框架的支援，以及通訊優化等特點，Horovod 被廣泛應用於資料並行的訓練中。在 Horovod 中，訓練程式是平等的參與者，每個程式既負責梯度的分發，也負責具體的梯度計算。如下圖所示，三個 Worker 中的梯度被均衡地劃分為三份，通過 4 次通訊，能夠完成叢集梯度的計算和同步。
依託 AllReduce 的分散式訓練由於其簡單易懂的程式設計邏輯和大幅提升的訓練速度，逐漸成為分散式訓練的主流方式。然而，當前這種模式依然存在一些問題：
- 首先，AI 訓練的成本問題顯著。藉助於 Kubernetes，大規模分散式訓練雖然已經不再複雜，但是高昂的訓練成本使得這項技術難以真正做到普惠。
- 其次，相比於單機的訓練，分散式訓練有更大的可能出現任務失敗的情況。在分散式訓練中，有多個程式同時參與訓練，而其中的某個程式出現了問題，整個訓練任務都會因此而失敗。尤其是當訓練任務需要持續幾天甚至幾個禮拜時，這個問題就會顯得尤為嚴重。
- 同時，由於一些混部的叢集存在業務壓力週期性波動的特性，在閒時 GPU 佔用率通常不到 40%。但是與之相對的是，在任務密集提交時，叢集的資源又會出現緊張的情況。資源利用在時間上的不均衡問題非常突出。
彈性訓練
為了解決上述問題，更好地向分散式訓練釋放雲原生的紅利，業界提出了彈性訓練這一概念。
在傳統的深度學習分散式訓練任務中，通常任務的例項配置是固定的。這很大程度上限制了任務的靈活性和訓練速度，對於整個叢集的資源利用率而言也不友好。而彈性訓練，就是指讓訓練任務能夠在執行時動態地調整參與計算的例項數量。這使得訓練更加靈活，同時可以配合叢集的負載進行更好的擴縮容和排程。這一特性為訓練場景帶來了諸多收益：
- 容錯性的提升。在這樣的選型下，所有例項的失敗都是可以容忍的。任務不再會因為某個程式出錯而導致任務整體的失敗。
- 資源利用率的提升。在叢集資源緊張時，通過減少低優先順序訓練任務的例項數量，能夠保證高優先順序訓練任務的資源配額，保證業務的 SLA。在叢集資源閒置時，又可以通過建立更多例項加入訓練的方式，將原本閒置的 GPU 等資源利用起來，加速訓練。這不僅使得任務的訓練速度得到了提升，同時也提高了叢集的資源利用率。
- 實現雲原生的 AI 訓練，配合競價例項等雲上資源更好地降低上雲成本。競價例項相比於按量付費等例項有著非常大的成本優勢，但是也面臨著隨時可能被回收的問題。彈性訓練能夠完美地契合這一場景，在競價例項可用時，在競價例項中建立訓練任務，在競價例項被回收時，訓練任務仍然能夠繼續下去。
彈性分散式訓練能夠很好地解決分散式訓練在成本、資源利用率和容錯等方面的問題。儘管看起來彈性訓練只是能夠將訓練任務的例項動態調整，但是它能夠與公有云提供的雲原生能力產生相互的作用，產生更大的價值。
在我們實際的測試中，基於 Horovod 的彈性訓練在競價例項上，可以將每 GPU 時的花費從 16.21 元降低到了 1.62 元，整個模型訓練的成本可以下降接近 70%。而如果在保持花費不變的情況下，競價例項上的彈性模型訓練可以購買到更多的 GPU 卡，訓練速度能夠提升 5 到 10 倍。原本需要一天的訓練任務，可以在幾個小時內完成。更進一步地，結合彈性訓練與叢集排程，有更多的可能性可以探索。
Horovod 是目前在資料並行的分散式訓練中應用最多的訓練框架之一，因此我們以訓練框架 Horovod 為例，介紹 Horovod 的彈性訓練方案如何在雲原生的環境下落地。
Horovod Elastic
Uber 開源的 Horovod 框架作為資料並行模式下廣泛使用的訓練框架，在 2020 年夏天也開始著手解決彈性訓練這個需求。最終 Elastic Horovod 在 Horovod v0.20.0 版本釋出中面世。
為了實現彈性訓練的能力，Horovod Elastic 對 Horovod 的架構和實現進行了一定的修改，其中主要包括：
- 聚合操作需要被定義在
hvd.elastic.run函式下 - 每個 worker 都有自身的狀態（state），且在訓練之前會被同步一次
- worker 的增減會觸發其他 worker 上的重置（reset）事件
- 重置事件會啟用以下幾個操作（不一定全部執行）：
a. worker 是否應該繼續執行
b. 將失效的 worker 列入黑名單
c. 在新的 hosts 上啟動 worker 程式
d. 更新 worker 的 rank 資訊 - 在重置事件之後，每個 worker 的狀態會被同步
在實際操作中，使用者需要向 horovodrun 提供一個 discover_hosts.sh 指令碼，用以實時反饋當前可用的 hosts 以及每個 hosts 上的 slots（以下用 discover_hosts.sh 指代該指令碼，但該指令碼無需命名為 discover_hosts.sh）。
Horovod Elastic on Kubernetes
在 Elastic 功能推出之前，Kubeflow 社群的 MPI-Operator 是將 Horovod 部署並執行在 Kubernetes 叢集上的主流方案。MPI-Operator 雖然經歷 v1alpha1、v1alpha2 和 v1 三個版本，但大體上的思想一致。其主要過程包括：
- MPIJob Controller 會根據每一份 MPIJob 的配置，生成一個 launcher pod 和對應個數的 worker pod
- MPIJob Controller 會針對每一份 MPIJob 生成一份 ConfigMap，其中包含兩份指令碼，一為反應該任務所有 worker pod 的
hostfile，一為kubexec.sh指令碼 - Launcher pod 上的
mpirun會利用由 ConfigMap 中的kubexel 在 worker pod 中拉起程式；需要注意的是，kubectl的執行有賴於 MPIJob Controller 預先建立的 RBAC 資源（如果對應的 Role 中沒有給 launcher pod 配置在 worker pod 上的執行許可權，launcher pod 在執行kubectl exec` 時會被拒絕）
此前，MPI-Operator 和 Elastic Horovod 存在幾個相容性上的問題。由於 MPI-Operator 的三個版本間存在些許差異，我們這裡只討論 v1 版本：
- MPI-Operator 尚不提供
discover_hosts.sh，這一點直接導致 Elastic Horovod 無法使用 - 當使用者將 worker replicas 調小之後，controller 不會對“額外”的 worker pod 採取任何措施，這會導致 worker pod 無法釋放，訓練任務的例項規模也就無法縮小
- 當使用者增大 worker replica 後，controller 並不會為 launcher pod 的 Role 配置新增 worker 的執行許可權，這會導致 launcher pod 上的 horovodrun 在試圖利用
kubectl在新建立的 worker pod 上執行程式時被 Kubernetes 的許可權管理機制拒絕
基於這些存在的相容性問題，我們在社群上提出了 Elastic Horovod on MPIJob：https://github.com/kubeflow/mpi-operator/pull/335 。配合對 Horovod 的修改 https://github.com/horovod/horovod/pull/2199 ，能夠在 Kubernetes 上實現 Horovod 的彈性訓練。
在該方案中，最關鍵的問題在於如何在 launcher pod 上實現 discover_hosts.sh 的功能。而在 Kubernetes 上實現該功能的關鍵，在於如何獲取當前處在 Running 狀態的 worker pods。這裡有兩種思路。
- MPIJob Controller 構建
discover_hosts.sh並通過 ConfigMap 同步至 launcher pod
- MPIJob Controller 本身就在監聽 pods 相關的資訊，利用 controller 內的 podLister，可以很快地列出每一個 MPIJob 的 worker pods；
- 根據 pods 的 status.phase，controller 在篩選出 Running 狀態的 worker pods 之後，就可以構建出一份反映當前 worker pods 狀態的
discover_hosts.sh； - 通過 ConfigMap，controller 可以將
discover_hosts.sh像hostfile、kubexec.sh指令碼一樣同步至 launcher pod。 - 利用 launcher pod 內已有的
kubectl向 APIServer 實時獲取 worker pod 資訊
2.Launcher pod 自身已經繫結了 pods 的 “get” 和 “list” 許可權，通過 kubectl 或者其他 Kubernetes client 的直接呼叫，即可獲取對應 pod 資訊，通過一樣的篩選標準也可以返回 Elastic Horovod 期待的資訊。
考慮到第二種思路無法限制使用者執行 discover_hosts.sh 的頻率，如果使用者執行過於頻繁或是 MPIJob 規模較大的情況下，會對 Kubernetes 叢集造成較大的壓力，第一種思路在管控上更為全面。
一種對思路二的修正是將 kubectl 或是 client 改為一個 podLister 執行在 launcher pod 中，從而降低對 APIServer 的壓力。然而這種方式使得 launcher pod 中執行了兩個程式。當這個 podLister 程式失效時，缺乏合適的機制將其重新拉起，會造成後續的彈性訓練失效。
因此，我們提議中選擇了第一種思路，這樣一來，controller 通過 ConfigMap 將 discover_hosts.sh 同步至 launcher pod 內，並掛載於 /etc/mpi/discover_hosts.sh 下。同時，該提議中也對 controller 針對另外兩個相容性問題做了相應的修改。這些修改並不會影響到非 Elastic 模式的 MPI 任務，使用者只需忽略 discover_hosts.sh 即可。
當然這種方案也存在一定的問題。ConfigMap 同步至 launcher pod 存在一定的延遲。然而一方面，這個延遲時間是 Kubernetes 管理員可以進行調整的。另一方面相比整個訓練所花的時間，同時也相比 Elastic Horovod 在重置上所花的時間，這一部分延遲也是可以接受的。
彈性訓練演示
最後，我們通過一個示例來演示如何在 Kubernetes 上執行 Horovod 彈性訓練任務。任務建立的過程與普通的訓練任務類似，即通過 MPIJob 建立。
bash-5.0$ kubectl create -f ./tensorflow-mnist-elastic.yaml
mpijob.kubeflow.org/tensorflow-mnist-elastic
createdbash-5.0$ kubectl get po
NAME READY STATUS RESTARTS AGE
tensorflow-mnist-elastic-launcher 1/1 Running 0 14s
tensorflow-mnist-elastic-worker-0 1/1 Running 0 14s
tensorflow-mnist-elastic-worker-1 1/1 Running 0 14s
在示例中，我們一共建立了兩個 worker 參與訓練。在訓練開始後，調整 MPIJob.Spec.MPIReplicaSpecs["Worker"].Replicas 例項數量，增加一個新的 worker 後，觀察例項數量。新的 worker 加入訓練，完成資料集的獲取和初始化之後，訓練任務會不中斷地繼續訓練。其中 discover_hosts.sh 的內容如下：
bash-5.0$ kubectl exec tensorflow-mnist-elastic-launcher -- /etc/mpi/discover_hosts.sh
tensorflow-mnist-elastic-worker-0:1
tensorflow-mnist-elastic-worker-1:1
bash-5.0$ kubectl edit mpijob/tensorflow-mnist-elastic
mpijob.kubeflow.org/tensorflow-mnist-elastic edited
bash-5.0$ kubectl exec tensorflow-mnist-elastic-launcher -- /etc/mpi/discover_hosts.sh
tensorflow-mnist-elastic-worker-0:1
tensorflow-mnist-elastic-worker-1:1
tensorflow-mnist-elastic-worker-2:1
最後，我們再嘗試把例項數量調整為一，訓練叢集中的兩個例項會被回收，而訓練仍然會繼續。
bash-5.0$ kubectl edit mpijob/tensorflow-mnist-elastic
mpijob.kubeflow.org/tensorflow-mnist-elastic edited
bash-5.0$ kubectl get po
NAME READY STATUS RESTARTS AGE
tensorflow-mnist-elastic-launcher 1/1 Running 0 4m48s
tensorflow-mnist-elastic-worker-0 1/1 Running 0 4m48s
tensorflow-mnist-elastic-worker-1 1/1 Terminating 0 4m48s
tensorflow-mnist-elastic-worker-2 1/1 Terminating 0 2m21s
Thu Mar 11 01:53:18 2021[1]<stdout>:Step #40 Loss: 0.284265
Thu Mar 11 01:53:18 2021[0]<stdout>:Step #40 Loss: 0.259497
Thu Mar 11 01:53:18 2021[2]<stdout>:Step #40 Loss: 0.229993
Thu Mar 11 01:54:27 2021[2]<stderr>:command terminated with exit code 137
Process 2 exit with status code 137.
Thu Mar 11 01:54:27 2021[0]<stderr>:command terminated with exit code 137
Process 0 exit with status code 137.
Thu Mar 11 01:54:57 2021[1]<stderr>:[2021-03-11 01:54:57.532928: E /tmp/pip-install-2jy0u7mn/horovod/horovod/common/operations.cc:525] Horovod background loop uncaught exception: [/tmp/pip-install-2jy0u7mn/horovod/third_party/compatible_gloo/gloo/transport/tcp/pair.cc:575] Connection closed by peer [10.244.2.27]:54432
WARNING:root:blacklist failing host: tensorflow-mnist-elastic-worker-2
WARNING:root:blacklist failing host: tensorflow-mnist-elastic-worker-1
Thu Mar 11 01:54:58 2021[1]<stdout>:Step #50 Loss: 0.207741
Thu Mar 11 01:55:00 2021[1]<stdout>:Step #60 Loss: 0.119361
Thu Mar 11 01:55:02 2021[1]<stdout>:Step #70 Loss: 0.131966
這說明通過 MPIJob 的支援，Horovod Elastic 能夠手動地擴縮容，滿足業務需要。在後續的工作中，我們會繼續支援配合 HorizontalPodAutoscaler 的自動擴縮容、指定例項的縮容等高階特性，以滿足更多的場景。
總結
在雲原生技術不斷擴充套件新的場景邊界的過程中，以彈性訓練為代表的 AI 基礎設施新潮流一定會逐漸在工業界落地，與雲原生進行更好地融合，在 AI 場景下充分發揮雲端計算的價值，幫助業務降本增效。
目前，騰訊雲原生 AI 團隊正在積極與 PyTorch、Horovod 等開源社群合作，推進彈性訓練能力在 Kubeflow 中的落地。在這一系列後續的文章中，我們會逐步介紹在 PS Worker 訓練的彈效能力，以及在資源管理和優先順序排程等方面的聯合優化，分享我們在這一方向上的探索和落地實踐。
如果您同樣關注雲原生 AI 方向的進展，可以掃描下面的二維碼加入我們的討論群，一起探討。
