Step 1: Load the base model and capture baseline outputs
Pretraining teaches a model capabilities — the ability to produce grammatical text that's loosely correlated with the prompt. Alignment teaches a model what kind of response is desired — helpful, honest, harmless, concise, faithful to your company's policy, whatever you define.
Every frontier model you've heard of — GPT-4, Claude, Gemini, Llama 3 — went through an alignment pipeline. Before alignment, a raw LLM might respond to "What's the capital of France?" with "What's the capital of France? What's the capital of Germany? What's the capital of..." (continuing the pattern), or with a 500-word essay, or with an irrelevant tangent. After alignment, it says "Paris."
The RLHF recipe (OpenAI, 2017–2022)
The method that made ChatGPT possible, introduced across Christiano et al. 2017, Ziegler et al. 2019, and OpenAI's InstructGPT paper (Ouyang et al., 2022). It has three stages:
- Supervised fine-tuning (SFT) — train the base model on high-quality demonstrations
- Reward model training — humans rank pairs of outputs; train a separate model to predict the ranking
- RL (PPO) fine-tuning — use the reward model to fine-tune the SFT model via reinforcement learning, with a KL penalty to prevent drift
It works, but it is famously fiddly. You're juggling three models (base, SFT, reward, policy), unstable RL, reward hacking, and a full PPO implementation. Reproducing InstructGPT-scale RLHF was an extreme engineering effort.
DPO: the 2023 simplification
Rafailov et al., 2023 — Direct Preference Optimization: Your Language Model Is Secretly a Reward Model. The insight hiding in the title: the closed-form mapping between reward functions and optimal policies lets you skip the reward model entirely. You can train directly on preference pairs (chosen vs rejected) with a simple classification-style loss.
DPO replaces the full RLHF stack with one training loop, one model being updated, one frozen reference model, one contrastive loss. Published results match or beat PPO-RLHF on summarization and single-turn dialogue. It's now the default alignment algorithm in open-source — Zephyr, Tulu, Orca, Mistral's Instruct variants, Qwen's chat variants all use DPO or close derivatives (IPO, KTO, SimPO).
What you'll do in this lab
- Load GPT-2 (base model) and generate baseline outputs on a test set of prompts.
- Build a preference dataset — 10 prompts with (chosen, rejected) response pairs encoding a concrete alignment preference.
- Run DPO with
trl.DPOTraineron the preference data — a few hundred steps, two minutes on a 3060 Ti. - Regenerate on the same prompts with the aligned model and observe the shift.
We'll use GPT-2 because it's small (124M params, <1s per step), and the alignment signal is visible even with this tiny preference set. The pipeline you build here scales unchanged to Llama-70B.
Read while this runs
- DPO paper (Rafailov et al., 2023) — sections 4 and 5 derive the loss; the math is surprisingly clean.
- Tyler Romero — DPO Explained In-Depth — best single blog post covering the intuition and the gradient.
- TRL docs — DPOTrainer — what we're driving here.
We pick three prompts that expose classic base-LM failure modes:
- Completion: a model trained on raw text often just continues a question rather than answering it
- Conciseness: base LMs ramble; aligned ones stay terse
- Directness: base LMs over-hedge
After DPO, if our preference dataset rewards concise direct answers, we should see movement on all three.
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
MODEL_ID = 'gpt2'
base_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if base_tokenizer.pad_token is None:
base_tokenizer.pad_token = base_tokenizer.eos_token
base_model = AutoModelForCausalLM.from_pretrained(MODEL_ID).to('cuda')
base_model.eval()
print(f'Loaded {MODEL_ID}: {base_model.num_parameters()/1e6:.0f}M params')
# Keep a frozen copy of the base model — DPO needs the reference model to stay put
# (we'll use this directly via TRL's ref_model=None default, which auto-copies)
test_prompts = [
'Q: What is the capital of France?\nA:',
'Q: How do I boil an egg?\nA:',
'Q: Explain gravity in one sentence.\nA:',
]
def generate(model, tok, prompt, max_new=25):
ids = tok(prompt, return_tensors='pt').input_ids.to('cuda')
with torch.no_grad():
out = model.generate(ids, max_new_tokens=max_new, do_sample=False, pad_token_id=tok.eos_token_id)
return tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True).strip()
baseline_outputs = {p: generate(base_model, base_tokenizer, p) for p in test_prompts}
print('\nBASELINE OUTPUTS (before alignment):')
for p, t in baseline_outputs.items():
print(f'\n PROMPT: {p.strip()}')
print(f' BASE: {t!r}')