System Design Lab

An RLHF pipeline is a training loop where rollout generation, reward scoring, and policy updates fight over the same GPUs.

Change the preference dataset, how many rollouts each PPO step generates, the policy and reward-model sizes, the PPO epochs per batch, the KL penalty, and how often you eval. The design moves from SFT-only, to adding a reward model, to a full PPO rollout loop, to scaling rollout generation (the usual bottleneck), to eval gates and a DPO alternative that drops the RL loop entirely.

Read the full tutorial

Guided walkthrough

Reason about it one step at a time

Normal evolution scenarios

Click left to right for the intended demo path. Each card changes the workload inputs.

Workload

These are inputs, not preset architecture stages.

Preference pairs Human comparison pairs (chosen vs rejected) used to train the reward model or, under DPO, the policy directly. Rollouts per step Completions the policy generates each PPO step; every one must be scored by the reward model and the reference model. Policy model size Parameters in the policy being optimized; it must run both generation (inference) and gradient updates. Reward model size Parameters in the reward model that scores each rollout; runs as an inference forward pass in the loop. PPO epochs per batch Optimization passes the policy takes over each batch of scored rollouts before generating again. KL penalty coefficient Strength of the KL constraint pulling the policy back toward the frozen reference model; too low and the policy reward-hacks. Eval frequency How often a held-out eval and safety check runs, measured in steps between evals (lower means more frequent). Use DPO (no RL loop) Direct Preference Optimization trains the policy straight from preference pairs, removing the rollout-generation loop and online reward model. Separate reward model Host the reward model on its own devices instead of co-locating it with the policy; isolates scoring throughput.

Recommended shape

Current architecture path

Data

Demonstrations human-written prompts and ideal responses used for supervised fine-tuning

Preference data human comparison pairs that define which responses are better

SFT

SFT trainer fine-tunes the base model on demonstrations to seed the policy

Reward model

RM trainer trains a reward model to score responses from preference pairs

Reward model serves a scalar reward for each generated rollout in the loop

RL loop

Rollout gen the policy generates many completions per step; usually the bottleneck

Policy update gradient step that pushes the policy toward higher reward

Reference model frozen reference whose KL divergence keeps the policy from drifting

Eval + registry

Eval harness runs held-out and safety evals to gate checkpoints

Model registry stores and versions promoted policy checkpoints

Metrics log collects reward, KL, and rollout metrics off the training path

Bottlenecks

Rollout generation load

Reward + reference scoring

Policy training memory

KL drift risk

Eval / safety coverage demand

Why this changes

Decision tradeoffs

SFT vs preference data

Reward model vs DPO

Rollout-generation throughput

Distributed policy training

KL / reference control

Eval + safety gating

Source-backed rules

These are the durable system-design claims behind the model. The exact slider thresholds are deliberately labeled as teaching assumptions.

Verified rule

InstructGPT: SFT, then a reward model, then PPO against a KL-penalized reference

The canonical RLHF recipe fine-tunes on demonstrations, trains a reward model from comparisons, and optimizes the policy with PPO using a per-token KL penalty against the SFT reference.

Ouyang et al., 2022 (InstructGPT)

Verified rule

DPO removes the explicit reward model and RL rollout loop

Direct Preference Optimization shows the RLHF objective can be optimized with a simple classification loss on preference pairs, eliminating online sampling and a separate reward model.

Rafailov et al., 2023 (DPO)

Verified rule

PPO is the clipped policy-gradient algorithm used in the RL stage

Proximal Policy Optimization takes several optimization epochs over each batch of sampled trajectories with a clipped objective, which is why rollouts are reused across PPO epochs.

Schulman et al., 2017 (PPO)

Verified rule

Online RLHF couples generation and training; generation throughput dominates

High-throughput batched inference engines are used to generate rollouts because, in the online RL loop, sampling completions from the policy is the throughput bottleneck.

vLLM docs

Teaching assumptions

Single-GPU rollout-generation throughput and memory budgets are conservative teaching numbers, not vendor limits.
Reward and reference scoring cost is modeled as one inference forward pass per rollout each; PPO reuses each rollout for the chosen number of epochs.
KL drift risk is approximated from the KL coefficient and rollout volume; real runs also track measured KL and reward-hacking signals directly.