大模型 RL 对齐（RLHF / DPO）
原理与配套极简源码片段

import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer

# ========== 1. SFT监督微调示意代码 ==========
# SFT的目标：标准自回归语言模型损失，只在response部分计算loss
model = AutoModelForCausalLM.from_pretrained("your-sft-base")
tokenizer = AutoTokenizer.from_pretrained("your-sft-base")

def sft_loss(prompt: str, response: str):
    full_text = prompt + response
    tokens = tokenizer(full_text, return_tensors="pt")
    outputs = model(**tokens)
    logits = outputs.logits
    labels = tokens["input_ids"]
    # 自回归损失：用第i个token预测第i+1个token
    # 只对回答部分计算loss是工程上常用优化，这里简化演示
    loss = nn.CrossEntropyLoss()(logits[:,:-1,:].reshape(-1, logits.size(-1)), labels[:,1:].reshape(-1))
    return loss

# ========== 2. 奖励模型RM训练示意代码 ==========
# 奖励模型：在基础大模型之上增加一个reward head，输出标量奖励
class RewardModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.base = base_model
        self.reward_head = nn.Linear(base_model.config.hidden_size, 1)

    def forward(self, input_ids):
        out = self.base(input_ids=input_ids, output_hidden_states=True)
        last_hidden = out.hidden_states[-1][:, -1, :] # 取最后一个token的hidden state
        reward = self.reward_head(last_hidden)
        return reward

def rm_pair_loss(rm, prompt, chosen, rejected):
    """成对偏好损失：保证chosen样本奖励 > rejected样本奖励"""
    tok_chosen = tokenizer(prompt + chosen, return_tensors="pt")
    tok_rej = tokenizer(prompt + rejected, return_tensors="pt")
    r_chosen = rm(**tok_chosen)
    r_rej = rm(**tok_rej)
    # 排序损失，拉开优选回答和劣选回答之间的奖励差距
    loss = torch.relu(-(r_chosen - r_rej)).mean()
    return loss

# ========== PPO Clip Loss 核心实现 ==========
def ppo_clip_loss(log_prob_new, log_prob_old, advantage, eps=0.2):
    """
    log_prob_new: 当前策略模型动作log概率
    log_prob_old: rollout阶段旧策略保存下来的log概率
    advantage: 优势函数，衡量当前动作比平均水平好多少
    eps: clip截断系数，通常取0.2
    """
    ratio = torch.exp(log_prob_new - log_prob_old)
    unclipped = ratio * advantage
    clipped = torch.clamp(ratio, 1 - eps, 1 + eps) * advantage
    # 取最小值，负号是因为梯度下降最小化损失
    loss = -torch.min(unclipped, clipped).mean()
    return loss

# ========== KL惩罚项实现 ==========
def kl_divergence(policy_logits, ref_logits):
    # KL(policy || ref)，计算策略分布相对参考SFT分布的散度
    logp = torch.log_softmax(policy_logits, dim=-1)
    logref = torch.log_softmax(ref_logits, dim=-1)
    kl = (torch.exp(logp) * (logp - logref)).sum(dim=-1)
    return kl

def ppo_total_loss(log_prob_new, log_prob_old, adv, policy_logits, ref_logits, kl_coeff=0.05):
    clip_loss = ppo_clip_loss(log_prob_new, log_prob_old, adv)
    kl = kl_divergence(policy_logits, ref_logits).mean()
    total_loss = clip_loss + kl_coeff * kl
    return total_loss

# ========== GAE（广义优势估计）示意代码 ==========
def compute_gae(rewards, values, gamma=0.99, lam=0.95):
    """
    rewards: 每一步的即时奖励 [T]
    values:  价值模型对每个状态的估计 V(s_t) [T+1]，最后一位是 bootstrap
    返回: advantages [T], returns [T]
    """
    T = len(rewards)
    advantages = torch.zeros(T)
    last_gae = 0.0
    for t in reversed(range(T)):
        # TD 误差 δ_t = r_t + γV(s_{t+1}) - V(s_t)
        delta = rewards[t] + gamma * values[t + 1] - values[t]
        # 递推累积：A_t = δ_t + γλ · A_{t+1}
        last_gae = delta + gamma * lam * last_gae
        advantages[t] = last_gae
    returns = advantages + values[:-1]   # 回报 = 优势 + 价值基线
    return advantages, returns

# ========== PPO 单轮迭代骨架 ==========
def ppo_train_step(policy, ref_model, reward_model, critic,
                   prompts, optimizer, gen_kwargs,
                   clip_eps=0.2, kl_coeff=0.05, gamma=0.99, lam=0.95):
    """一次完整的 PPO 迭代：采样 → 打分 → 优势 → 更新"""
    # ---- ① Rollout：策略模型批量采样，保存 log_prob_old ----
    with torch.no_grad():
        responses, logp_old, values_old = [], [], []
        for p in prompts:
            tokens = tokenizer(p, return_tensors="pt")
            gen = policy.generate(**tokens, **gen_kwargs)
            resp = tokenizer.decode(gen[0], skip_special_tokens=True)
            responses.append(resp)

            full = tokenizer(p + resp, return_tensors="pt")
            # 策略模型在采样时的 log 概率（旧策略）
            out_old = policy(**full)
            logp_old.append(policy_logprob(out_old.logits, full["input_ids"]))
            # Critic 估计状态价值
            values_old.append(critic(**full))

    # ---- ② Reward Model 打分 ----
    with torch.no_grad():
        rewards = torch.stack([
            reward_model(tokenizer(p + r, return_tensors="pt")["input_ids"])
            for p, r in zip(prompts, responses)
        ]).squeeze(-1)

    # ---- ③ GAE 优势估计 ----
    values_cat = torch.cat([v.squeeze(-1) for v in values_old] + [torch.zeros(1)])
    advantages, returns = compute_gae(rewards, values_cat, gamma, lam)
    # 优势标准化，稳定训练
    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

    # ---- ④ 多轮 minibatch 更新 ----
    total_loss_val = 0.0
    for _ in range(4):  # 同一批数据复用 4 个 epoch
        out_new = policy(**tokenizer(prompts[0] + responses[0], return_tensors="pt"))
        logp_new = policy_logprob(out_new.logits, tokenizer(prompts[0] + responses[0])["input_ids"])

        # PPO-Clip 损失
        clip_loss = ppo_clip_loss(logp_new, logp_old[0], advantages[0], eps=clip_eps)
        # KL 惩罚（对齐冻结的参考模型）
        with torch.no_grad():
            out_ref = ref_model(**tokenizer(prompts[0] + responses[0], return_tensors="pt"))
        kl = kl_divergence(out_new.logits, out_ref.logits).mean()
        # 价值模型回归损失
        value_loss = torch.nn.functional.mse_loss(critic_out, returns[0])

        loss = clip_loss + kl_coeff * kl + 0.5 * value_loss

        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(policy.parameters(), max_norm=1.0)
        optimizer.step()
        total_loss_val += loss.item()

    return {"loss": total_loss_val / 4, "reward_mean": rewards.mean().item()}

# ========== DPO Loss 实现 ==========
def get_log_prob(model, input_ids):
    """辅助函数：计算给定序列的对数概率之和"""
    out = model(input_ids=input_ids)
    logits = out.logits
    logp_tokens = torch.log_softmax(logits, dim=-1)
    token_logp = logp_tokens[:, :-1].gather(-1, input_ids[:,1:].unsqueeze(-1)).squeeze(-1)
    return token_logp.sum(dim=-1)

def dpo_loss(policy_model, ref_model, prompt, chosen, rejected, beta=0.1):
    tok_chosen = tokenizer(prompt + chosen, return_tensors="pt")
    tok_rej = tokenizer(prompt + rejected, return_tensors="pt")

    logp_chosen_pol = get_log_prob(policy_model, **tok_chosen)
    logp_rej_pol = get_log_prob(policy_model, **tok_rej)

    logp_chosen_ref = get_log_prob(ref_model, **tok_chosen)
    logp_rej_ref = get_log_prob(ref_model, **tok_rej)

    log_ratio_chosen = logp_chosen_pol - logp_chosen_ref
    log_ratio_rej = logp_rej_pol - logp_rej_ref

    loss = -torch.nn.functional.logsigmoid(beta * (log_ratio_chosen - log_ratio_rej))
    return loss.mean()

# ========== ORPO Loss 示意 ==========
def orpo_loss(policy_model, prompt, chosen, rejected, beta=0.1):
    # ORPO：SFT损失 + odds ratio偏好损失合并
    sft_loss_val = sft_loss(prompt, chosen)
    tok_chosen = tokenizer(prompt + chosen, return_tensors="pt")
    tok_rej = tokenizer(prompt + rejected, return_tensors="pt")
    logp_chosen = get_log_prob(policy_model, **tok_chosen)
    logp_rej = get_log_prob(policy_model, **tok_rej)
    odds_ratio = torch.log(torch.sigmoid(beta * (logp_chosen - logp_rej)))
    total_loss = sft_loss_val - odds_ratio
    return total_loss.mean()

# ========== Rollout 采样示意代码 ==========
def rollout(policy_model, rm, prompts, gen_kwargs):
    """rollout阶段：策略模型批量生成回答，RM打分"""
    responses = []
    rewards = []
    for p in prompts:
        tokens = tokenizer(p, return_tensors="pt")
        gen_out = policy_model.generate(**tokens, **gen_kwargs)
        resp = tokenizer.decode(gen_out[0], skip_special_tokens=True)
        responses.append(resp)
        # RM打分得到reward
        reward = rm(tokenizer(p+resp, return_tensors="pt"))
        rewards.append(reward)
    return responses, torch.cat(rewards)

# ========== GRPO 组内相对奖励示意 ==========
def grpo_group_reward(rewards_group):
    """对同一个prompt的N个回答做组内归一化"""
    mean_r = rewards_group.mean()
    std_r = rewards_group.std() + 1e-8
    normalized_reward = (rewards_group - mean_r) / std_r
    return normalized_reward