Модуль 8. RLHF и прямые методы оптимизации предпочтений

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

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

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()))

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

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]),
})

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

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()

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"]))

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.")

# 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.")