在工业级推荐模型训练中,Embedding Table 往往包含数十亿甚至数百亿个 Key,参数规模可达到数十 TB。如此庞大的模型状态使 Checkpoint 成为分布式训练中的重要性能瓶颈:全量保存会带来巨大的存储 I/O 开销,并迫使训练暂停以保证一致性;而现有的增量或差分 Checkpoint 方案虽能减少单次保存开销,却又分别面临恢复链过长或检查点持续膨胀的问题。因此,如何在训练吞吐、存储开销与故障恢复效率之间取得平衡,成为工业级推荐模型容错训练中的关键挑战。
一、研究背景:工业级推荐模型的 Checkpoint 困境
工业级深度推荐模型通常包含规模巨大的 Embedding Table,参数量可达到数十 TB。传统全量 Checkpoint 需要周期性保存整个 Embedding Table,不仅带来巨大的存储 I/O 开销,还可能因数据一致性要求而阻塞训练。与此同时,Embedding 的访问频率具有明显的幂律分布,约 90% 的访问仅集中在约 0.1% 的高频 Embedding 上。
为降低全量保存开销,现有方法主要采用增量 Checkpoint 和差分 Checkpoint。增量方式单次保存量较小,但恢复时需要依次加载多个检查点;差分方式恢复更直接,但保存的数据量会随训练持续增长。
更重要的是,我们发现这两种方案都没有充分利用 Embedding 的冷热差异。差分 Checkpoint 对 Cold Embedding 容易产生重复保存,而增量 Checkpoint 对 Hot Embedding 则会带来较高的恢复开销,因此单一策略难以同时兼顾保存效率与恢复速度。
二、方法简介:AdaptCP 的三大核心模块
针对以上问题,AdaptCP 从 Checkpoint 策略、训练流水线和底层 I/O 三个层面进行了系统性优化:
冷热感知的自适应混合 Checkpoint 策略
SoU 更新时保存机制
基于 SSD 缓冲的并行 I/O 引擎
三者分别解决:存哪些参数、什么时候存,以及怎么更快地存。
2.1 冷热感知的自适应混合 Checkpoint
AdaptCP 的第一个核心创新,是根据 Embedding 的访问热度,为不同参数动态选择更合适的 Checkpoint 策略。
对于频繁访问和更新的 Hot Embedding,采用增量 Checkpoint 时,同一批高频参数会分散在多个增量文件中,故障恢复需要依次加载,容易带来较高的恢复开销,因此更适合采用差分 Checkpoint。相反,对于访问和更新频率较低的 Cold Embedding,采用差分 Checkpoint 时,已发生变化的参数会在后续检查点中被重复保存,造成额外的写入和存储开销,因此更适合采用增量 Checkpoint。
基于这一观察,AdaptCP 对 Hot Embedding 采用差分 Checkpoint,对 Cold Embedding 采用增量 Checkpoint,从而兼顾保存效率与恢复速度。
论文进一步建立了统一的 I/O 成本模型,并从理论上证明:随着 Embedding 访问频率变化,系统中存在唯一的最优 Hot/Cold 划分阈值,可使 Checkpoint 保存与故障恢复的总体 I/O 成本达到最低,且该阈值可通过二分法快速求解。该理论推导为在线训练过程中动态选择冷热参数边界提供了明确依据。
在线训练过程中,Embedding 的访问分布并非固定不变。随着用户行为持续变化,原有的冷热关系可能发生迁移,固定阈值难以长期保持最优。 因此,AdaptCP 进一步设计了在线自适应冷热划分机制,由三个组件协同完成:
Tracker:记录当前周期内发生更新的 Embedding;
Access Counter:实时统计不同 Embedding 的访问频率;
Threshold Solver:根据当前访问分布在线求解最优冷热阈值。
系统随后按照阈值将 Embedding 动态划分为 Hot Set 和 Cold Set,并分别采用差分 Checkpoint 和增量 Checkpoint,从而使 Checkpoint 策略能够随运行时访问模式变化持续调整。
为此,AdaptCP 提出了 SoU 机制,其设计思想类似于操作系统中的 Copy-on-Write。由于推荐模型每个 Batch 实际只会更新少量 Embedding,SoU 无需在 Checkpoint 触发时立即对全部参数进行同步保存,而是在某个 Embedding 即将被更新时,先保留其更新前的旧值,再继续执行参数更新。未发生变化的参数则可以继续由训练过程和 Checkpoint 共享,从而将一次大规模同步保存拆解为细粒度的异步保存任务。
在 SoU 基础上,AdaptCP 进一步结合 Embedding Prefetch 设计了 Prefetch-triggered SoU。系统在参数被预取后,就提前启动高优先级的保存任务,从而利用 Prefetch 到参数更新之间的计算窗口完成后台 Checkpoint I/O。等参数真正进入更新阶段时,大部分保存任务已经完成,可以进一步减少训练线程的等待时间,将 Checkpoint 延迟尽可能隐藏在正常训练过程之中。
三、实验结果
为了全面验证 AdaptCP 的有效性,团队分别在公开数据集和真实工业训练环境中进行了实验。
3.1 整体性能
在训练吞吐方面,相比 Check-N-Run,AdaptCP 在 Criteo Terabyte 场景下实现 4.59× 提升,在真实生产数据场景下进一步达到 13.87×。随着模型和 Checkpoint 规模进一步扩大,AdaptCP 的性能优势更加明显,说明其对工业级大规模 Checkpoint 场景具有更好的扩展性。
在故障恢复方面,AdaptCP 同样显著降低了 Checkpoint 加载时间。在生产数据集上,Check-N-Run、Incr CKPT-N 和原生增量 Checkpoint 的恢复时间分别是 AdaptCP 的 9.41×、12.70× 和 13.62×。这一优势主要来自并行恢复机制,可以同时从远端存储加载多个 Checkpoint 分片,充分利用读取带宽。
3.3 高频 Checkpoint 下依然保持高吞吐
Checkpoint 频率直接影响系统的容错能力:保存越频繁,发生故障时需要回退和重算的数据越少,但也会带来更高的 I/O 压力。实验将 Checkpoint 间隔从每 10,000 Batch 逐步缩短至每 100 Batch,以测试不同方案在高频保存场景下的表现。
当 Checkpoint 提高到每 100 Batch 一次时,Check-N-Run 和原生增量 Checkpoint 的训练吞吐分别下降 92.6% 和 74.9%,而 AdaptCP 仅下降 19.0%,并在所有 Checkpoint 频率下始终保持最高吞吐。其原因在于 SoU 能够将保存过程与计算重叠,而并行 Writer Engine 又能够在有限时间窗口内完成大部分数据持久化,使高频 Checkpoint 对训练的影响显著降低。
四、总结与展望
AdaptCP 的价值在于将推荐模型自身的参数访问特征、训练过程与底层存储系统结合起来,重新设计大规模推荐模型训练中的 Checkpoint 机制。相比将 Checkpoint 作为独立于训练之外的周期性保存任务,AdaptCP 根据参数热度动态选择保存策略,并通过 Save-on-Update 与并行 I/O 将保存过程尽可能融入正常训练流水线,为 TB 级推荐模型提供了一种兼顾训练吞吐与故障恢复效率的系统方案。
面向正在快速演进的推荐大模型,Checkpoint 系统仍将面临新的挑战。随着模型规模持续扩大,除海量 Embedding 参数外,Dense 部分也在不断 Scaling,模型状态、优化器状态以及训练集群规模都将进一步增长,对训练容错和系统稳定性提出更高要求。未来,团队将继续围绕推荐大模型的高效、稳定训练展开探索,在更大参数规模、更复杂训练架构和更高容错要求下,进一步降低 Checkpoint 对训练过程的干扰,并提升故障恢复效率,为推荐大模型持续 Scaling 提供更加可靠的基础设施支撑。
五、团队介绍
快手算法引擎部 AI Platform 团队长期深耕大规模模型训练基础设施建设,支撑搜广推等核心业务模型及自研多模态大模型的高效训练,并持续优化大规模训推资源池的调度能力。团队先后支撑了 OneRec 端到端生成式推荐大模型、KeyeVL 多模态大模型的发布上线与全量推广,并推动相关模型取得了较为突出的业务收益。
团队在多个技术方向发力,在推荐训练方面,率先打造了业界先进针对大规模生成式推荐模型的超算软硬一体化框架 SKAI,将 MFU 提升至 30% 以上;在训推一体化资源调度平台上,团队通过多种形式的混部、潮汐、统一调度等多种业界先进手段持续提升资源利用率。近 3 年来,团队在 EuroSys、SoCC、ASE 等国际顶级学术会议发表多篇论文,持续推动大规模训练系统在工业场景中的研究与落地。
六、招聘职位
1、大模型训练引擎研发工程师
职位描述
深度参与多模态/大语言模型训练全链路开发,包括数据、预训练、后训练全流程优化;设计和优化分布式训练框架,通过混合并行,通信计算 overlap、低精度训练等方法解决超长序列、超大规模 moe 场景下的训练效率问题;
参与通用高性能 RL 框架的开发和优化;
算法工程 co-design,探索最优的训练范式。
任职要求
熟悉 C++/Python 编程,有较好的编程风格和代码管理能力;
有优秀的逻辑分析能力,能够对业务逻辑进行合理的抽象和拆分;
有大模型算法领域知识者优先;
有开源分布式训练框架,PyTorch/Megatron/Verl 等经验优先;
有设计和优化大语言模型预训练、微调、强化学习流程经验者优先;
有大模型训练基础设施相关从业经验者优先。
职位描述
深度参与快手自研生成式推荐大模型(OneRec)训练全链路开发和优化,以及快手广告、电商、直播、搜索等全域模型的训练全链路研发与优化;
设计和优化分布式训练框架,通过算子优化、通信计算 overlap、低精度训练等方法解决超长序列、超大规模 moe 场景下的训练效率问题;
算法工程 co-design,探索最优的训练范式。
任职要求
熟悉 C++ / Python 编程,有较好的编程风格和代码管理能力;
有优秀的逻辑分析能力,有较好的数学基础;
至少对 pytorch / tensorflow 等训练框架有较丰富的使用经验和二次开发经验;
有大模型算法领域知识者优先;
有大模型训练 Infra 相关从业经验者优先。
投递方式:扫描下方二维码,投递相关岗位。
【推荐阅读】