通过阶段融合提升大模型强化学习训练效率
Optimizing RLHF Training for Large Language Models with Stage Fusion
- 将生成与训练任务细分为样本级和微批次子任务，实现跨阶段重叠执行
- 相比现有系统，训练吞吐量最高提升3.7倍，有效缓解长尾样本瓶颈
- 适合大规模语言模型强化学习训练优化，尤其关注算力利用率的团队
我们提出RLHFuse，一种针对基于人类反馈的强化学习（RLHF）的高效训练系统，采用阶段融合策略。由于RLHF训练中生成阶段的数据偏斜和训练阶段的流水线空洞问题，现有系统普遍存在GPU利用率低的问题。RLHFuse打破传统将RLHF流程视为独立任务组合的认知，将每个任务拆分为更细粒度的子任务，并通过阶段融合提升GPU利用率。其核心思想包括：对生成与推理任务，分解为样本级子任务，实现生成与推理阶段的高效重叠执行，缓解由长尾样本主导的生成瓶颈；对训练任务，将微批次作为子任务，采用融合流水线调度在训练阶段并发执行，有效消除流水线空洞。实验表明，相比现有系统，RLHFuse将训练吞吐量最高提升3.7倍。
原文摘要 · Abstract (English)
We present RLHFuse, an efficient training system with stage fusion for Reinforcement Learning from Human Feedback (RLHF). Due to the intrinsic nature of RLHF training, i.e., the data skewness in the generation stage and the pipeline bubbles in the training stage, existing RLHF systems suffer from low GPU utilization. RLHFuse breaks the traditional view of RLHF workflow as a composition of individual tasks, splitting each task into finer-grained subtasks, and performing stage fusion to improve GPU utilization. RLHFuse contains two key ideas. First, for generation and inference tasks, RLHFuse splits them into sample-level subtasks, enabling efficient inter-stage fusion to overlap the execution of generation and inference stages, thus mitigating the original generation bottleneck dominated by long-tailed samples. Second, for training tasks, RLHFuse breaks them into subtasks of micro-batches and performs intra-stage fusion to concurrently execute these subtasks in the training stage with a fused pipeline schedule, effectively mitigating the pipeline bubbles. The experiments show that RLHFuse increases the training throughput by up to $3.7\times$, compared to existing systems.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。