🔥 简介:深度推荐模型统一训练系统Primus

搜索、广告、推荐是互联网业务的重大组成部分。在模型日益庞大、用户基数不断增长的背景下，如何高效完成数万亿样本的训练，已成为业界公认的技术挑战。字节跳动基础架构-基础技术-训练设施团队联合AML-Engine团队，与浙江大学、上海交通大学合作，并在搜索、广告、推荐等团队的大力支持下，设计并开发了一个面向大规模深度学习推荐模型（Deep Learning Recommendation Model, DLRM）的全自研统一训练系统Primus，相关成果已被计算机系统领域顶级会议 USENIX ATC录用 ，标志着 Primus 在工业界落地和学术界影响上的双重认可。

论文标题：Primus: Unified Training System for Large-Scale Deep Learning Recommendation Models

论文作者：Jixi Shan, Xiuqi Huang, Yang Guo, Hongyue Mao, Ho-Pang Hsu, Hang Cheng, Can Wang, Jun Song, Rui Shi, Xiaofeng Gao, Jingwei Xu, Shiru Ren, Jiaxiao Zheng, Hua Huang, Lele Yu, Peng Xu, Guihai Chen

论文地址：
https://www.usenix.org/conference/atc25/presentation/shan-jixi

📌 背景:大规模 DLRM 训练的三重困境

如图所示，过去5年，随着字节跳动内部推荐业务的迅猛发展，单个DLRM的日均训练数据量从最初的 20TB 快速增长到 160TB，CPU虚拟核心数使用总量也从 150 万核扩充到 900 万核，涉及的 GPU 使用量更是达到数万张。

图：过去 5 年 DLRM 训练规模增长趋势

现有的训练系统面临三大痛点：

• 资源调度割裂：
  大规模部署一般跨越由不同调度系统管理的异构资源池，弹性训练策略缺乏跨环境操作的能力，限制了资源的高效利用。
• 数据编排复杂：
  海量样本分散在各种存储系统中且具有多样化的格式，读取与预处理多源批式和流式数据容易形成瓶颈。
• 模型时效性不足：

  用户行为和偏好的不断演变，单一的离线、在线、或先离线后在线的训练模式难以平衡“长期记忆”与“实时适应”。

🚀 创新:统一的资源调度/数据编排/训练范式

Primus 针对上述三大痛点，提出了以统一资源调度、统一数据编排、统一训练范式为设计核心的训练系统架构。

图：Primus系统整体框架

🌟 01 | 统一资源调度

图：Primus系统跨环境多策略资源调度

Primus 通过统一资源控制器和动态扩缩容管理器进行资源整合。

统一资源控制器引入统一资源抽象技术，即用户资源定义（Custom Resource Definitions, CRD）标准化 DLRM 训练作业和异构资源，提供跨调度系统如YARN和Kubernetes的统一API，在多集群资源间实现多策略动态扩缩容。

动态扩缩容管理器根据实时指标选择合适策略，水平扩缩支持在大规模集群场景下实现动态并行，自动调整数据执行器与训练执行器的比例，垂直扩缩支持调整单个执行器的 CPU 和内存资源，提高资源利用率。

🌟02 | 统一数据编排

图：Primus系统多源批流数据编排

Primus 通过抽象数据定义和获取逻辑，优化数据编排速度。

Primus 提出统一的三层数据定义，包括数据集、数据流和数据源层，有效管理 DLRM 的训练数据集。通过这必定义，进一步设计了数据任务图生成（Data Task Graph Generation, DTGG）机制以进行数据编排，提高多源批式和流式数据的处理效率，适应训练作业的多样化数据需求。

数据任务图生成（DTGG）技术包含 Timer 、Data Source、Joiner和Sink四类算子，数据驱动器通过DTGG快速按时间顺序跨数据源生成数据任务，并由数据执行器运行任务提供数据，洗牌机制也做了优化以更好解决数据处理的瓶颈。

🌟 03 | 统一训练范式

图：Primus系统在离线混合训练范式

Primus 提供混合训练推荐模型框架（Mixture Training Recommendation Model, MTRM），系统性支持记忆塔（Memory Tower） + 适应塔（Adapt Tower）的双塔结构，通过不同的计算图控制批式数据和流式数据的学习顺序，对两个塔分别进行参数更新，从而既保持对历史分布的记忆，又能快速适应新趋势。

同时，也引入了混合数据优先级机制，使用缓冲区长度和重大性计算并提供行级优先级控制，避免流式数据激增无法及时处理导致的模型效果下降。

✅ 成效:支撑字节搜广推核心服务

自 2019 年投入生产以来，Primus 已稳定支撑字节跳动内部抖音、今日头条、西瓜视频等业务，带动搜广推收益持续增长。

在实际DLRM训练作业上的充分评估也论证了Primus系统承担DLRM训练的卓越能力。

图：不用 (w/o) 和用 (w/) MTRM 的各种模型的 AUC 和收入表现比较

训练设施团队为字节内部大规模机器学习包括大模型场景提供计算/编排/数据管理的基础框架以及解决方案，包括Primus/Ray/Magnus等，欢迎感兴趣同学加入我们！联系方式：wuyixin.yx@bytedance.com。