Skip to content

shipbehaves/reward-model-ppo

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

6 Commits
 
 
 
 
 
 
 
 
 
 

Repository files navigation

reward-model-ppo

The classic RLHF loop: train an explicit reward model on preference pairs, then optimize a policy against it with online RL. DPO collapses these two steps into one; this does them the original way, so the reward model and the RL loop are visible and debuggable.

Method

  1. Reward model: a transformer with a scalar head, trained on (chosen outscores rejected) pairs (Bradley-Terry loss).
  2. Online RL: the policy generates responses, the reward model scores them, and the policy is updated to raise reward under a KL penalty to a frozen reference. PPO is the canonical algorithm; RLOO (REINFORCE leave-one-out) is the simpler, more stable variant TRL maintains at top level.

Plan (small steps)

  1. smoke (free, single 8 GB GPU, e.g. RTX 4070): RewardTrainer + the RL loop run end to end on a 0.5B model.
  2. reward model: train on a public preference set (UltraFeedback), report held-out accuracy.
  3. RL: optimize a small policy against the reward model, watching reward and KL together.
  4. eval + failure analysis: the canonical failure of this loop is reward hacking / over- optimization (reward climbs while real quality stalls and KL blows up). Characterize it, do not just report the reward number.

Result

Qwen2.5-1.5B, UltraFeedback. Two parts, and they tell different stories.

Reward model (the solid half): held-out accuracy 0.757 (n=1000) - it scores the preferred response higher 76% of the time. A clean, reusable artifact: yavuz-ai/qwen2.5-1.5b-rm-ultrafeedback.

PPO against it (the hard half): the loop runs, but reward did not cleanly improve.

episode KL to ref rlhf reward
8 12.0 -4.2
640 13.0 -2.4
1280 21.7 -5.7

The failure analysis (the substance)

PPO hit two walls. Memory: it holds policy + reference + reward model + value head plus rollout buffers, so it OOM'd on an 8 GB GPU entirely and OOM'd on an 80 GB A100 at batch 16 / response 128; it fit only at batch 8 / response 64 / rollout-forward 8 with capped prompts. Instability: in a short, lightly-tuned run (1,280 episodes, kl_coef 0.05) reward bounced (-2 to -6) with no clear trend while KL drifted up (12 -> 22). PPO is famously a multi-run tuning grind.

The point: DPO (used in the sibling projects) is a single, stable step that worked first time; building the classic two-step RM + PPO loop here surfaces exactly the cost and instability DPO was designed to avoid. Having done both makes that tradeoff concrete. Full teardown in NOTES.md.

Reproducibility

uv + uv.lock pin every dependency. Seeds fixed. Artifacts pushed to the Hub. Each step is a single documented command:

hf jobs uv run src/train_rm.py --flavor a100-large --secrets HF_TOKEN --env BASE_MODEL=Qwen/Qwen2.5-1.5B-Instruct --env OUTPUT_REPO=<rm>
hf jobs uv run src/ppo.py      --flavor a100-large --secrets HF_TOKEN --env REWARD_MODEL=<rm> --env BATCH=8 --env RESPONSE_LENGTH=64 --env OUTPUT_REPO=<policy>

Status

Complete: reward model trained (0.757 held-out) and the classic RM + PPO loop built and run, with the memory + instability difficulties characterized. Models and datasets on the Hub under yavuz-ai/.

About

Classic RLHF: train a reward model (0.757 held-out) then PPO a policy against it, with an honest teardown of PPO's memory + instability cost vs DPO.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages