Lesson 12: RLHF — Reinforcement Learning from Human Feedback

Reinforcement Learning: From Basics to Advanced

Part 4: RLHF, LLM Alignment & Production

xdev.asia

Introduction

RLHF (Reinforcement Learning from Human Feedback) is the technique that turned GPT-3 into ChatGPT — aligning LLM with human preferences through reward model and PPO.

---

1. RLHF Pipeline — 3 Steps

Step 1: Supervised Fine-Tuning (SFT)

from transformers import AutoModelForCausalLM, TrainingArguments
from trl import SFTTrainer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B")

trainer = SFTTrainer(
    model=model,
    train_dataset=demo_dataset,  # (prompt, response) pairs
    args=TrainingArguments(
        output_dir="./sft_model",
        num_train_epochs=3,
        per_device_train_batch_size=4,
        learning_rate=2e-5,
    ),
)
trainer.train()

Step 2: Reward Model Training

from trl import RewardTrainer, RewardConfig

reward_model = AutoModelForSequenceClassification.from_pretrained(
    "meta-llama/Llama-3.1-8B", num_labels=1
)

trainer = RewardTrainer(
    model=reward_model,
    train_dataset=preference_dataset,
    # Format: {prompt, chosen_response, rejected_response}
    args=RewardConfig(
        output_dir="./reward_model",
        per_device_train_batch_size=4,
        num_train_epochs=1,
    ),
)
trainer.train()

Step 3: PPO Fine-Tuning

from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

model = AutoModelForCausalLMWithValueHead.from_pretrained("./sft_model")
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained("./sft_model")

config = PPOConfig(
    batch_size=16,
    learning_rate=1e-5,
    ppo_epochs=4,
    mini_batch_size=4,
)

trainer = PPOTrainer(config, model, ref_model, tokenizer)

for batch in dataloader:
    queries = batch["query"]
    responses = model.generate(queries)
    rewards = reward_model(queries, responses)
    
    # KL penalty to prevent reward hacking
    trainer.step(queries, responses, rewards)

---

2. Reward Hacking & KL Divergence

Total Reward = RM_score(response) - β * KL(π || π_ref)

KL penalty keeps model near SFT model → avoid reward hacking.

---

3. Constitutional AI (Anthropic)

1. Generate responses
2. Ask model to critique based on principles
3. Ask model to revise
4. Train on revised responses (RLAIF)

---

Summary

Step	Input	Output	Purpose
SFT	Demonstrations	Fine-tuned LLM	Learn format
Reward Model	Human preferences	Reward scorer	Learn preferences
PPO	RM rewards	Aligned LLM	Optimize for preferences