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

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

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

日均作業(yè)百萬(wàn)核的字節(jié)跳動(dòng)實(shí)踐

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

1. 自研訓(xùn)練框架：目前除了業(yè)界開(kāi)源的 Tensorflow、Pytorch，為了滿足用戶的各種需求，字節(jié)也在機(jī)器學(xué)習(xí)場(chǎng)景進(jìn)行了深入探索，自研多個(gè)訓(xùn)練框架用于滿足用戶需求；
2. 擁有底層資源調(diào)度系統(tǒng) YARN 和 Kubernetes：目前階段處在 YARN 向 Kubernetes 的遷移過(guò)程中，在此過(guò)程中需要對(duì)用戶屏蔽底層的調(diào)度系統(tǒng)，提供一致的體驗(yàn)；
3. 大規(guī)模應(yīng)用混部資源：由于混部資源不穩(wěn)定的特點(diǎn)，對(duì)訓(xùn)練的容錯(cuò)和穩(wěn)定有著更高的要求；
4. 支持復(fù)雜調(diào)度編排語(yǔ)義：為了使集群資源利用率最大化，需要將合適的容器放在適當(dāng)?shù)奈恢蒙?，并需要能夠?dòng)態(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ù)，每天運(yùn)行在 Primus 上的作業(yè)總核數(shù)達(dá)到幾百萬(wàn)。相比舊的基于Hadoop Streaming的框架，訓(xùn)練性能提升3倍，單次訓(xùn)練總數(shù)據(jù)量從TB級(jí)提升到了PB級(jí)，訓(xùn)練準(zhǔn)備時(shí)間由幾十分鐘降低到秒級(jí)。

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

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

架構(gòu)介紹

Primus 整體架構(gòu)

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

1. 提交階段

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

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

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

2. 執(zhí)行階段

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

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

Primus 支持讀取數(shù)據(jù)輸出到訓(xùn)練器，能夠?qū)崿F(xiàn)數(shù)據(jù)的負(fù)載均衡并時(shí)刻記錄數(shù)據(jù)的狀態(tài)，降低長(zhǎng)尾問(wèn)題，在訓(xùn)練器有問(wèn)題時(shí)也能夠切換到正常的訓(xùn)練器繼續(xù)進(jìn)行訓(xùn)練。AM 通過(guò)掃描需要訓(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. 多編排策略：支持同時(shí)啟動(dòng)，逐個(gè)啟動(dòng)，基于角色的按順序啟動(dòng)等（如先啟動(dòng) PS，再啟動(dòng) Worker）；
5. 容錯(cuò)處理：Worker 失敗自動(dòng)拉起新 Worker，PS 失敗整體失?。?/li>
6. 動(dòng)態(tài)調(diào)度：例如支持動(dòng)態(tài)擴(kuò)大縮小 Worker 數(shù)；
7. 多數(shù)據(jù)源數(shù)據(jù)類型支持：HDFS、Kafka 等；
8. 數(shù)據(jù)負(fù)載均衡與狀態(tài)保存：支持按 Worker 負(fù)載動(dòng)態(tài)分配 Task，如在 Worker 失敗時(shí)支持回收 Task 并進(jìn)行重新分配；
9. 多線程高速數(shù)據(jù)讀?。褐С侄嗑€程讀取 HDFS 和 Kafka 后輸出到訓(xùn)練器，提高單訓(xùn)練器的吞吐。

部署情況

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

未來(lái)規(guī)劃

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

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

項(xiàng)目地址：https://github.com/bytedance/primus