字節(jié)跳動(dòng)正式開(kāi)源分布式訓(xùn)練調(diào)度框架 Primus
隨著機(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è)界有很多類(lèi)似的框架，如 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)單，不支持文件讀取等。
將算法工程師從此類(lèi)繁重的底層細(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ù)需求：
- 自研訓(xùn)練框架：目前除了業(yè)界開(kāi)源的 Tensorflow、Pytorch，為了滿(mǎn)足用戶(hù)的各種需求，字節(jié)也在機(jī)器學(xué)習(xí)場(chǎng)景進(jìn)行了深入探索，自研多個(gè)訓(xùn)練框架用于滿(mǎn)足用戶(hù)需求；
- 擁有底層資源調(diào)度系統(tǒng) YARN 和 Kubernetes：目前階段處在 YARN 向 Kubernetes 的遷移過(guò)程中，在此過(guò)程中需要對(duì)用戶(hù)屏蔽底層的調(diào)度系統(tǒng)，提供一致的體驗(yàn)；
- 大規(guī)模應(yīng)用混部資源：由于混部資源不穩(wěn)定的特點(diǎn)，對(duì)訓(xùn)練的容錯(cuò)和穩(wěn)定有著更高的要求；
- 支持復(fù)雜調(diào)度編排語(yǔ)義：為了使集群資源利用率最大化，需要將合適的容器放在適當(dāng)?shù)奈恢蒙希⑿枰軌騽?dòng)態(tài)調(diào)整并發(fā)和容器大小。
- 支持復(fù)雜數(shù)據(jù)源和數(shù)據(jù)調(diào)度需求：支持多種類(lèi)型數(shù)據(jù)源和數(shù)據(jù)類(lèi)型的混合訓(xùn)練，如 HDFS、Kafka 等批式流式數(shù)據(jù)源，Text、PB、Parquet 等數(shù)據(jù)類(lèi)型。
目前 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í)行階段。
- 提交階段
用戶(hù)需要描述整個(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ù)用戶(hù)的配置，向 YARN 或 Kubernetes 集群提交 Primus Application Master（之后簡(jiǎn)稱(chēng) AM）用于管理 Primus 作業(yè)集群的運(yùn)行。
- 執(zhí)行階段
當(dāng)作業(yè)提交到集群后，AM 會(huì)根據(jù)用戶(hù)的配置向 Resource Manager 申請(qǐng)若干 Executor，并上拉起對(duì)應(yīng)的角色，在持續(xù)監(jiān)控這些角色的健康狀態(tài)過(guò)程中，如果發(fā)現(xiàn)狀態(tài)異常，AM 則會(huì)根據(jù)用戶(hù)配置進(jìn)行相應(yīng)的操作保證訓(xùn)練正常運(yùn)行。
- 數(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)。
功能介紹
- 多訓(xùn)練框架支持：Tensorflow、PyTorch、Monolith 等；
- 多調(diào)度器支持：YARN、Kubernetes 等；
- 多角色支持：如 PS-Chief-CPU、Worker-GPU、Worker-Evaluator 等，并支持多角色之間的親和反親和等特殊調(diào)度策略；
- 多編排策略：支持同時(shí)啟動(dòng)，逐個(gè)啟動(dòng)，基于角色的按順序啟動(dòng)等（如先啟動(dòng) PS，再啟動(dòng) Worker）；
- 容錯(cuò)處理：Worker 失敗自動(dòng)拉起新 Worker，PS 失敗整體失??；
- 動(dòng)態(tài)調(diào)度：例如支持動(dòng)態(tài)擴(kuò)大縮小 Worker 數(shù)；
- 多數(shù)據(jù)源數(shù)據(jù)類(lèi)型支持：HDFS、Kafka 等；
- 數(shù)據(jù)負(fù)載均衡與狀態(tài)保存：支持按 Worker 負(fù)載動(dòng)態(tài)分配 Task，如在 Worker 失敗時(shí)支持回收 Task 并進(jìn)行重新分配；
- 多線(xiàn)程高速數(shù)據(jù)讀取：支持多線(xiàn)程讀取 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ī)劃
- 開(kāi)源 Primus 更多能力后續(xù)將陸續(xù)開(kāi)放，詳見(jiàn)：https://mp.weixin.qq.com/s/uGBy-WpdjTMUy-7MQAZiww
- 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
