分享嘉賓：劉柏芳 虎牙

編輯整理：吳春梅?格靈深瞳

出品平臺：DataFunTalk

導讀：本次分享題目為彈性分布式訓練在虎牙的實踐，主要包括以下幾方面內容：

為什么要彈性

彈性分布式訓練的設計

落地效果

未來展望

首先介紹下虎牙AI平臺的發展歷程。

虎牙AI平臺在2019年前還處于混沌狀態：AI開發各自為政、硬件資源不共享、技術棧不統一。為了解決這些問題，2019年我們開始虎牙AI平臺的開發，基于k8s云原生平臺實現資源統一調度，統一開發、訓練、推理流程。從2021年開始，平臺逐步走向降本增效，通過提升任務調度編排效率和AI框架層面的優化，提升了訓練效率以及資源利用率，降低任務隊列等待時長。從2022年虎牙AI平臺開始整體轉向更加精細化運營，比如AI CI/CD、訓練過程可視量化跟蹤等。

本次分享的彈性分布式訓練，就是在降本增效階段提出的一個解決方案。接下來詳細介紹下為什么要彈性。

我們面對的主要問題包括：

首先虎牙直播平臺有一個明顯的高低峰流量趨勢：從天的角度來看，晚上是高峰期、凌晨白天是低峰期；從周的角度看，周末是高峰期，工作日是低峰期。因為推理應用高低峰的存在，GPU使用也存在很明顯的高低峰，在低峰期很多資源是閑置的。

第二個問題是碎片化GPU資源。因為我們調度分配機器給AI算法同事，訓練任務的啟動和停止完全由算法同事主觀來決定。比如有兩臺8卡機器，每臺機器有一張空閑顯卡，即總共有2張空閑顯卡，假設有一個2卡待訓練任務，但因為我們是機器級別分配顯卡資源，所以這個2卡的待訓練任務也無法使用這2張分布在不同機器上的空閑顯卡。如果待訓練任務排隊時間較長，平臺側可能就會被迫增加機器，進而增加了成本。

第三個問題是如果機器異常宕機，必須要算法同事人工參與才能繼續訓練，除了耽誤時間，也可能會丟失一些訓練狀態，給算法同事帶來不小的麻煩。

為了方便大家理解，我們舉例來進一步說明為什么要彈性。

假設任務A正在節點1的兩張顯卡上執行訓練任務，任務B和C在節點2的兩張顯卡上執行訓練任務。一段時間后，任務B和C相繼結束。在過去傳統訓練集群里，如果沒有新任務來，節點2的兩張顯卡將一直是空閑浪費狀態。但是如果平臺支持彈性分布式訓練，那么可以將任務A動態彈性擴容到節點2上，即總共4張顯卡上加速訓練。一段時間后，更高優先級的任務D來臨，任務A可以先縮容自己到最小的2卡，空出節點2給任務D。另外一種情況就是，如果任務A處于4卡訓練時，節點1突然宕機，在過去傳統訓練平臺里，任務A就會報錯退出中斷訓練。但是如果支持彈性分布式訓練，平臺就可以將任務A自動縮至節點2繼續運行，訓練任務不會被中斷。另外整個訓練調度過程的擴縮對算法同事都是無感的，不需要他們人工參與。

理解了為什么需要彈性，接下來我們就討論怎么來設計彈性分布式訓練。彈性分布式訓練與傳統分布式訓練的核心改變是當節點擴縮時，每個節點需要能感知到變化，然后重新建立通信，恢復之前訓練狀態繼續訓練。

接下來我們演示下這個訓練過程：

我們通過ETCD的注冊、選舉和watch監聽功能來感知集群節點的狀態變化。每個節點在啟動時首先將IP、端口號、GPU卡等信息注冊到ETCD，注冊完后，每個節點都可以從ETCD獲取訓練任務相關的所有其他節點信息。

假設某個訓練任務初始階段有3個節點：

1. 啟動時，首先每個節點將自己的信息注冊到ETCD上去，包括IP地址、端口號、GPU卡信息等等，注冊完成之后，節點再從ETCD獲取所有跟它相同任務的其它節點的信息，這樣每個節點都互相知道相同任務所有節點的信息；

2. 然后利用ETCD進行選舉每個節點的角色，比如節點1的角色是rank0，節點2的角色是rank1，節點3的角色是rank2；

3. 有了這些信息后，接下來就可以利用傳統的分布式訓練方式對所有節點建立通訊，比如通過RingAllReduce建立環狀通訊等，建立通訊之后就可以進行分布式訓練了；

4. 訓練進行一段時間后，新節點4加入到集群，它首先也是將自己注冊到ETCD；

5. 其它節點監聽到有新節點加入進來，跑完各自當前訓練step后，會暫停訓練，重新從ETCD獲取最新的節點信息；

6. 重復執行步驟2、3，因為節點4新加入，會有一個狀態同步，同步后訓練任務將會在4個節點上繼續執行。

以上步驟是資源擴容的情況，縮容的步驟與其類似。

接下來介紹平臺整體模塊設計架構圖：

平臺建立在K8S集群基礎之上，通過自定義的k8s operator進行訓練任務啟停操作。具體訓練進程在k8s pod里執行，其中Rendezvous是訓練進程的一個核心組件，它負責與ETCD交互，注冊容器IP、端口號、GPU卡等信息，并且從ETCD同步獲取訓練任務的所有節點信息。

然后通過選舉賦予每張顯卡的角色：rank_0、rank_1……rank_n。接下來通過launcher為每張顯卡啟動一個訓練進程。此時就可以開始分布式訓練了。在訓練過程中，如果某張顯卡上的訓練進程因為內存溢出等原因崩潰，其它訓練進程可以實時捕獲到此異常(通信崩潰)，暫停訓練進程，從ETCD獲取訓練任務的最新所有節點信息，然后重新建立通訊并繼續訓練任務。

另外平臺還有一個組件叫Remote Cache，它緩存訓練過程中的一些中間數據，例如訓練參數、超參數等。在某些場景下，比如有一個低優先級訓練任務使用的全部都GPU被更高優先級任務搶占，此時低優先級訓練任務是掛起狀態。當有空閑資源時，之前被掛起的低優先級任務會被重新調度，平臺會從Remote Cache讀取緩存的數據，從而恢復之前被中斷的低優先級訓練任務的訓練狀態，繼續訓練。

平臺設計完后，我們進行了大量測試驗證效果是否符合預期。下面是我們關心的兩個核心指標：

精度：確保彈性分布式訓練對算法精度的影響在可接受范圍內

訓練耗時：期望彈性分布式訓練利用閑置可以縮短訓練時長

我們測試了模型Resnet50在ImageNet數據集上的表現。下圖是測試結果：

其中EFDL彈性數據為在實際集群中不同時間段分別跑了4次的數據。從表格可以看出彈性分布式訓練和單機多卡訓練相比，精度和總卡時接近，符合預期。從卡數使用趨勢圖可以看到彈性分布式訓練可以充分利用閑置資源，一開始如果集群資源不足，可以利用比較少的卡數將訓練任務先啟動起來，然后等有空閑資源時，可以動態使用更多顯卡以加快訓練速度。

測試符合預期，接下來就是實際落地給算法同學來使用了：

算法同學將傳統訓練任務改成彈性分布式訓練任務時，只需要改大概十幾行代碼。上圖提到引入了EFDL包，它就是我們實現的彈性分布式訓練框架。EFDL框架適配了虎牙內部常用的幾種訓練框架：支持Pytorch DDP分布式訓練、Pytorch Horovod分布式訓練；支持Tensorflow Horovod分布式訓練，包括Keras、Estimator這些高層API。

代碼修改完后，算法同學在AI平臺上新增一個訓練任務時，他們只需要修改訓練模式為EFDL彈性訓練，并設置最小和最大worker數。

下面分享一下我們彈性分布式訓練的落地效果。

對于算法同學來說最大的收益是縮短訓練時長。平臺通過彈性利用低峰空閑GPU資源以及碎片化資源，可無感的成倍的縮短部分訓練時間。同時也可以避免因為機器宕機對訓練造成的中斷。

除了對業務模型訓練帶來收益，平臺的收益也很顯著。傳統的訓練調度策略導致排隊等待時間太長，被迫加機器，造成成本浪費。而彈性訓練調度策略可以通過更加靈活的調度策略，更易于基于優先級進行全公司的訓練任務的編排，從而最大限度的壓榨GPU資源池。最終達到降低成本的目的。

最后介紹一下未來展望。

更易用，目前如果將傳統訓練改成彈性分布式訓練，大概需要修改十多行代碼，并且還需要對彈性有一定認知。為了降低門檻，我們正在迭代重構一個新的版本，新版本需要算法同學改動的代碼會更少，并且也不需要了解彈性概念。

更穩定，這也是每個平臺的基本要求。努力做到運行更穩定，故障可轉移。

更高效，提升分布式訓練的效率是我們不斷持續追求的目標。

更開放，平臺在彈性分布式訓練框架、提升分配效率等方面都借鑒了大量的開源代碼和設計思路。因此我們也希望通過開源平臺來回饋社區，一起分享、一起學習，一起進步，和社區一起將平臺做得更好。