文｜沙劍
螞蟻集團高階技術專家
專注分散式深度學習領域
主要負責螞蟻大規模分散式訓練引擎的設計和開發本文 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
本週推薦閱讀
