필사 모드: DPO（Direct Preference Optimization）论文深度解析 — 无需 RLHF 对齐 LLM

# RLHF (PPO) 训练循环伪代码
import torch

# 需要同时加载 4 个模型
policy_model = load_model("sft_checkpoint")
reference_model = load_model("sft_checkpoint")  # frozen
reward_model = load_model("reward_model_checkpoint")
value_model = load_model("value_head_checkpoint")

for batch in dataloader:
    prompts = batch["prompt"]

    # 1. 用策略模型生成响应(需要在线生成)
    responses = policy_model.generate(prompts)

    # 2. 用奖励模型计算分数
    rewards = reward_model(prompts, responses)

    # 3. 用 Value function 估计 advantage (GAE)
    values = value_model(prompts, responses)
    advantages = compute_gae(rewards, values, gamma=0.99, lam=0.95)

    # 4. 优化 PPO 裁剪目标函数
    ratio = policy_model.log_prob(responses) - reference_model.log_prob(responses)
    clipped_ratio = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps)
    policy_loss = -torch.min(ratio * advantages, clipped_ratio * advantages).mean()

    # 5. 加入 KL 惩罚项
    kl_penalty = beta * kl_divergence(policy_model, reference_model, responses)
    total_loss = policy_loss + kl_penalty

    total_loss.backward()
    optimizer.step()

# DPO 训练循环伪代码
import torch
import torch.nn.functional as F

# 只需要 2 个模型
policy_model = load_model("sft_checkpoint")
reference_model = load_model("sft_checkpoint")  # frozen

for batch in dataloader:
    # 直接从离线数据中训练(不需要在线生成)
    prompts = batch["prompt"]
    chosen = batch["chosen"]
    rejected = batch["rejected"]

    # 1. 计算各模型的 log probability
    pi_chosen = policy_model.log_prob(chosen, prompts)
    pi_rejected = policy_model.log_prob(rejected, prompts)
    ref_chosen = reference_model.log_prob(chosen, prompts)
    ref_rejected = reference_model.log_prob(rejected, prompts)

    # 2. DPO 损失函数(只用 3 行!)
    log_ratio_chosen = pi_chosen - ref_chosen
    log_ratio_rejected = pi_rejected - ref_rejected
    loss = -F.logsigmoid(beta * (log_ratio_chosen - log_ratio_rejected)).mean()

    loss.backward()
    optimizer.step()

import torch
import torch.nn.functional as F
from torch import nn


def dpo_loss(
    policy_chosen_logps: torch.Tensor,
    policy_rejected_logps: torch.Tensor,
    reference_chosen_logps: torch.Tensor,
    reference_rejected_logps: torch.Tensor,
    beta: float = 0.1,
    label_smoothing: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    DPO 损失函数实现。

    Args:
        policy_chosen_logps: 策略模型对 chosen 响应的 log probability
        policy_rejected_logps: 策略模型对 rejected 响应的 log probability
        reference_chosen_logps: 参考模型对 chosen 响应的 log probability
        reference_rejected_logps: 参考模型对 rejected 响应的 log probability
        beta: KL 约束强度(temperature)
        label_smoothing: 标签平滑系数(为 0 时是标准 DPO)

    Returns:
        loss: DPO 损失值
        chosen_rewards: chosen 响应的隐式奖励
        rejected_rewards: rejected 响应的隐式奖励
    """
    # 计算隐式奖励: r(x,y) = beta * log(pi/pi_ref)
    chosen_rewards = beta * (policy_chosen_logps - reference_chosen_logps)
    rejected_rewards = beta * (policy_rejected_logps - reference_rejected_logps)

    # 计算奖励边际
    reward_margin = chosen_rewards - rejected_rewards

    # 应用标签平滑(可选)
    if label_smoothing > 0:
        loss = (
            -F.logsigmoid(reward_margin) * (1 - label_smoothing)
            - F.logsigmoid(-reward_margin) * label_smoothing
        )
    else:
        loss = -F.logsigmoid(reward_margin)

    return loss.mean(), chosen_rewards.mean(), rejected_rewards.mean()


# 使用示例
batch_size = 4
seq_len = 128

# 模拟的 log probability 值
policy_chosen_logps = torch.randn(batch_size) * 0.1 - 2.0
policy_rejected_logps = torch.randn(batch_size) * 0.1 - 2.5
reference_chosen_logps = torch.randn(batch_size) * 0.1 - 2.0
reference_rejected_logps = torch.randn(batch_size) * 0.1 - 2.5

loss, chosen_r, rejected_r = dpo_loss(
    policy_chosen_logps,
    policy_rejected_logps,
    reference_chosen_logps,
    reference_rejected_logps,
    beta=0.1,
)
print(f"DPO Loss: {loss.item():.4f}")
print(f"Chosen Reward: {chosen_r.item():.4f}")
print(f"Rejected Reward: {rejected_r.item():.4f}")

import torch
import torch.nn.functional as F


def compute_log_probs(
    logits: torch.Tensor,
    labels: torch.Tensor,
    attention_mask: torch.Tensor,
) -> torch.Tensor:
    """
    计算逐 token 的 log probability，并按序列汇总。

    Args:
        logits: 模型输出的 logits (batch_size, seq_len, vocab_size)
        labels: 正确的 token ID (batch_size, seq_len)
        attention_mask: attention mask (batch_size, seq_len)

    Returns:
        按序列汇总的 log probability (batch_size,)
    """
    # 把 logits 向前移动 1 个 token(next token prediction)
    shift_logits = logits[:, :-1, :]
    shift_labels = labels[:, 1:]
    shift_mask = attention_mask[:, 1:]

    # 逐 token 的 log probability
    log_probs = F.log_softmax(shift_logits, dim=-1)
    per_token_logps = torch.gather(
        log_probs, dim=2, index=shift_labels.unsqueeze(2)
    ).squeeze(2)

    # 被掩码的 token 处理为 0，再按序列求和
    per_token_logps = per_token_logps * shift_mask
    return per_token_logps.sum(dim=-1)

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import DPOConfig, DPOTrainer
from peft import LoraConfig

# 1. 加载模型和分词器
model_name = "Qwen/Qwen2.5-7B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="bfloat16",
    attn_implementation="flash_attention_2",
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 2. LoRA 配置(节省内存)
peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    task_type="CAUSAL_LM",
)

# 3. 加载偏好数据(prompt、chosen、rejected 结构)
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")

# 4. DPO 训练配置
training_args = DPOConfig(
    output_dir="./dpo-qwen2.5-7b",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    learning_rate=5e-6,         # 比 SFT 低 10~100 倍
    beta=0.1,                   # KL 约束强度
    max_length=1024,
    max_prompt_length=512,
    num_train_epochs=1,
    bf16=True,
    logging_steps=10,
    save_strategy="steps",
    save_steps=500,
    warmup_ratio=0.1,
    gradient_checkpointing=True,
    remove_unused_columns=False,
)

# 5. 创建 DPO Trainer 并开始训练
trainer = DPOTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    processing_class=tokenizer,
    peft_config=peft_config,
)
trainer.train()
trainer.save_model("./dpo-qwen2.5-7b-final")

# 安装必需的包
pip install trl>=0.12.0 transformers>=4.46.0 peft>=0.13.0 \
    datasets accelerate bitsandbytes flash-attn

# 多 GPU 训练执行(DeepSpeed ZeRO-3)
accelerate launch --config_file deepspeed_zero3.yaml \
    --num_processes 4 \
    train_dpo.py

# 训练监控
tensorboard --logdir ./dpo-qwen2.5-7b/runs

# 生成合成偏好数据的流水线示例
from openai import OpenAI

client = OpenAI()

def generate_preference_pair(prompt: str, model_response_a: str, model_response_b: str) -> dict:
    """使用 GPT-4 进行偏好判定"""
    judge_prompt = f"""请比较以下两个响应，判断哪一个更有帮助、
更准确、更安全。

Prompt: {prompt}

响应 A: {model_response_a}

响应 B: {model_response_b}

只输出更好响应的编号(A 或 B):"""

    result = client.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": judge_prompt}],
        max_tokens=1,
    )
    choice = result.choices[0].message.content.strip()

    if choice == "A":
        return {"prompt": prompt, "chosen": model_response_a, "rejected": model_response_b}
    else:
        return {"prompt": prompt, "chosen": model_response_b, "rejected": model_response_a}

# DPO 训练过程中应监控的核心指标
def log_dpo_metrics(chosen_rewards, rejected_rewards, loss):
    """记录训练过程中核心指标的函数"""
    reward_margin = (chosen_rewards - rejected_rewards).mean()
    accuracy = (chosen_rewards > rejected_rewards).float().mean()

    metrics = {
        "dpo/loss": loss.item(),
        "dpo/reward_margin": reward_margin.item(),
        "dpo/accuracy": accuracy.item(),
        "dpo/chosen_reward_mean": chosen_rewards.mean().item(),
        "dpo/rejected_reward_mean": rejected_rewards.mean().item(),
    }
    # reward_margin 持续增大是过拟合的征兆
    # accuracy 达到 1.0 时应考虑停止训练
    return metrics