分享嘉賓：劉柏芳 虎牙

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

出品平臺：DataFunTalk

導(dǎo)讀：本次分享題目為彈性分布式訓(xùn)練在虎牙的實(shí)踐，主要包括以下幾方面內(nèi)容：

為什么要彈性

彈性分布式訓(xùn)練的設(shè)計(jì)

落地效果

未來展望

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

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

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

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

首先虎牙直播平臺有一個(gè)明顯的高低峰流量趨勢：從天的角度來看，晚上是高峰期、凌晨白天是低峰期；從周的角度看，周末是高峰期，工作日是低峰期。因?yàn)橥评響?yīng)用高低峰的存在，GPU使用也存在很明顯的高低峰，在低峰期很多資源是閑置的。

第二個(gè)問題是碎片化GPU資源。因?yàn)槲覀冋{(diào)度分配機(jī)器給AI算法同事，訓(xùn)練任務(wù)的啟動(dòng)和停止完全由算法同事主觀來決定。比如有兩臺8卡機(jī)器，每臺機(jī)器有一張空閑顯卡，即總共有2張空閑顯卡，假設(shè)有一個(gè)2卡待訓(xùn)練任務(wù)，但因?yàn)槲覀兪菣C(jī)器級別分配顯卡資源，所以這個(gè)2卡的待訓(xùn)練任務(wù)也無法使用這2張分布在不同機(jī)器上的空閑顯卡。如果待訓(xùn)練任務(wù)排隊(duì)時(shí)間較長，平臺側(cè)可能就會被迫增加機(jī)器，進(jìn)而增加了成本。

第三個(gè)問題是如果機(jī)器異常宕機(jī)，必須要算法同事人工參與才能繼續(xù)訓(xùn)練，除了耽誤時(shí)間，也可能會丟失一些訓(xùn)練狀態(tài)，給算法同事帶來不小的麻煩。

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

假設(shè)任務(wù)A正在節(jié)點(diǎn)1的兩張顯卡上執(zhí)行訓(xùn)練任務(wù)，任務(wù)B和C在節(jié)點(diǎn)2的兩張顯卡上執(zhí)行訓(xùn)練任務(wù)。一段時(shí)間后，任務(wù)B和C相繼結(jié)束。在過去傳統(tǒng)訓(xùn)練集群里，如果沒有新任務(wù)來，節(jié)點(diǎn)2的兩張顯卡將一直是空閑浪費(fèi)狀態(tài)。但是如果平臺支持彈性分布式訓(xùn)練，那么可以將任務(wù)A動(dòng)態(tài)彈性擴(kuò)容到節(jié)點(diǎn)2上，即總共4張顯卡上加速訓(xùn)練。一段時(shí)間后，更高優(yōu)先級的任務(wù)D來臨，任務(wù)A可以先縮容自己到最小的2卡，空出節(jié)點(diǎn)2給任務(wù)D。另外一種情況就是，如果任務(wù)A處于4卡訓(xùn)練時(shí)，節(jié)點(diǎn)1突然宕機(jī)，在過去傳統(tǒng)訓(xùn)練平臺里，任務(wù)A就會報(bào)錯(cuò)退出中斷訓(xùn)練。但是如果支持彈性分布式訓(xùn)練，平臺就可以將任務(wù)A自動(dòng)縮至節(jié)點(diǎn)2繼續(xù)運(yùn)行，訓(xùn)練任務(wù)不會被中斷。另外整個(gè)訓(xùn)練調(diào)度過程的擴(kuò)縮對算法同事都是無感的，不需要他們?nèi)斯⑴c。

理解了為什么需要彈性，接下來我們就討論怎么來設(shè)計(jì)彈性分布式訓(xùn)練。彈性分布式訓(xùn)練與傳統(tǒng)分布式訓(xùn)練的核心改變是當(dāng)節(jié)點(diǎn)擴(kuò)縮時(shí)，每個(gè)節(jié)點(diǎn)需要能感知到變化，然后重新建立通信，恢復(fù)之前訓(xùn)練狀態(tài)繼續(xù)訓(xùn)練。

接下來我們演示下這個(gè)訓(xùn)練過程：

我們通過ETCD的注冊、選舉和watch監(jiān)聽功能來感知集群節(jié)點(diǎn)的狀態(tài)變化。每個(gè)節(jié)點(diǎn)在啟動(dòng)時(shí)首先將IP、端口號、GPU卡等信息注冊到ETCD，注冊完后，每個(gè)節(jié)點(diǎn)都可以從ETCD獲取訓(xùn)練任務(wù)相關(guān)的所有其他節(jié)點(diǎn)信息。

假設(shè)某個(gè)訓(xùn)練任務(wù)初始階段有3個(gè)節(jié)點(diǎn)：

1. 啟動(dòng)時(shí)，首先每個(gè)節(jié)點(diǎn)將自己的信息注冊到ETCD上去，包括IP地址、端口號、GPU卡信息等等，注冊完成之后，節(jié)點(diǎn)再從ETCD獲取所有跟它相同任務(wù)的其它節(jié)點(diǎn)的信息，這樣每個(gè)節(jié)點(diǎn)都互相知道相同任務(wù)所有節(jié)點(diǎn)的信息；

2. 然后利用ETCD進(jìn)行選舉每個(gè)節(jié)點(diǎn)的角色，比如節(jié)點(diǎn)1的角色是rank0，節(jié)點(diǎn)2的角色是rank1，節(jié)點(diǎn)3的角色是rank2；

3. 有了這些信息后，接下來就可以利用傳統(tǒng)的分布式訓(xùn)練方式對所有節(jié)點(diǎn)建立通訊，比如通過RingAllReduce建立環(huán)狀通訊等，建立通訊之后就可以進(jìn)行分布式訓(xùn)練了；

4. 訓(xùn)練進(jìn)行一段時(shí)間后，新節(jié)點(diǎn)4加入到集群，它首先也是將自己注冊到ETCD；

5. 其它節(jié)點(diǎn)監(jiān)聽到有新節(jié)點(diǎn)加入進(jìn)來，跑完各自當(dāng)前訓(xùn)練step后，會暫停訓(xùn)練，重新從ETCD獲取最新的節(jié)點(diǎn)信息；

6. 重復(fù)執(zhí)行步驟2、3，因?yàn)楣?jié)點(diǎn)4新加入，會有一個(gè)狀態(tài)同步，同步后訓(xùn)練任務(wù)將會在4個(gè)節(jié)點(diǎn)上繼續(xù)執(zhí)行。

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

接下來介紹平臺整體模塊設(shè)計(jì)架構(gòu)圖：

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

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

另外平臺還有一個(gè)組件叫Remote Cache，它緩存訓(xùn)練過程中的一些中間數(shù)據(jù)，例如訓(xùn)練參數(shù)、超參數(shù)等。在某些場景下，比如有一個(gè)低優(yōu)先級訓(xùn)練任務(wù)使用的全部都GPU被更高優(yōu)先級任務(wù)搶占，此時(shí)低優(yōu)先級訓(xùn)練任務(wù)是掛起狀態(tài)。當(dāng)有空閑資源時(shí)，之前被掛起的低優(yōu)先級任務(wù)會被重新調(diào)度，平臺會從Remote Cache讀取緩存的數(shù)據(jù)，從而恢復(fù)之前被中斷的低優(yōu)先級訓(xùn)練任務(wù)的訓(xùn)練狀態(tài)，繼續(xù)訓(xùn)練。

平臺設(shè)計(jì)完后，我們進(jìn)行了大量測試驗(yàn)證效果是否符合預(yù)期。下面是我們關(guān)心的兩個(gè)核心指標(biāo)：

精度：確保彈性分布式訓(xùn)練對算法精度的影響在可接受范圍內(nèi)

訓(xùn)練耗時(shí)：期望彈性分布式訓(xùn)練利用閑置可以縮短訓(xùn)練時(shí)長

我們測試了模型Resnet50在ImageNet數(shù)據(jù)集上的表現(xiàn)。下圖是測試結(jié)果：

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

測試符合預(yù)期，接下來就是實(shí)際落地給算法同學(xué)來使用了：

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

代碼修改完后，算法同學(xué)在AI平臺上新增一個(gè)訓(xùn)練任務(wù)時(shí)，他們只需要修改訓(xùn)練模式為EFDL彈性訓(xùn)練，并設(shè)置最小和最大worker數(shù)。

下面分享一下我們彈性分布式訓(xùn)練的落地效果。

對于算法同學(xué)來說最大的收益是縮短訓(xùn)練時(shí)長。平臺通過彈性利用低峰空閑GPU資源以及碎片化資源，可無感的成倍的縮短部分訓(xùn)練時(shí)間。同時(shí)也可以避免因?yàn)闄C(jī)器宕機(jī)對訓(xùn)練造成的中斷。

除了對業(yè)務(wù)模型訓(xùn)練帶來收益，平臺的收益也很顯著。傳統(tǒng)的訓(xùn)練調(diào)度策略導(dǎo)致排隊(duì)等待時(shí)間太長，被迫加機(jī)器，造成成本浪費(fèi)。而彈性訓(xùn)練調(diào)度策略可以通過更加靈活的調(diào)度策略，更易于基于優(yōu)先級進(jìn)行全公司的訓(xùn)練任務(wù)的編排，從而最大限度的壓榨GPU資源池。最終達(dá)到降低成本的目的。

最后介紹一下未來展望。

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

更穩(wěn)定，這也是每個(gè)平臺的基本要求。努力做到運(yùn)行更穩(wěn)定，故障可轉(zhuǎn)移。

更高效，提升分布式訓(xùn)練的效率是我們不斷持續(xù)追求的目標(biāo)。

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