RLHFSpec: Breaking the Efficiency Bottleneck in RLHF Training via Adaptive Drafting

2025-12-04

AI Summary

This work addresses the low generation efficiency during RLHF training, which constitutes a critical bottleneck in end-to-end workflows. To tackle this, we propose RLHFSpec—a novel system that introduces speculative decoding into the RLHF generation pipeline for the first time. Its core innovations include a workload-aware dynamic drafting strategy selection mechanism, coupled with validation cost modeling and adaptive sample redistribution, enabling efficient GPU resource scheduling and load balancing. Compared to state-of-the-art methods, RLHFSpec significantly improves throughput in the generation phase, alleviating the primary bottleneck and accelerating end-to-end RLHF training by 1.8–2.3×, while preserving policy optimization quality without degradation.

Technology Category

• Natural Language Processing: Generation
• Machine Learning: Hardware-aware ML
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Reinforcement Learning from Human Feedback (RLHF) is an important fine-tuning technique for large language models (LLMs) and comprises three stages: generation, inference, and training. The generation stage generates samples that are then used to infer learnable experiences for training. We observe that the generation stage is the bottleneck of the entire execution process and consider it a key point for optimization. Specifically, we realize the first attempt to integrate speculative decoding into the RLHF generation stage and propose RLHFSpec, an RLHF system that accelerates generation execution with adaptive speculative decoding and sample reallocation. To fully exploit the performance potential provided by speculative decoding, especially dealing with the dynamic workload of the generation stage, RLHFSpec proposes a workload-aware drafting strategy selection mechanism, which selects the near-optimal strategy by jointly considering the verification cost and the number of accepted tokens. Moreover, RLHFSpec also proposes sample reallocation to fully utilize the GPU resources, and optimizes it with an efficient sample migration mechanism. The experimental results show that the RLHFSpec can achieve higher throughput in the generation stage compared to state-of-the-art works. Moreover, due to the effective alleviation of the generation bottleneck, RLHFSpec also shows significant performance speedup in the entire RLHF execution.

Problem

• Optimizes RLHF generation stage efficiency via adaptive speculative decoding.
• Addresses dynamic workload with workload-aware drafting strategy selection.
• Enhances GPU resource utilization through sample reallocation and migration.

Innovation

• Integrates speculative decoding into RLHF generation stage
• Uses workload-aware drafting strategy for optimal token acceptance
• Implements sample reallocation with efficient migration for GPU utilization