字節跳動正式開源分布式訓練調度框架 Primus
隨著機器學習的發展，模型及訓練模型所需的數據量越來越大，也都趨向于通過分布式訓練實現。而算法工程師通常需要對這些分布式框架涉及到的底層文件存儲和調度系統有較深的理解，才能夠快速批量開啟模型訓練，保證資源利用率。
目前業界有很多類似的框架，如 TonY、TensorFlowOnSpark，Kubeflow 中的 Training Operators 等，但這些框架或多或少存在某些問題，如與固定的機器學習框架（ Tensorflow，Pytorch ）耦合需要寫明例如 PS、Worker 等角色，容錯和彈性調度支持不友好，不支持異構調度，調度語義較為簡單，不支持文件讀取等。
將算法工程師從此類繁重的底層細節中解脫出來、更多地關注到算法層面，即為 Primus 解決的問題。
日均作業百萬核的字節跳動實踐
經過字節跳動在不斷實踐中調整打磨的 Primus，擁有以下能力支撐業務需求：
- 自研訓練框架：目前除了業界開源的 Tensorflow、Pytorch，為了滿足用戶的各種需求，字節也在機器學習場景進行了深入探索，自研多個訓練框架用于滿足用戶需求；
- 擁有底層資源調度系統 YARN 和 Kubernetes：目前階段處在 YARN 向 Kubernetes 的遷移過程中，在此過程中需要對用戶屏蔽底層的調度系統，提供一致的體驗；
- 大規模應用混部資源：由于混部資源不穩定的特點，對訓練的容錯和穩定有著更高的要求；
- 支持復雜調度編排語義：為了使集群資源利用率最大化，需要將合適的容器放在適當的位置上，并需要能夠動態調整并發和容器大小。
- 支持復雜數據源和數據調度需求：支持多種類型數據源和數據類型的混合訓練，如 HDFS、Kafka 等批式流式數據源，Text、PB、Parquet 等數據類型。
目前 Primus 流批一體訓練框架在字節內部支持了抖音、頭條、Tiktok、廣告等大部分業務，每天運行在 Primus 上的作業總核數達到幾百萬。相比舊的基于Hadoop Streaming的框架，訓練性能提升3倍，單次訓練總數據量從TB級提升到了PB級，訓練準備時間由幾十分鐘降低到秒級。
分布式訓練調度框架 Primus
Primus 是一個通用的分布式訓練調度框架，管理了機器學習訓練框架（如 Tensorflow、Pytorch）的生命周期和數據分發，幫助訓練框架獲得更好的分布式能力。
架構介紹
Primus 整體架構
整個 Primus 生命周期分為兩階段，提交階段和執行階段。
- 提交階段
用戶需要描述整個任務的訓練資源，數據輸入以及容錯策略。
- 訓練資源包括需要的角色（如 PS，Worker 等）以及各角色所需的資源，包括其需要的 CPU、內存以及運行腳本、環境變量等。
- 數據輸入用于描述如何把數據提供給訓練器。
- 容錯策略用于描述遇到錯誤時，Primus 需要進行的操作。
Primus Client 通過根據用戶的配置，向 YARN 或 Kubernetes 集群提交 Primus Application Master（之后簡稱 AM）用于管理 Primus 作業集群的運行。
- 執行階段
當作業提交到集群后，AM 會根據用戶的配置向 Resource Manager 申請若干 Executor，并上拉起對應的角色，在持續監控這些角色的健康狀態過程中，如果發現狀態異常，AM 則會根據用戶配置進行相應的操作保證訓練正常運行。
- 數據讀取
Primus 支持讀取數據輸出到訓練器，能夠實現數據的負載均衡并時刻記錄數據的狀態，降低長尾問題，在訓練器有問題時也能夠切換到正常的訓練器繼續進行訓練。AM 通過掃描需要訓練的數據并切分為 Task，可以將 Task 分發給 Executor，并與 Executor 通信記錄 Task 的狀態。
功能介紹
- 多訓練框架支持：Tensorflow、PyTorch、Monolith 等；
- 多調度器支持：YARN、Kubernetes 等；
- 多角色支持：如 PS-Chief-CPU、Worker-GPU、Worker-Evaluator 等，并支持多角色之間的親和反親和等特殊調度策略；
- 多編排策略：支持同時啟動，逐個啟動，基于角色的按順序啟動等（如先啟動 PS，再啟動 Worker）；
- 容錯處理：Worker 失敗自動拉起新 Worker，PS 失敗整體失敗；
- 動態調度：例如支持動態擴大縮小 Worker 數；
- 多數據源數據類型支持：HDFS、Kafka 等；
- 數據負載均衡與狀態保存：支持按 Worker 負載動態分配 Task，如在 Worker 失敗時支持回收 Task 并進行重新分配；
- 多線程高速數據讀取：支持多線程讀取 HDFS 和 Kafka 后輸出到訓練器，提高單訓練器的吞吐。
部署情況
Primus 支撐了字節跳動內部“推薦”“廣告”“搜索”等場景，如頭條推薦、抖音視頻推薦、穿山甲廣告、千川圖文廣告、抖音搜索等業務的超大規模深度學習訓練，日均可達上萬任務的訓練，450W Core資源的使用。
未來規劃
- 開源 Primus 更多能力后續將陸續開放，詳見：https://mp.weixin.qq.com/s/uGBy-WpdjTMUy-7MQAZiww
- Primus 目前與字節內部訓練框架集成較多，Tensorflow 和 Pytorch 可以使用 Primus 的基本能力，后續將開發 Tensorflow 和 Pytorch 適配 Primus API 實現更深度的集成，賦能開源訓練框架。
目前，Primus 已在 Github 上開源，歡迎大家一同參與共建！
項目地址：https://github.com/bytedance/primus
