字节跳动正式开源分布式训练调度框架 Primus-优快云博客

动手点关注

干货不迷路

项目地址：https://github.com/bytedance/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