隨著機器學(xué)習(xí)的發(fā)展，模型及訓(xùn)練模型所需的數(shù)據(jù)量越來越大，也都趨向于通過分布式訓(xùn)練實現(xiàn)。而算法工程師通常需要對這些分布式框架涉及到的底層文件存儲和調(diào)度系統(tǒng)有較深的理解，才能夠快速批量開啟模型訓(xùn)練，保證資源利用率。

目前業(yè)界有很多類似的框架，如 TonY、TensorFlowOnSpark，Kubeflow 中的 Training Operators 等，但這些框架或多或少存在某些問題，如與固定的機器學(xué)習(xí)框架（ Tensorflow，Pytorch ）耦合需要寫明例如 PS、Worker 等角色，容錯和彈性調(diào)度支持不友好，不支持異構(gòu)調(diào)度，調(diào)度語義較為簡單，不支持文件讀取等。

將算法工程師從此類繁重的底層細(xì)節(jié)中解脫出來、更多地關(guān)注到算法層面，即為 Primus 解決的問題。

日均作業(yè)百萬核的字節(jié)跳動實踐

經(jīng)過字節(jié)跳動在不斷實踐中調(diào)整打磨的 Primus，擁有以下能力支撐業(yè)務(wù)需求：

1. 自研訓(xùn)練框架：目前除了業(yè)界開源的 Tensorflow、Pytorch，為了滿足用戶的各種需求，字節(jié)也在機器學(xué)習(xí)場景進行了深入探索，自研多個訓(xùn)練框架用于滿足用戶需求；
2. 擁有底層資源調(diào)度系統(tǒng) YARN 和 Kubernetes：目前階段處在 YARN 向 Kubernetes 的遷移過程中，在此過程中需要對用戶屏蔽底層的調(diào)度系統(tǒng)，提供一致的體驗；
3. 大規(guī)模應(yīng)用混部資源：由于混部資源不穩(wěn)定的特點，對訓(xùn)練的容錯和穩(wěn)定有著更高的要求；
4. 支持復(fù)雜調(diào)度編排語義：為了使集群資源利用率最大化，需要將合適的容器放在適當(dāng)?shù)奈恢蒙?，并需要能夠動態(tài)調(diào)整并發(fā)和容器大小。
5. 支持復(fù)雜數(shù)據(jù)源和數(shù)據(jù)調(diào)度需求：支持多種類型數(shù)據(jù)源和數(shù)據(jù)類型的混合訓(xùn)練，如 HDFS、Kafka 等批式流式數(shù)據(jù)源，Text、PB、Parquet 等數(shù)據(jù)類型。

目前 Primus 流批一體訓(xùn)練框架在字節(jié)內(nèi)部支持了抖音、頭條、Tiktok、廣告等大部分業(yè)務(wù)，每天運行在 Primus 上的作業(yè)總核數(shù)達到幾百萬。相比舊的基于Hadoop Streaming的框架，訓(xùn)練性能提升3倍，單次訓(xùn)練總數(shù)據(jù)量從TB級提升到了PB級，訓(xùn)練準(zhǔn)備時間由幾十分鐘降低到秒級。

分布式訓(xùn)練調(diào)度框架 Primus

Primus 是一個通用的分布式訓(xùn)練調(diào)度框架，管理了機器學(xué)習(xí)訓(xùn)練框架（如 Tensorflow、Pytorch）的生命周期和數(shù)據(jù)分發(fā)，幫助訓(xùn)練框架獲得更好的分布式能力。

架構(gòu)介紹

Primus 整體架構(gòu)

整個 Primus 生命周期分為兩階段，提交階段和執(zhí)行階段。

1. 提交階段

用戶需要描述整個任務(wù)的訓(xùn)練資源，數(shù)據(jù)輸入以及容錯策略。

• 訓(xùn)練資源包括需要的角色（如 PS，Worker 等）以及各角色所需的資源，包括其需要的 CPU、內(nèi)存以及運行腳本、環(huán)境變量等。
• 數(shù)據(jù)輸入用于描述如何把數(shù)據(jù)提供給訓(xùn)練器。
• 容錯策略用于描述遇到錯誤時，Primus 需要進行的操作。

Primus Client 通過根據(jù)用戶的配置，向 YARN 或 Kubernetes 集群提交 Primus Application Master（之后簡稱 AM）用于管理 Primus 作業(yè)集群的運行。

2. 執(zhí)行階段

當(dāng)作業(yè)提交到集群后，AM 會根據(jù)用戶的配置向 Resource Manager 申請若干 Executor，并上拉起對應(yīng)的角色，在持續(xù)監(jiān)控這些角色的健康狀態(tài)過程中，如果發(fā)現(xiàn)狀態(tài)異常，AM 則會根據(jù)用戶配置進行相應(yīng)的操作保證訓(xùn)練正常運行。

3. 數(shù)據(jù)讀取

Primus 支持讀取數(shù)據(jù)輸出到訓(xùn)練器，能夠?qū)崿F(xiàn)數(shù)據(jù)的負(fù)載均衡并時刻記錄數(shù)據(jù)的狀態(tài)，降低長尾問題，在訓(xùn)練器有問題時也能夠切換到正常的訓(xùn)練器繼續(xù)進行訓(xùn)練。AM 通過掃描需要訓(xùn)練的數(shù)據(jù)并切分為 Task，可以將 Task 分發(fā)給 Executor，并與 Executor 通信記錄 Task 的狀態(tài)。

功能介紹

1. 多訓(xùn)練框架支持：Tensorflow、PyTorch、Monolith 等；
2. 多調(diào)度器支持：YARN、Kubernetes 等；
3. 多角色支持：如 PS-Chief-CPU、Worker-GPU、Worker-Evaluator 等，并支持多角色之間的親和反親和等特殊調(diào)度策略；
4. 多編排策略：支持同時啟動，逐個啟動，基于角色的按順序啟動等（如先啟動 PS，再啟動 Worker）；
5. 容錯處理：Worker 失敗自動拉起新 Worker，PS 失敗整體失敗；
6. 動態(tài)調(diào)度：例如支持動態(tài)擴大縮小 Worker 數(shù)；
7. 多數(shù)據(jù)源數(shù)據(jù)類型支持：HDFS、Kafka 等；
8. 數(shù)據(jù)負(fù)載均衡與狀態(tài)保存：支持按 Worker 負(fù)載動態(tài)分配 Task，如在 Worker 失敗時支持回收 Task 并進行重新分配；
9. 多線程高速數(shù)據(jù)讀?。褐С侄嗑€程讀取 HDFS 和 Kafka 后輸出到訓(xùn)練器，提高單訓(xùn)練器的吞吐。

部署情況

Primus 支撐了字節(jié)跳動內(nèi)部“推薦”“廣告”“搜索”等場景，如頭條推薦、抖音視頻推薦、穿山甲廣告、千川圖文廣告、抖音搜索等業(yè)務(wù)的超大規(guī)模深度學(xué)習(xí)訓(xùn)練，日均可達上萬任務(wù)的訓(xùn)練，450W Core資源的使用。

未來規(guī)劃

1. 開源 Primus 更多能力后續(xù)將陸續(xù)開放，詳見：https://mp.weixin.qq.com/s/uGBy-WpdjTMUy-7MQAZiww
2. Primus 目前與字節(jié)內(nèi)部訓(xùn)練框架集成較多，Tensorflow 和 Pytorch 可以使用 Primus 的基本能力，后續(xù)將開發(fā) Tensorflow 和 Pytorch 適配 Primus API 實現(xiàn)更深度的集成，賦能開源訓(xùn)練框架。

目前，Primus 已在 Github 上開源，歡迎大家一同參與共建！

項目地址：https://github.com/bytedance/primus