Reinforcement Learning from Human Feeedback (RLHF) aligns language models with human preferences through a multi-stage process. Modern large language model pipelines commonly incorporate RLHF, often combining techniques like Direct Preference Optimization (DPO) followed by Proximal Policy Optimization (PPO). This guide focuses on implementing the full RLHF pipeline, particular using PPO and RLOO algorithms.

The implementation is part of a lightweight RL framework available at: https://github.com/mst272/LLM-Dojo/tree/main/rlhf. It supports single- or multi-GPU training via DeepSpeed and is designed for clarity and extensibility, covering methods including PPO, RLOO, and SimPO.

The RLHF workflow consists of three core components:

• Supervised Fine-Tuned (SFT) Model: The base policy to be optimized.
• Reward Model: Evaluates responses based on human preference data.
• Reinforcement Learning Loop: Optimizes the policy using feedback from the reward model.

Reward Modeling

The reward model acts as a scorer for (instruction, response) pairs. Typical, it’s implemented as a sequence classification model that outputs a scalar reward.

For smaller models like GPT-2 or OPT-350M, one can directly use AutoModelForSequenceClassification. For larger architectures such as LLaMA or Qwen, Hugging Face provides dedicated classes like Qwen2ForSequenceClassification or Phi3ForSequenceClassification. These append a linear head (self.score) to the base transformer. The final token’s hidden state is used to produce a single reward per sequence:

from transformers import Phi3ForSequenceClassification

reward_model = Phi3ForSequenceClassification.from_pretrained(
    "microsoft/Phi-3-mini-4k-instruct",
    num_labels=1,
    trust_remote_code=True
)

This setup yields logits of shape [batch_size, seq_len, 1], which are reduced to [batch_size, 1] by selecting the last valid token’s output—commonly the end-of-sequence token.

Lightweight Custom Reward Models

To reduce computational cost, one can instantiate a smaller version of a large architecture by modifying its configuration:

from transformers import LlamaConfig, LlamaForSequenceClassification
import torch

torch.manual_seed(42)

# Define a minimal LLaMA-like config
config = LlamaConfig(
    vocab_size=100,
    hidden_size=256,
    intermediate_size=512,
    num_hidden_layers=4,
    num_attention_heads=4,
    max_position_embeddings=512
)

light_reward_model = LlamaForSequenceClassification(config, num_labels=1)

This approach allows experimentation with custom-sized reward models while retaining architectural consistency with the policy model.