VKVlad KulikovPublic / lectures
RL for LLM8 / 13

Module 8. RLHF and direct preference optimization

Read the explanations and code in the existing notebooks, or download the files. Code does not run on this site.

Download course files (.zip) 382.1 KB

The archive preserves the source folder structure, both languages, and supporting files.

Module 8. RLHF and direct preference optimization

Main module notebook

Download .ipynb module-08-rlhf-dpo.en.ipynb · 17.1 KB

Show notebook text

Module 8. RLHF and direct preference optimization

This notebook links the module's derivations to two levels of practice. The default route is fully reproducible on CPU: it solves a finite KL-regularized problem, checks the DPO loss and gradient, constructs an exact likelihood-displacement counterexample, and trains a small autoregressive policy using genuine sequence log-probabilities and completion masks. The final optional branch uses Qwen3-0.6B and TRL; it is disabled by default and is not presented as an executed course result.

0. Reproducible setup

The code locates the extracted course root automatically. PROFILE="smoke" runs the full teaching workflow in a few seconds; PROFILE="full" increases the training budget. Every result reported below belongs only to this synthetic CPU experiment.

Python
from __future__ import annotations

from pathlib import Path
import sys
import json
import subprocess

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
from IPython.display import display

_candidates = [Path.cwd(), Path.cwd().parent, Path.cwd().parent.parent]
ROOT = next((p for p in _candidates if (p / "modules" / "module-08").exists()), None)
if ROOT is None:
    raise FileNotFoundError("Could not locate the RL for LLM repository root")
MODULE_DIR = ROOT / "modules" / "module-08"
sys.path.insert(0, str(MODULE_DIR))

from module_8_reference import (
    conservative_dpo_loss,
    constrained_lagrangian,
    dpo_gradient_weight,
    dpo_loss_from_margin,
    ipo_loss,
    likelihood_displacement_curve,
    optimal_kl_policy,
    orpo_total_loss,
    reward_kl_frontier,
    shape_token_rewards,
    simpo_loss,
)
from module_8_training import (
    build_vocabulary,
    pair_batch,
    run_toy_dpo_experiment,
    toy_preference_examples,
)

PROFILE = "smoke"  # "full" uses a larger CPU training budget
RUN_REAL_MODEL = False
SEED = 7
np.random.seed(SEED)
torch.manual_seed(SEED)
print({"root": str(ROOT), "profile": PROFILE, "real_model": RUN_REAL_MODEL})

1. Exact KL-regularized optimum

On a finite response space, π∗∝πrefexp⁡(r/β)\pi^*\propto\pi_{ref}\exp(r/\beta) can be computed without stochastic optimization. This separates exact variational geometry from neural-training error. We also verify invariance to a constant reward shift.

Python
reference = np.array([0.55, 0.30, 0.15])
rewards = np.array([0.0, 1.0, -0.5])
betas = np.array([4.0, 2.0, 1.0, 0.5, 0.25])
frontier = reward_kl_frontier(reference, rewards, betas)
table = pd.DataFrame({
    "beta": frontier["beta"],
    "expected_reward": frontier["reward"],
    "KL_to_reference": frontier["kl"],
})
display(table.round(6))

fig, ax = plt.subplots(figsize=(7.5, 4.5))
ax.plot(frontier["kl"], frontier["reward"], marker="o")
ax.set_xlabel(r"$D_{KL}(\pi^*\Vert\pi_{ref})$")
ax.set_ylabel("expected reward")
ax.set_title("Exact reward–KL frontier")
ax.grid(True, alpha=0.25)
plt.show()

base = optimal_kl_policy(reference, rewards, beta=1.0)
shifted = optimal_kl_policy(reference, rewards + 100.0, beta=1.0)
assert np.allclose(base, shifted)
print("beta=1 optimum:", np.round(base, 6))

2. Token-level reward shaping

The penalty −β(log⁡π−log⁡πref)-\beta(\log\pi-\log\pi_{ref}) is applied only to active response tokens. A task verifier or reward model supplies a separate terminal signal. Prompt positions must not leak into the sum.

Python
policy_logp = np.array([-0.8, -1.1, -0.4, -0.7])
reference_logp = np.array([-0.9, -0.9, -0.5, -0.6])
response_mask = np.array([0.0, 1.0, 1.0, 1.0])
shaped = shape_token_rewards(
    policy_logp,
    reference_logp,
    terminal_reward=2.0,
    beta=0.15,
    mask=response_mask,
)
display(pd.DataFrame({
    "policy_logp": policy_logp,
    "reference_logp": reference_logp,
    "response_mask": response_mask,
    "shaped_reward": shaped,
}).round(4))
assert shaped[0] == 0.0
print("sum of active shaped rewards:", float(shaped.sum()))

3. DPO loss, gradient weight, cDPO, and IPO

DPO emphasizes misranked pairs and saturates on pairs already separated. The coefficient β\beta changes both logit scale and the derivative. cDPO and IPO encode different targets rather than cosmetic renamings of DPO.

Python
margins = np.linspace(-5.0, 5.0, 401)
fig, ax = plt.subplots(figsize=(7.5, 4.5))
for beta in (0.5, 1.0, 2.0):
    ax.plot(margins, dpo_gradient_weight(margins, beta), label=fr"$\beta={beta}$")
ax.set_xlabel("reference-corrected margin h")
ax.set_ylabel(r"$\beta\sigma(-\beta h)$")
ax.set_title("Which pairs receive the largest DPO update?")
ax.legend()
ax.grid(True, alpha=0.25)
plt.show()

h = 0.8
display(pd.DataFrame([{
    "DPO": float(dpo_loss_from_margin(h, beta=0.2)),
    "cDPO eps=0.1": float(conservative_dpo_loss(h, beta=0.2, label_smoothing=0.1)),
    "IPO tau=0.5": float(ipo_loss(h, tau=0.5)),
}]).round(6))

4. Exact likelihood-displacement counterexample

The pairwise objective controls a relative margin. In a three-response model, the chosen–rejected margin rises as 2t2t even while chosen-response probability falls because mass moves to an unlabeled third response.

Python
t = np.linspace(0.0, 1.0, 301)
disp = likelihood_displacement_curve(t)
fig, ax = plt.subplots(figsize=(8.0, 4.8))
ax.plot(t, disp["chosen"], label="chosen response")
ax.plot(t, disp["rejected"], label="rejected response")
ax.plot(t, disp["other"], label="unlabeled response")
ax.set_xlabel("update parameter t")
ax.set_ylabel("probability")
ax.set_title("A larger pair margin need not raise chosen probability")
ax.legend()
ax.grid(True, alpha=0.25)
plt.show()

assert np.all(np.diff(disp["log_margin"]) >= 0)
assert disp["chosen"][-1] < disp["chosen"][0]
print({
    "initial_chosen_probability": float(disp["chosen"][0]),
    "final_chosen_probability": float(disp["chosen"][-1]),
    "final_log_margin": float(disp["log_margin"][-1]),
})

5. Sequence-pair contract

One prompt is encoded with two complete continuations. Current and reference models must share the same template, tokenization, and truncation. After causal shifting, the mask includes completion tokens and EOS only.

Python
examples = toy_preference_examples()
vocab = build_vocabulary(examples)
batch = pair_batch(examples, vocab)
first = examples[0]
inspection = {
    "prompt": " ".join(first.prompt),
    "chosen": " ".join(first.chosen),
    "rejected": " ".join(first.rejected),
    "chosen_ids": batch["chosen_ids"][0].tolist(),
    "chosen_response_mask": batch["chosen_mask"][0].int().tolist(),
    "rejected_ids": batch["rejected_ids"][0].tolist(),
    "rejected_response_mask": batch["rejected_mask"][0].int().tolist(),
}
print(json.dumps(inspection, ensure_ascii=False, indent=2))
assert batch["chosen_mask"].shape[1] == batch["chosen_ids"].shape[1] - 1
assert batch["chosen_mask"][0].sum().item() == len(first.chosen) + 1
assert batch["rejected_mask"][0].sum().item() == len(first.rejected) + 1

6. A realistic small DPO loop

A reference GRU language model is first trained on both response families. The policy is cloned from it and optimized with DPO. We log not only loss and pair accuracy, but also absolute log-probabilities, KL, and free greedy-generation quality.

Python
experiment = run_toy_dpo_experiment(profile=PROFILE, seed=SEED)
summary = pd.DataFrame([
    {"regime": "reference", **experiment["before"], "generation_match": experiment["reference_generation_match"]},
    {"regime": "moderate DPO", **experiment["after"], "generation_match": experiment["policy_generation_match"]},
    {"regime": "high-pressure DPO", **experiment["stress"], "generation_match": experiment["stress_generation_match"]},
])
display(summary[[
    "regime", "loss", "pair_accuracy", "mean_margin", "mean_chosen_logp",
    "mean_rejected_logp", "generation_match"
]].round(6))
assert np.isfinite(summary.select_dtypes(include="number").to_numpy()).all()

7. Moderate training versus high optimization pressure

Both regimes can rank the fixed pair set perfectly. Their behavior still differs: under high pressure, chosen log-probability falls and free generation degrades. Pair accuracy alone is therefore not a sufficient diagnostic.

Python
stable = pd.DataFrame(experiment["history"])
stress = pd.DataFrame(experiment["stress_history"])
fig, ax = plt.subplots(figsize=(8.5, 5.0))
ax.plot(stable["step"], stable["mean_chosen_logp"], label="moderate: chosen")
ax.plot(stable["step"], stable["mean_rejected_logp"], label="moderate: rejected")
ax.plot(stress["step"], stress["mean_chosen_logp"], linestyle="--", label="high pressure: chosen")
ax.plot(stress["step"], stress["mean_rejected_logp"], linestyle="--", label="high pressure: rejected")
ax.set_xlabel("DPO update")
ax.set_ylabel("mean sequence log-probability")
ax.set_title("Pair accuracy alone misses absolute likelihood displacement")
ax.legend()
ax.grid(True, alpha=0.25)
plt.show()

print("teacher-forced KL on chosen prefixes:", round(experiment["teacher_forced_kl_chosen"], 6))
print("teacher-forced KL on rejected prefixes:", round(experiment["teacher_forced_kl_rejected"], 6))
display(pd.DataFrame(experiment["stress_generations"]))

8. Methods differ by data contract

SimPO and ORPO are placed side by side not to collapse them into one formula, but to expose differences in reference anchoring, length normalization, supervised terms, and feedback type. We also evaluate a Lagrangian objective with multiple cost constraints.

Python
contracts = pd.DataFrame([
    {
        "method": "SimPO",
        "loss": simpo_loss(-4.0, 4, -5.4, 6, beta=2.0, target_margin=0.3),
        "reference_model": "no",
        "feedback": "pairs",
    },
    {
        "method": "ORPO",
        "loss": orpo_total_loss(0.8, -0.7, -1.5, lam=0.2),
        "reference_model": "no",
        "feedback": "pairs + chosen NLL",
    },
])
display(contracts.round(6))
lagrangian = constrained_lagrangian(
    expected_reward=1.4,
    expected_costs=[0.10, 0.03],
    limits=[0.06, 0.04],
    multipliers=[2.0, 1.5],
)
print("constrained Lagrangian:", round(lagrangian, 6))
print("Inspect feedback, reference anchor, normalization, and constraints before choosing a method.")

9. Optional Qwen3-0.6B + TRL branch

The current official TRL example uses Qwen/Qwen3-0.6B and trl-lib/ultrafeedback_binarized. The branch below provides a short LoRA template and a pinned version snapshot. It needs network access and a suitable accelerator, so the release CPU gate did not execute it. UltraFeedback is synthetic AI feedback, not direct human annotation.

Python
# Optional real-model branch. It is deliberately disabled in the release artifact.
if RUN_REAL_MODEL:
    packages = [
        "trl==1.9.2",
        "transformers==5.14.1",
        "peft==0.20.0",
        "accelerate==1.14.0",
        "datasets==5.0.1",
    ]
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *packages])

    from datasets import load_dataset
    from peft import LoraConfig
    from trl import DPOConfig, DPOTrainer

    dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train").select(range(64))
    args = DPOConfig(
        output_dir="module-08-qwen3-dpo-smoke",
        max_steps=20,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,
        learning_rate=1e-5,
        logging_steps=1,
        report_to="none",
    )
    trainer = DPOTrainer(
        model="Qwen/Qwen3-0.6B",
        args=args,
        train_dataset=dataset,
        peft_config=LoraConfig(r=8, lora_alpha=16, lora_dropout=0.05),
    )
    trainer.train()
else:
    print("RUN_REAL_MODEL=False: no download or GPU training was attempted.")