项目地址:
https://github.com/bytedance/primus
随着机器学习的发展,模型及训练模型所需的数据量越来越大,也都趋向于通过分布式训练实现。而算法工程师通常需要对这些分布式框架涉及到的底层文件存储和调度系统有较深的理解,才能够快速批量开启模型训练,保证资源利用率。
目前业界有很多类似的框架,如 TonY、TensorFlowOnSpark,Kubeflow 中的 Training Operators 等,但这些框架或多或少存在某些问题,如与固定的机器学习框架( Tensorflow,Pytorch )耦合需要写明例如 PS、Worker 等角色,容错和弹性调度支持不友好,不支持异构调度,调度语义较为简单,不支持文件读取等。
将算法工程师从此类繁重的底层细节中解脱出来、更多地关注到算法层面,即为 Primus 解决的问题。
经过字节跳动在不断实践中调整打磨的 Primus,拥有以下能力支撑业务需求:
目前 Primus 流批一体训练框架在字节内部支持了抖音、头条、Tiktok、广告等大部分业务,每天运行在 Primus 上的作业总核数达到几百万。相比旧的基于Hadoop Streaming的框架,训练性能提升3倍,单次训练总数据量从TB级提升到了PB级,训练准备时间由几十分钟降低到秒级。
Primus 是一个通用的分布式训练调度框架,管理了机器学习训练框架(如 Tensorflow、Pytorch)的生命周期和数据分发,帮助训练框架获得更好的分布式能力。
架构介绍
Primus 整体架构
整个 Primus 生命周期分为两阶段,提交阶段和执行阶段。
用户需要描述整个任务的训练资源,数据输入以及容错策略。
Primus Client 通过根据用户的配置,向 YARN 或 Kubernetes 集群提交 Primus Application Master(之后简称 AM)用于管理 Primus 作业集群的运行。
当作业提交到集群后,AM 会根据用户的配置向 Resource Manager 申请若干 Executor,并上拉起对应的角色,在持续监控这些角色的健康状态过程中,如果发现状态异常,AM 则会根据用户配置进行相应的操作保证训练正常运行。
Primus 支持读取数据输出到训练器,能够实现数据的负载均衡并时刻记录数据的状态,降低长尾问题,在训练器有问题时也能够切换到正常的训练器继续进行训练。AM 通过扫描需要训练的数据并切分为 Task,可以将 Task 分发给 Executor,并与 Executor 通信记录 Task 的状态。
功能介绍
部署情况
Primus 支撑了字节跳动内部“推荐”“广告”“搜索”等场景,如头条推荐、抖音视频推荐、穿山甲广告、千川图文广告、抖音搜索等业务的超大规模深度学习训练,日均可达上万任务的训练,450W Core资源的使用。
目前,Primus 已在 Github 上开源,欢迎大家一同参与共建!
项目地址:https://github.com/bytedance/primus