Workflow:OpenRLHF OpenRLHF Rejection Sampling
| Knowledge Sources | |
|---|---|
| Domains | LLMs, RLHF, Rejection_Sampling, Alignment |
| Last Updated | 2026-02-07 10:00 GMT |
Overview
Iterative best-of-N training pipeline that generates multiple candidate responses, scores them with a reward model, selects the best, and trains on top-ranked responses.
Description
This workflow implements iterative rejection sampling (also called best-of-N training or iterative conditional SFT). In each iteration, the current policy generates multiple candidate responses per prompt using vLLM, a reward model scores all candidates, and the highest-scored response for each prompt is selected. The policy is then fine-tuned (SFT) on these best responses. The loop repeats for multiple iterations, progressively improving response quality. This approach avoids the complexity of PPO while still leveraging a reward model for alignment.
Usage
Execute this workflow when you have a trained reward model and want to iteratively improve your policy without the infrastructure complexity of PPO. Rejection sampling is simpler to implement and debug than RL-based methods, making it suitable for teams with limited distributed training infrastructure. It works well when the reward model is reliable and the base model already produces some high-quality responses.
Execution Steps
Step 1: Prepare initial models and data
Load the SFT-trained policy model and the reward model. Prepare the prompt dataset that will be used for generating candidate responses in each iteration. Set the iteration counter and output directory structure.
Key considerations:
- The policy model starts as the SFT checkpoint and is updated each iteration
- The reward model remains fixed across all iterations
- Prompts should be diverse and representative of the target domain
Step 2: Generate candidate responses
Using vLLM batch inference, generate N candidate responses per prompt (typically N=4-8) with temperature sampling. Each prompt produces multiple diverse responses through high temperature (e.g., 0.9) and top-p sampling.
Key considerations:
- Higher best_of_n values increase the quality ceiling but require more compute
- Temperature and top_p control diversity of candidates
- vLLM enables efficient batched generation across many prompts
Step 3: Score candidates with reward model
Run batch reward model inference on all generated responses. The reward model assigns a scalar quality score to each candidate. Apply the rejection sampling post-processor to select the highest-scored response for each prompt.
Key considerations:
- The rejection sampling post-processor (rs) selects the best response per prompt group
- Reward normalization can improve consistency of selection across prompts
- Output format matches the SFT training input format
Step 4: Train on best responses
Fine-tune the current policy model on the selected best responses using standard SFT training. This supervised step pushes the model toward generating higher-quality outputs as scored by the reward model.
Key considerations:
- Training uses the same SFT infrastructure with standard cross-entropy loss
- Typically 1 epoch of training per iteration to avoid overfitting
- Learning rate may be lower than initial SFT (e.g., 2e-6)
Step 5: Iterate or terminate
Check if the maximum number of iterations has been reached. If not, update the policy model checkpoint to the newly trained version and return to Step 2. If done, save the final model.
Key considerations:
- Typical iteration counts range from 5 to 10
- Monitor reward scores across iterations for improvement
- Diminishing returns typically appear after 5-8 iterations