hands-on labon a real NVIDIA GPU

RLHF & DPO Alignment

Run real Direct Preference Optimization on a small language model with TRL's DPOTrainer.

time
55 min
checked steps
4
level
Advanced
setup
None
NCP-GENLNCP-ADSNCA-GENLNCA-GENMPart of the AI Engineer Course

$29.99a month with Pro

Pro opens this lab and every other one. Cancel any time.

Lab cockpit55 min · 4 stepsSession running
3 / 4 steps passingBefore vs after: did alignment work? · step 4 of 4
lab.ipynb

Step 4 — Before vs after: did alignment work?

The moment of truth. We regenerate on the exact same three prompts we used in Step 1, using the DPO-trained model, and compare side by side.

[ ]:
aligned_outputs = {p: generate(aligned_model, base_tokenizer, p) for p in test_prompts} print('=' * 90)print('BEFORE vs AFTER DPO ALIGNMENT')print('=' * 90)for p in test_prompts:    print(f'\n PROMPT: {p.strip()}')    print(f'  BEFORE: {baseline_outputs[p]!r}')    print(f'  AFTER : {aligned_outputs[p]!r}')    delta_len = len(aligned_outputs[p]) - len(baseline_outputs[p])    print(f'  length delta: {delta_len:+d} chars')
[ ]:
lab.check(4)
step 4 of 4, replayed: Before vs after: did alignment work?. The code, the output and the check are the lab's own.

4 steps, each checked when you finish it

A check runs your work at the end of every step. Hints and the full solution are there if you get stuck.

  1. 1

    Load the base model and capture baseline outputs

    read this step below ↓

  2. 2

    Build the preference dataset

  3. 3

    Train with DPOTrainer

  4. 4

    Before vs after: did alignment work?

Step 1 as it appears in the lab

The lab’s own text. The hint and the solution stay inside the lab.

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:

  1. Supervised fine-tuning (SFT) — train the base model on high-quality demonstrations
  2. Reward model training — humans rank pairs of outputs; train a separate model to predict the ranking
  3. 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

  1. Load GPT-2 (base model) and generate baseline outputs on a test set of prompts.
  2. Build a preference dataset — 10 prompts with (chosen, rejected) response pairs encoding a concrete alignment preference.
  3. Run DPO with trl.DPOTrainer on the preference data — a few hundred steps, two minutes on a 3060 Ti.
  4. 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


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}')

what you get

  • A JupyterLab notebook on a real NVIDIA GPU
  • A check at the end of every step
  • Hints and the full solution
  • Nothing to install or configure

before you start

  • Fluent with Hugging Face transformers (tokenizer + causal LM generation)
  • Understanding of supervised fine-tuning loss
  • Basic grasp of what RLHF is trying to accomplish vs pretraining

Exam domains covered

AlignmentPrompt EngineeringLLM Fundamentals & ArchitectureExperimentation

Part of the LLM training and fine-tuning track

Lab 4 of 14. LoRA and QLoRA fine-tuning, DPO alignment, and a tokenizer and tiny GPT built from scratch.

Questions about this lab

How is DPO mathematically different from PPO-based RLHF?

PPO-RLHF needs three moving parts: a reward model trained on preferences, a policy being optimised, and a reference policy with a KL penalty. DPO rewrites the optimisation so the reward model disappears: the Bradley-Terry likelihood of preferring chosen over rejected factors into a closed-form loss over log-ratios of the policy vs the reference, scaled by a temperature β. You still have a reference model (ref_model in DPOConfig), but you skip reward-model training entirely and use plain gradient descent, no rollouts, no value network, no importance sampling. The tradeoff is that DPO is less expressive when the reward landscape is complex but dramatically simpler and more stable for preference-shaping.

Why does the check require only one changed output, not all of them?

Because with eight preference pairs and a handful of training steps, the alignment signal is sparse. Some prompts are far enough from the training distribution that neither the chosen nor rejected gradient meaningfully touches them, the model's output stays numerically close. At least one changing proves the DPO loss actually flowed through backward and updated weights; that's the floor the check enforces. In a real alignment run you'd instead evaluate on held-out preference pairs and report win-rate against the reference policy as your metric, not a binary "did anything change".

Is eight preference pairs enough to 'align' a model?

No, it's enough to verify the pipeline runs end-to-end and produce measurable behavior change on your specific prompts. Production DPO uses thousands to millions of pairs: UltraFeedback is ~64K, HH-RLHF is ~161K, Nectar crosses a million. Anthropic, OpenAI, and Meta all layer multiple preference datasets, iterative DPO rounds, and frequently blend SFT data in. The lab's eight-pair setup is deliberately the smallest thing that still lets you see DPO's mechanics; the reflection asks you to reason about what breaks when you scale it up.

Why use GPT-2 instead of a bigger model?

Because this lab is about DPO mechanics, not about producing a useful aligned assistant. GPT-2 at 124M parameters trains fast enough to converge on eight preference pairs in under a minute, leaves plenty of VRAM for the reference model to sit alongside, and makes the before/after shift in baseline_outputs vs aligned_outputs easy to eyeball. The exact same DPOTrainer call scales up to Llama 3 or Mistral 7B, you'd just add LoRA or QLoRA (DPO-on-LoRA is standard for 7B-class alignment on a single 24 GB card).

What's `alignment tax' and does this lab exhibit it?

Alignment tax is the capability drop that often follows preference optimisation, the model gets better at the preferred behaviour and worse at general reasoning, factual recall, or code generation. It happens because the preference gradient pulls the policy toward a narrow distribution defined by your labelers, away from the richer pretraining distribution. You can catch a miniature version of it in Step 4 by picking one baseline prompt far from your preference pairs and comparing the text, often slightly less coherent after training. Mitigations in production: KL regularisation (what β controls), mixing in SFT loss, and iterative rounds that re-anchor to the reference.

When would I reach for KTO, IPO, or ORPO instead of DPO?

KTO (Kahneman-Tversky Optimization) takes per-sample thumbs-up / thumbs-down labels instead of paired preferences, useful when you have binary feedback from production traffic rather than ranked pairs. IPO (Identity Preference Optimization) replaces DPO's log-sigmoid with a squared loss to reduce overfitting on small datasets. ORPO (Odds Ratio Preference Optimization) folds SFT and preference optimisation into a single loss, skipping the reference model entirely. Pick by what your data looks like: pairs → DPO or IPO; thumbs → KTO; no SFT-then-align split available → ORPO. DPO remains the default because the library support is strongest and the hyperparameters are best understood.

What you'll build in this DPO alignment lab

DPO replaced PPO-based RLHF as the default alignment recipe across the open-weight ecosystem, Zephyr, Tulu, Llama 3 Instruct's preference stages, and most community instruction-tunes all shipped on trl.DPOTrainer in the last two years, because it throws out the reward model entirely and still lands measurably aligned behavior. In about 55 minutes on a real NVIDIA GPU we provision, you'll run a real DPO pass on a causal LM, build a preference dataset from scratch, and produce a measurable before/after shift, not a fake "loss went down" signal but actual text that differs on held prompts. You'll walk away understanding why DPO is simpler than PPO (no rollouts, no value network, no reward model), why it's more fragile in a specific way (uniquely sensitive to preference-data quality), and how to read the trainer.state.log_history to tell the difference between a healthy run and alignment theater.

Technically you'll package {prompt, chosen, rejected} triplets into a datasets.Dataset, instantiate trl.DPOTrainer(model, args=DPOConfig(...), train_dataset, tokenizer), and let it handle the Bradley-Terry log-sigmoid loss over the β-scaled log-ratio of policy vs reference policy under the hood. The lab deliberately uses GPT-2 and eight preference pairs, not because that's production-realistic but because it collapses a 30-minute mechanics demonstration into a 55-minute lab where you can see the alignment tax (the general-capability drift that follows narrow preference optimisation) on specific prompts. The production gap is explicit: real DPO runs use UltraFeedback (~64k pairs), HH-RLHF (~161k), or Nectar (~1M), layer multiple datasets, and frequently blend in SFT loss or iterate DPO rounds to manage capability regression. You'll also get the comparative map: KTO for thumbs-up/down data, IPO for small-dataset stability, ORPO for skipping the SFT-then-align split, PPO-RLHF for when reward-model expressivity matters more than implementation simplicity.

Prerequisites: fluency with Hugging Face causal-LM generation (.generate(), sampling params), a working sense of supervised fine-tuning loss, and a rough grasp of what RLHF is trying to accomplish versus pretraining. The sandbox has trl, transformers, datasets, GPT-2, and a compatible TRL release preinstalled. If you're Googling "DPO vs PPO RLHF", "DPOTrainer tutorial", "how to align a small LLM", "preference dataset format TRL", or "KTO vs DPO vs ORPO", this is the hands-on answer, and the code ports directly to Llama 3 or Mistral 7B with LoRA/QLoRA for production-scale alignment.