雲原生的彈性 AI 訓練系列之一：基於 AllReduce 的彈性分散式訓練實踐

引言

隨著模型規模和數據量的不斷增大，分散式訓練已經成為了工業界主流的 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：//github.com/kubeflow/mpi-operator/pull/335 。配合對 Horovod 的修改 //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 方向的進展，可以掃描下面的二維碼加入我們的討論群，一起探討。

Tags: Kubernetes介紹&性能優化系列