基于阶段融合的大型语言模型 RLHF 训练优化
机器学习
2025-04-23 v3 计算与语言
分布式、并行与集群计算
摘要
我们提出 RLHFuse，一种针对强化学习从人类反馈 (RLHF) 训练的高效系统，实现阶段融合。由于 RLHF 训练的内在特性——生成阶段的数据偏斜以及训练阶段的管线气泡，现有 RLHF 系统面临 GPU 利用率低的问题。RLHFuse 打破传统将 RLHF 工作流程视为独立任务组合的观念，将每个任务划分为更细粒度的子任务，通过阶段融合提升 GPU 利用率。RLHFuse 的核心思想包括：首先，针对生成与推理任务，将其划分为样本级子任务，实现高效的跨阶段融合，重叠生成与推理阶段的执行，从而缓解由长尾样本主导的生成瓶颈；其次，针对训练任务，将其划分为微批次级子任务，在训练阶段进行内部融合，以融合管线方案并发执行这些子任务，有效缓解管线气泡问题。实验表明，RLHFuse 的训练吞吐量相较于现有系统提升最高可达 。
引用
@article{arxiv.2409.13221,
title = {Optimizing RLHF Training for Large Language Models with Stage Fusion},
author = {Yinmin Zhong and Zili Zhang and Bingyang Wu and Shengyu Liu and Yukun Chen and Changyi Wan and Hanpeng Hu and Lei Xia and Ranchen Ming and Yibo Zhu and Xin Jin},
journal= {arXiv preprint arXiv:2409.13221},
year = {2025}
}