文｜沙劍
螞蟻集團高級技術專家
專注分佈式深度學習領域 主要負責螞蟻大規模分佈式訓練引擎的設計和開發
本文 4491 字 閱讀 12分鐘
本文整體介紹了 DLRover 的項目動機與核心能力，未來我們會發布一系列文章，來從同步/異步彈性訓練，優化策略服務，多種集羣和訓練框架對接，策略定製開發等多個角度來介紹 DLRover 的更多細節，敬請期待。
01
技術背景
2022 年 6 月，螞蟻集團決定全面引入 ESG 框架，啓動並確立了“數字普惠”、“綠色低碳”、“科技創新”、“開放生態”四位一體的可持續發展戰略。針對“綠色低碳”，設立了 4 個子議題，包括綠色運營、科技助力產業碳中和、生態保護與修復綠色低碳生活。
在此背景下，綠色 AI 也成爲螞蟻 AI Infra 團隊的一個重要工作方向。作爲綠色 AI 的重要板塊，工程提效項目致力於打造高性能離在線 AI 工程體系，通過提升算力效率和資源利用率，最終達到節省資源降低碳排放的目的。
當前，用戶提交分佈式訓練作業的工具有 Yarn 或者 KubeFlow/Training-Operator。在提交作業時，用戶需要在作業中指定作業資源，包括不同角色的節點數量和資源規格（CPU 核數、內存、GPU 等）。
在訓練作業提交後，作業可能遇到如下問題：
- 集羣資源不足以啓動作業的所有節點，作業只能等待。
- 訓練作業的節點可能會出錯，比如被高優任務搶佔、機器故障、IO 故障等，導致作業失敗。
出現這些問題後，用戶只能修改作業資源來重新提交作業。
針對這兩個問題，螞蟻集團早期基於 Kubernetes 開源了 ElasticDL 項目來支持 K8s 上 TF 2.x 分佈式訓練的彈性容錯。在項目落地過程中我們又發現瞭如下問題：
- 用戶配置的資源可能過少引起 OOM 和訓練性能差。
- 用戶爲了保障作業成功率和速度，通常會配置超額資源導致利用率低。
- 越來越多的用戶使用 PyTorch 或其他 TF 之外的框架來開發和訓練模型。
- 越來越多的分佈式集羣開始支持 AI 作業，比如 Ray、Spark 集羣，能否適配任意計算集羣？
- 在線學習越來越被廣泛採用的情況下，如何運用一套系統同時解決兼容離在線訓練？
前兩個問題使得集羣 CPU 利用率通常只有 20% 上下，同時算法開發人員需要投入很多人工運維成本，爲了解決訓練端資源提效的需求，支持在不同集羣上針對在離線多種訓練模式，給不同框架的分佈式訓練作業自動地尋找最優資源配置。
螞蟻 AI Infra 團隊基於 ElasticDL 彈性容錯的思路，升級擴展並開源了 DLRover，其目標在於提升分佈式模型訓練的智能性，目前很多公司的訓練作業都是跑在混部的集羣中，運行環境複雜多變，正如其名，DLRover 作爲分佈式訓練領域的 “路虎”，不管多麼崎嶇的地形，都可以輕鬆駕馭。
02
整體方案
DLRover 提出了 “ML for System” 的理念來提升分佈式訓練的智能性，那麼這樣的系統應該具備哪些能力呢？
我們認爲主要體現在如下幾個方面：
- 解耦：不和底層訓練框架耦合在一起，只依賴接口抽象，遵循依賴倒置原則。（i.e. Elastic Runtime）
- 資源調度：具備上帝視角的資源調度管控能力。和建立在對作業精準畫像的決策能力。
- 數據驅動：同時收集掌握集羣資源數據，也掌握訓練作業數據。以數據驅動智能。
- 作業交互：以對訓練作業以及模型白盒化的理解，動態根據實際情況，對訓練作業進行優化調整。超越簡單機械的彈性容錯！
- 智能：通過對集羣以及作業信息的收集，結合算法模型+固定策略產出精準的作業優化策略。
我們希望設計並實現一個系統，讓用戶完全擺脫資源配置的束縛，專注於模型訓練本身。在沒有任何資源配置輸入的情況下，DLRover 仍然可以爲每個訓練作業提供最佳資源配置。考慮到用戶可能會以不同的方式運行他們的訓練作業，DLRover 除了面向訓練平臺進行作業統一管理的 Cluster Mode，也提供 Single-Job Mode 方便獨立的算法開發者也能享受到彈性容錯等基本特性。
03
系統架構
DLRover 由四個主要組件組成：ElasticJob、Elastic Trainer、Brain 服務和 Cluster Monitor。
上圖顯示了 DLRover 如何在 K8s 集羣上管理深度學習訓練作業。DLRover 以 ElasticJob CRD 的形式將作業提交到集羣。收到 CRD 後，ElasticJob Operator 會拉起一個 Master Pod 作爲 Elastic Trainer。其從 Brain 服務中獲取初始資源計劃。Elastic Trainer 用它來創建 Scale CRD，並應用 Scale CRD 通知 ElasticJob Controller 啓動所需的 Pod，每個 Pod 將在其上啓動一個 Elastic Agent。
在訓練過程中，Elastic Trainer 的 Training Master 將數據分片分發給 Worker。同時，Cluster Monitor 監控每個作業的運行狀態（i.e.每個節點的 Workload）和集羣狀態（i.e. 資源水位）。這些數據將定期報告給 Brain，Brain 將數據持久化到數據庫中。
然後 DLRover Brain 根據作業的運行狀態，選擇合適的算法生成新的資源計劃，並通知 Elastic Trainer 開始資源調整。
總的來講，DLRover 可以幫助分佈式訓練作業自動化運行在集羣中，可以看作分佈式作業的自動駕駛，模型開發者只需要關注模型的算法設計，DLRover 目前開源版則可以爲用戶提供如下能力：
- 自動資源推導：幫助用戶自動初始化訓練資源，提升資源利用率與作業穩定性。
- 動態訓練數據分片：針對不同 Worker 性能不通造成的木桶效應，根據實際消費速度分配訓練數據，可配合 Failover 記錄消費位點，數據不丟失。
- 單點容錯：提供單點容錯的能力，不需要完整重啓作業。
- 資源彈性：支持運行時 Pod 級和 CPU/Memory 級的資源彈性擴縮容，動態全局優化決策。
04
DLRover 能帶來什麼
1.作業零資源參數配置
用戶提交分佈式作業時無需提供任何資源信息，DLRover 會自動對作業進行畫像，推導出最優的資源配置，同時運行時可以根據實際情況（集羣資源、樣本流量、當前利用率、...）自動對資源進行調整。下面展示了兩種提交腳本的配置對比：
2. 單點容錯提升作業穩定性與恢復效率
DLRover 支持單點恢復 Parameter Server 和 Worker 角色的失敗退出而不需要整體作業重啓，對於非用戶代碼和數據類型的錯誤可以實現用戶無感的重啓。例如集羣中，很常見的一類錯誤是由於用戶配置了不足的內存，導致訓練 OOM。在 DLRover 的幫助下，我們可以自動拉起一個優化配置的節點來恢復失敗的 Node。在真實環境下，DLRover 管理的訓練作業，相比基線的 Kubeflow TF-Operator 作業，訓練成功率從 84% 提升到了 95% 以上。
3. 自動擴縮容提升作業訓練性能
DLRover 針對 Parameter Server 和 Worker 級別都支持在訓練運行時進行自動的調節訓練資源以提升訓練性能。通過監控作業節點的 Workload，DLRover 可以分析資源配置的瓶頸。常見的資源瓶頸有：節點搶佔、Workload 不平衡、CPU 不足導致算力低下、節點數目不足。DLRover 可以通過動態的資源熱更新來持續優化訓練性能。
4. 自動擴縮容提升作業資源利用率
通常不同的模型訓練作業，需要不同的資源配置。然而用戶傾向於超額配置作業的資源以保障作業的成功率。這通常會導致大量的資源浪費。DLRover 的自動擴縮容能力，可以自動根據作業的真實需求配置資源，以最少的資源達到最優的訓練性能，從而減少資源浪費。下圖顯示了自動資源對比手動資源的資源利用率曲線對比：
5. 動態數據分發解決慢節點問題
混部集羣存在資源超賣和搶佔的情況，部分節點消費數據慢，快節點需要等待慢節點，降低訓練速度。DLRover 可以通過數據動態分發給慢節點少分發一些數據，減少等待。此外 DLRover 應該保證訓練任務儘可能按照用戶配置參數消費數據，避免重複消費/丟失數據，這會給訓練帶來不確定性，影響模型性能。
當擴容或者縮容時，需要有個全局協調者知道記錄節點當前消費數據詳情。當節點失敗重啓後，全局協調者需要知道節點已經消費和尚未消費的數據。如果這些邏輯讓訓練節點來做，訓練節點和訓練節點之間需要交互，增加訓練節點邏輯的複雜性。DLRover Master 充當了這個全局協調者的角色。
總而言之，在我們看來，通過動態數據可以簡化訓練節點邏輯的複雜性，訓練節點只管從 DLRover Master 獲取 Shard，然後讀取數據，不需要處理其他的邏輯。
6. 統一離線與在線學習範式
上述動態數據分片特性，實際上幫助我們將 Data Source 和訓練作業進行了解耦，在此基礎上 DLRover 可以同時支持離線訓練，也可以支持消費實時樣本流的在線學習作業。（可以通過 Dlrover.trainer 直接對接樣本流，也可以作爲流計算引擎的訓練 Sink 節點）
在螞蟻的實踐中，DLRover 可以作爲一個理想的組件，來幫助我們構建出一個端到端的在線學習系統。DLRover 可以提供數據源消費位點記錄與恢復，在線學習長跑作業穩定性與性能保障，資源利用率保障等一系列實際問題。我們的開源倉庫中也提供了簡單的範例，後續我們也會開放更多周邊組件。
7. 支持異步和同步訓練模式
訓練集羣中每天都運行着不同業務域性質各異的訓練作業：推薦系統的大規模稀疏模型通常運行在 PS/Worker 架構的訓練模式下進行異步參數更新，資源也多以 CPU 計算爲主。CV/NLP 領域的稠密模型則多以數據並行的方式在 GPU 服務器上進行同步訓練，這時只有 Worker 一種角色。
DLRover 在設計上，可以同時支持同步和異步更新模式，做到針對各種訓練範式的統一。
8. 同訓 DL 練框架解耦
DLRover 支持用戶使用任何自己的訓練框架，底層訓練代碼通過提供約定的 API 接口以實現自動彈性擴縮等需要同底層分佈式代碼深度交互。集羣中部署完成後，終端算法同學基本可以無感接入。
05
總結 & 未來計劃
DLRover 目前已經在螞蟻大規模落地，集羣資源利用率相對於基線穩定獲得了 15% 以上的提升。同時也有效解決了由於資源配置不合理造成的訓練吞吐不及預期的問題。我們希望通過 DLRover 的開源可以幫助更多同行一起推行低碳、綠色、AI 的理念。同時也切實降低模型開發中的運維成本，釋放更多的生產力去解決業務的問題。
當前 DLRover 的調優算法，以及資源，作業畫像策略主要針對螞蟻內部技術棧優化。考慮到不同機構實際技術棧的多樣性，在設計上，DLRover 在 API 層做了統一接口抽象，具體調優算法與作業畫像策略則可靈活自定義。我們歡迎不同機構的開發者也能根據自身特點，同我們一起共建 DLRover 項目，將其發展壯大。
瞭解更多...
DLRover Star 一下✨： https://github.com/intelligent-machine-learning
本週推薦閱讀
