1
class RewardModel(torch.nn.Module):
2
    def __init__(self, base_model):
3
        super().__init__()
4
        self.base_model = base_model
5
        # 替换 LM head 为 reward head
6
        self.value_head = torch.nn.Linear(
7
            base_model.config.hidden_size,
8
            1,
9
            bias=False
10
        )
11

12
    def forward(self, input_ids, attention_mask):
13
        outputs = self.base_model(
14
            input_ids=input_ids,
15
            attention_mask=attention_mask
16
        )
17

18
        # 使用最后一个 token 的 hidden state 预测 reward
19
        last_hidden = outputs.last_hidden_state[:, -1, :]
20
        reward = self.value_head(last_hidden)
21

22
        return reward

1
class PreferenceDataBuilder:
2
    def __init__(self, annotation_interface):
3
        self.interface = annotation_interface
4

5
    def build_preference_data(self, prompts: list, responses: list) -> list:
6
        """
7
        构建偏好数据：
8
        给定一个 prompt，收集多个 response，由人类标注偏好
9
        """
10
        preference_data = []
11

12
        for prompt in prompts:
13
            # 生成多个候选 response
14
            candidates = self.generate_candidates(prompt, num_choices=4)
15

16
            # 人类偏好标注
17
            # 标注界面展示 prompt + candidates，标注哪个更好
18
            annotations = self.interface.annotate(prompt, candidates)
19

20
            # 构建 pairwise 偏好数据
21
            for chosen, rejected in annotations:
22
                preference_data.append({
23
                    "prompt": prompt,
24
                    "chosen": chosen,
25
                    "rejected": rejected,
26
                    "preference": 1  # chosen 优于 rejected
27
                })
28

29
        return preference_data
30

31
    def generate_candidates(self, prompt, num_choices=4):
32
        """使用模型生成多样化的候选回复"""
33
        candidates = []
34
        for _ in range(num_choices):
35
            response = self.model.generate(
36
                prompt,
37
                temperature=0.8,  # 高温度增加多样性
38
                top_p=0.95
39
            )
40
            candidates.append(response)
41
        return candidates

1
def reward_model_loss(reward_chosen, reward_rejected):
2
    """
3
    Reward Model 损失函数：Bradley-Terry 模型
4
    目标：P(chosen > rejected) = sigmoid(reward_chosen - reward_rejected)
5
    """
6
    # 偏好差异
7
    diff = reward_chosen - reward_rejected
8

9
    # 对数损失
10
    loss = -torch.log(torch.sigmoid(diff)).mean()
11

12
    # 添加对比损失（可选）
13
    contrastive_loss = torch.relu(1 - diff).mean()
14

15
    return loss + 0.1 * contrastive_loss
16

17
class RewardTrainer:
18
    def training_step(self, batch):
19
        # 计算 chosen response 的 reward
20
        reward_chosen = self.reward_model(
21
            input_ids=batch["chosen_input_ids"],
22
            attention_mask=batch["chosen_attention_mask"]
23
        )
24

25
        # 计算 rejected response 的 reward
26
        reward_rejected = self.reward_model(
27
            input_ids=batch["rejected_input_ids"],
28
            attention_mask=batch["rejected_attention_mask"]
29
        )
30

31
        loss = reward_model_loss(reward_chosen, reward_rejected)
32

33
        self.backward(loss)
34
        return loss

1
class PPOTrainer:
2
    def __init__(self, policy, ref_policy, reward_model, kl_coef=0.04):
3
        self.policy = policy
4
        self.ref_policy = ref_policy  # SFT 模型作为参考
5
        self.reward_model = reward_model
6
        self.kl_coef = kl_coef
7

8
    def compute_rewards(self, logits, responses, log_probs_old):
9
        """
10
        计算广义优势估计（GAE）
11
        """
12
        # 1. 从 Reward Model 获取即时奖励
13
        rewards = self.reward_model(responses)
14

15
        # 2. KL 散度惩罚（防止 policy 偏离 SFT 太远）
16
        kl_penalty = self.kl_coef * self.compute_kl_penalty(
17
            logits, log_probs_old
18
        )
19

20
        # 3. 广义优势估计
21
        advantages = self.compute_gae(rewards, kl_penalty)
22

23
        return advantages
24

25
    def compute_kl_penalty(self, logits, log_probs_old):
26
        """
27
        KL(π_new || π_old) 惩罚
28
        """
29
        log_probs_new = torch.log_softmax(logits, dim=-1)
30
        kl = torch.exp(log_probs_new) * (log_probs_new - log_probs_old)
31
        return kl.sum(dim=-1)
32

33
    def ppo_update(self, batches):
34
        """
35
        PPO 核心更新
36
        """
37
        for batch in batches:
38
            # 1. 计算 ratio = π_new / π_old
39
            ratio = torch.exp(batch["log_probs"] - batch["log_probs_old"])
40

41
            # 2. PPO-Clip 目标函数
42
            surr1 = ratio * batch["advantages"]
43
            surr2 = torch.clamp(
44
                ratio,
45
                1 - self.epsilon,  # 0.2
46
                1 + self.epsilon
47
            ) * batch["advantages"]
48

49
            # 3. 取最小值（clip 外的部分被忽略）
50
            policy_loss = -torch.min(surr1, surr2).mean()
51

52
            # 4. 添加 KL 惩罚
53
            kl_penalty = batch["kl_penalty"].mean()
54

55
            total_loss = policy_loss + kl_penalty
56

57
            self.optimizer.zero_grad()
58
            self.backward(total_loss)
59
            self.optimizer.step()

1
class DPOConfig:
2
    def __init__(
3
        self,
4
        beta: float = 0.1,  # KL 惩罚系数
5
        lr: float = 1e-6,
6
        batch_size: int = 8,
7
    ):
8
        self.beta = beta
9
        self.lr = lr
10
        self.batch_size = batch_size
11

12
def dpo_loss(policy_chosen_logps, policy_rejected_logps,
13
             ref_chosen_logps, ref_rejected_logps, beta=0.1):
14
    """
15
    DPO 损失函数
16

17
    直观理解：
18
    - policy_chosen - ref_chosen: chosen response 在新策略下比参考策略好的程度
19
    - policy_rejected - ref_rejected: rejected response 在新策略下比参考策略差的程度
20

21
    目标：最大化 (chosen 比 rejected 好的程度)
22
    """
23
    # 策略概率比
24
    chosen_ratio = torch.exp(policy_chosen_logps - ref_chosen_logps)
25
    rejected_ratio = torch.exp(policy_rejected_logps - ref_rejected_logps)
26

27
    # DPO 公式
28
    # log(sigmoid(chosen_advantage - rejected_advantage))
29
    chosen_advantage = torch.log(chosen_ratio) / beta
30
    rejected_advantage = torch.log(rejected_ratio) / beta
31

32
    loss = -torch.log(torch.sigmoid(chosen_advantage - rejected_advantage))
33

34
    return loss.mean()

1
from transformers import AutoTokenizer, AutoModelForCausalLM
2
from torch.utils.data import DataLoader
3

4
class DPOTrainer:
5
    def __init__(self, model_name, beta=0.1):
6
        self.model = AutoModelForCausalLM.from_pretrained(model_name)
7
        self.ref_model = AutoModelForCausalLM.from_pretrained(model_name)
8
        self.ref_model.eval()  # 冻结参考模型
9

10
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
11
        self.beta = beta
12
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-6)
13

14
    def compute_log_probs(self, model, input_ids, attention_mask, label_ids=None):
15
        """计算序列的平均 log probability"""
16
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
17

18
        # 实际计算时需要根据 label_ids 偏移
19
        logits = outputs.logits
20

21
        # 计算每个 token 的 log prob
22
        log_probs = torch.log_softmax(logits, dim=-1)
23

24
        # 获取 target tokens 的 log prob
25
        # 注意：需要处理 label shift
26
        gathered_log_probs = log_probs[..., :-1, :].gather(-1, label_ids[..., 1:].unsqueeze(-1))
27

28
        # 返回平均 log prob（排除 padding）
29
        return gathered_log_probs.squeeze(-1).mean()
30

31
    def training_step(self, batch):
32
        # 1. 计算 policy 在 chosen/rejected 上的 log prob
33
        policy_chosen_logps = self.compute_log_probs(
34
            self.model,
35
            batch["chosen_input_ids"],
36
            batch["chosen_attention_mask"],
37
            batch["chosen_labels"]
38
        )
39

40
        policy_rejected_logps = self.compute_log_probs(
41
            self.model,
42
            batch["rejected_input_ids"],
43
            batch["rejected_attention_mask"],
44
            batch["rejected_labels"]
45
        )
46

47
        # 2. 计算 reference model 的 log prob
48
        with torch.no_grad():
49
            ref_chosen_logps = self.compute_log_probs(
50
                self.ref_model,
51
                batch["chosen_input_ids"],
52
                batch["chosen_attention_mask"],
53
                batch["chosen_labels"]
54
            )
55

56
            ref_rejected_logps = self.compute_log_probs(
57
                self.ref_model,
58
                batch["rejected_input_ids"],
59
                batch["rejected_attention_mask"],
60
                batch["rejected_labels"]
61
            )
62

63
        # 3. 计算 DPO 损失
64
        loss = dpo_loss(
65
            policy_chosen_logps, policy_rejected_logps,
66
            ref_chosen_logps, ref_rejected_logps,
67
            beta=self.beta
68
        )
69

70
        # 4. 反向传播
71
        self.optimizer.zero_grad()
72
        loss.backward()
73
        self.optimizer.step()
74

75
        return loss.item()

1
# 偏好数据质量检查
2
data_quality_rules = {
3
    "偏好一致性": "同一 prompt 的多次标注应一致",
4
    "辨别力": "模型能轻易区分 chosen vs rejected 则数据价值低",
5
    "多样性": "rejected 应该是合理的但次优的答案，而非明显错误",
6
    "长度平衡": "避免长度成为偏好因素",
7
}
8

9
def filter_preference_data(data):
10
    """过滤低质量偏好数据"""
11
    filtered = []
12
    for item in data:
13
        # 长度惩罚：避免模型学会"越长越好"
14
        len_ratio = len(item["chosen"]) / max(len(item["rejected"]), 1)
15
        if 0.5 < len_ratio < 2.0:  # 长度差异过大则过滤
16
            filtered.append(item)
17
    return filtered

1
# 1. Reward Model 预训练
2
# 先在大量偏好数据上训练 reward model
3
rm_pretrain_data = load_large_preference_dataset()
4
reward_model.pretrain(rm_pretrain_data)
5

6
# 2. 学习率调度
7
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
8
    optimizer,
9
    T_max=len(dataloader) * num_epochs,
10
    eta_min=1e-7
11
)
12

13
# 3. 梯度裁剪
14
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
15

16
# 4. 早停
17
best_reward_margin = 0
18
patience = 3
19
for epoch in range(num_epochs):
20
    val_margin = evaluate_reward_margin(model)
21
    if val_margin > best_reward_margin + 0.01:
22
        best_reward_margin = val_margin
23
        patience_counter = 0
24
    else:
25
        patience_counter += 1
26
        if patience_counter >= patience:
27
            break

1
# 对齐税：RLHF 后模型在某些任务上能力下降
2
# 解决方案：混合预训练
3

4
class AlignmentTaxMitigator:
5
    def mix_pretraining(self, model, ratio=0.1):
6
        """
7
        在 RLHF 训练数据中混入一定比例的预训练数据
8
        ratio: 预训练数据占比（通常 5-15%）
9
        """
10
        mixed_dataloader = []
11

12
        for batch in rlhf_dataloader:
13
            # 90% RLHF 数据
14
            mixed_dataloader.append(batch)
15

16
            # 10% 预训练数据
17
            if random.random() < ratio:
18
                pretrain_batch = sample_pretrain_batch()
19
                mixed_dataloader.append(pretrain_batch)
20

21
        return mixed_dataloader