第一板塊：開篇
近日，阿里雲機器學習平臺PAI和北京大學楊智老師團隊合作的論文《GoldMiner: Elastic Scaling of Training Data Pre-Processing Pipelines for Deep Learning》被數據庫領域頂會SIGMOD 2023接收。
GoldMiner觀察到深度學習任務中的數據預處理流水線具有無狀態的特點，具有內在的資源彈性。基於此，GoldMiner將數據預處理流水線和模型訓練部分分離執行，通過自動計算圖分析來識別無狀態的數據預處理計算，並對其實現高效的並行加速和彈性伸縮，從而緩解數據預處理瓶頸，提升訓練性能。通過與集羣調度器的協同設計，GoldMiner進一步發揮了數據預處理計算的資源彈性，大幅提升集羣調度效率。實驗顯示GoldMiner可提升訓練性能達12.1倍，提升GPU集羣利用率達2.5倍。
第二板塊：背景
近年來，隨着GPU加速器的不斷進化，以及各類軟件優化技術的層出不窮，深度學習訓練的計算效率正不斷被提升到新的層次。但與此同時，深度學習本質上仍是一種多階段、多資源的任務類型：不僅需要在GPU上進行大量的訓練計算，同時往往也需要CPU端的數據預處理流水線（如數據增強、特徵轉換等），這類預處理計算是訓練出高質量模型的必要步驟。因此，GPU端訓練性能的提升也帶來了更大的數據預處理壓力，使後者成爲新的性能瓶頸。
我們觀察到數據預處理瓶頸對任務訓練性能和集羣資源利用效率都有深刻的影響。一方面，對於單個訓練任務而言，數據預處理瓶頸意味着訓練性能的受損。我們在一臺配備8塊V100 GPU和64個vCPU核的虛擬機上使用一塊GPU和不同的vCPU數進行了性能測試，觀察不同模型需要多少vCPU才能達到最優性能。結果顯示（下圖）大部分模型都需要超過8個vCPU（即一塊GPU平均能得到的vCPU數）才能達到最優性能，甚至部分模型需要消耗掉整機八卡的64個vCPU。這就意味着這類模型在共享集羣中很可能無法得到足夠的CPU資源，導致數據預處理部分性能下降，最終影響訓練效率（下圖右側縱軸表示僅使用8個vCPU時能達到的相對性能）。
而另一方面，上述的問題在雲上場景中會更爲嚴重，影響共享集羣中的資源分配效率。目前企業一般都會構建或購買共享GPU集羣來運行訓練任務，在這種GPU集羣中GPU利用率至關重要。而用戶爲了避免自己任務CPU不足，可能主動提高任務的CPU-GPU配比，然而這些用戶自定義的CPU-GPU配比很容易導致集羣資源碎片。比如某臺機器上由於跑了一些CPU-GPU配比較高的任務，最終CPU比GPU先被分配完，這樣機器上空閒的GPU會無法被分配出去，不僅導致昂貴的GPU資源被浪費，也會提高任務的等待時間。我們在阿里巴巴內部GPU集羣中的觀察發現有近40%的任務等待時間浪費在了這種“GPU足夠但CPU不足”的情況上。
解決上述的兩方面問題的一種手段是將GPU端訓練和CPU端數據預處理分離，使這兩部分計算的資源分配不必捆綁在同一臺機器上。這樣當機器CPU資源不足時可以使用其他機器的資源，一來可以爲單個任務分配更多的CPU資源達到加速效果，同時也緩解了碎片GPU無法分配的問題。實際上這一思路並不是首次被提出，然而要用這種方式提高任務和集羣效率仍存在一系列技術挑戰。
第三板塊：挑戰
儘管已有一些方案（如tf.data service、PyTorch DPP）支持數據預處理計算的分離執行，現有技術仍有以下幾方面挑戰需要解決：
- 計算切分效率：現有技術簡單地通過深度學習框架提供的Dataset/DataLoader API，將這類API內部封裝的計算作爲數據預處理計算分離出去。然而我們發現即便在這類API之外，仍可能存在一部分計算是可以分離執行的，簡單的切分方式錯失了這部分並行加速的機會。
- 用戶代碼侵入性：tf.data service [1]、PyTorch DPP [2]等技術要達到數據預處理的分離執行，均需要用戶重構這部分代碼邏輯，有較大的代碼侵入性。我們希望達到以用戶透明的方式實現這一分離的效果。
- 與集羣調度的結合：在與訓練分離後，數據預處理計算實際上蘊含了內在的資源彈性，而現有技術均沒有在集羣調度的層次挖掘這部分彈性，以提升集羣整體資源利用效率。
第四板塊：破局
GoldMiner是一個自動、彈性的數據預處理服務。如圖所示，GoldMiner使用data worker（DW）和training worker（TW）兩類角色來分別進行數據預處理和訓練兩部分計算。GoldMiner可從原始用戶代碼中自動化地識別出data worker部分的計算（也包含沒有被封裝在Dataset/DataLoader API中的計算）。同時GoldMiner實現了數據預處理計算的彈性伸縮，並通過與集羣調度器的協同設計，進一步提升了集羣效率。
GoldMiner實現這一效果的關鍵在於利用了數據預處理計算的無狀態性。這裏的無狀態是指數據預處理不依賴於模型參數，而模型參數需要在訓練的每一次迭代中被反覆更新，因此不依賴於模型參數的計算即可與訓練部分異步化地執行。我們知道深度學習計算可以被表達爲數據流圖（DFG），GoldMiner通過對於用戶DFG的分析，自動找出其中無狀態的子圖。下圖是一個典型推薦模型的DFG，不同於直接僅切分Dataset的做法（圖中的Simple partition），GoldMiner通過識別與模型參數的依賴關係，將切分範圍自動擴展到後續的一些特徵轉換操作（Expected partition）。實驗顯示通過這一拓展我們可以將data worker並行加速的效果再提升1.6倍。
基於自動化的圖切分，GoldMiner進一步實現了data worker部分子圖在多個data workers間的並行加速以及data workers和training workers間的數據傳遞。利用data worker的無狀態性，這一分佈式執行實現了data workers數量的動態伸縮，從而在集羣空閒資源不斷變化的過程中更高效地利用資源。
爲了充分發揮data worker的資源彈性，GoldMiner提供了一個data worker調度器，調度器會在任務和集羣兩個層次進行動態資源調整。對於每個任務，GoldMiner調整其DW和TW的尺寸以搜索到一個擴展效率最高的配置；在集羣維度上，GoldMiner在不同任務之間動態調整data workers數量分配，以優化某些全局調度目標（如最小化任務完成時間）。這兩層調整利用了一個統一的性能指標，即在DW和TW之間傳遞數據的隊列，該隊列的狀態體現了DW和TW的相對速度，以及增加DW的潛在收益。在一個64-GPU的集羣中進行的實驗顯示GoldMiner調度器可以縮短平均任務完成時間達2.5倍，提升GPU分配率達2.1倍。
第五板塊：應用
我們此前已經在客戶的真實推薦模型上評測了GoldMiner的效果，結果顯示GoldMiner可爲用戶模型加速1.43倍，並可削減13%的訓練成本。目前在落地部署中。
同時我們也開發了PyTorch版本實現，並即將與PAI-DLC集成，向用戶提供數據預處理加速的能力。
第六板塊：
- 論文名字：GoldMiner: Elastic Scaling of Training Data Pre-Processing Pipelines for Deep Learning
- 論文作者：趙漢宇，楊智，程羽，田超，任仕儒，肖文聰，袁滿，陳浪石，劉愷博，張楊，李永，林偉
- 論文pdf鏈接：https://dl.acm.org/doi/pdf/10.1145/3589773
- 參考文獻：
[1] Andrew Audibert, Yang Chen, Dan Graur, Ana Klimovic, Jiri Simsa, Chandramohan A. Thekkath. A Case for Disaggregation of ML Data Processing. https://arxiv.org/abs/2210.14826
[2] Mark Zhao, Niket Agarwal, Aarti Basant, Bugra Gedik, Satadru Pan, Mustafa Ozdal, Rakesh Komuravelli, Jerry Pan, Tianshu Bao, Haowei Lu, Sundaram Narayanan, Jack Langman, Kevin Wilfong, Harsha Rastogi, Carole-Jean Wu, Christos Kozyrakis, Parik Pol. Understanding Data Storage and Ingestion for Large-Scale Deep Recommendation Model Training. ISCA'22
點擊立即免費試用雲產品 開啓雲上實踐之旅！
本文爲阿里雲原創內容，未經允許不得轉載。