分享嘉賓：劉柏芳 虎牙
編輯整理：吳春梅 格靈深瞳
出品平臺：DataFunTalk
導讀：
本次分享題目為彈性分散式訓練在虎牙的實踐，主要包括以下幾方面內容：
為什麼要彈性
彈性分散式訓練的設計
落地效果
未來展望
01
為什麼要彈性
首先介紹下虎牙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繼續執行，訓練任務不會被中斷。另外整個訓練排程過程的擴縮對演算法同事都是無感的，不需要他們人工參與。
02
彈性分散式訓練的設計
理解了為什麼需要彈性，接下來我們就討論怎麼來設計彈性分散式訓練。彈性分散式訓練與傳統分散式訓練的核心改變是當節點擴縮時，每個節點需要能感知到變化，然後重新建立通訊，恢復之前訓練狀態繼續訓練。
接下來我們演示下這個訓練過程：
我們透過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數。
03
落地效果
下面分享一下我們彈性分散式訓練的落地效果。
對於演算法同學來說最大的收益是縮短訓練時長。平臺透過彈性利用低峰空閒GPU資源以及碎片化資源，可無感的成倍的縮短部分訓練時間。同時也可以避免因為機器宕機對訓練造成的中斷。
除了對業務模型訓練帶來收益，平臺的收益也很顯著。傳統的訓練排程策略導致排隊等待時間太長，被迫加機器，造成成本浪費。而彈性訓練排程策略可以透過更加靈活的排程策略，更易於基於優先順序進行全公司的訓練任務的編排，從而最大限度的壓榨GPU資源池。最終達到降低成本的目的。
04
未來展望
最後介紹一下未來展望。
更易用
，目前如果將傳統訓練改成彈性分散式訓練，大概需要修改十多行程式碼，並且還需要對彈性有一定認知。為了降低門檻，我們正在迭代重構一個新的版本，新版本需要演算法同學改動的程式碼會更少，並且也不需要了解彈性概念。
更穩定
，這也是每個平臺的基本要求。努力做到執行更穩定，故障可轉移。
更高效
，提升分散式訓練的效率是我們不斷持續追求的目標。
更開放
，平臺在彈性分散式訓練框架、提升分配效率等方面都借鑑了大量的開原始碼和設計思路。因此我們也希望透過開源平臺來回饋社群，一起分享、一起學習，一起進步，和社群一起將平臺做得更好。
今天的分享就到這裡，謝謝大家。
炒股開戶享福利，入金抽188元紅包，100%中獎！
開啟App看更多精彩內容