1
2
3

def grpo_kl(pi_logprob, pi_ref_logprob):
    """GRPO 的无偏低方差 KL 散度计算"""
    return pi_ref_logprob.exp() / pi_logprob.exp() - (pi_ref_logprob - pi_logprob) - 1

1
2
3
4
5
6

阶段一：SFT（监督微调）
    ↓ 得到基础对话模型
阶段二：Reward Model 训练
    ↓ 得到能打分的奖励模型
阶段三：PPO 策略优化
    ↓ 得到对齐后的最终模型

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

# 手写 LoRA 模块
class Lora(torch.nn.Module):
    def __init__(self, linear):
        super().__init__()
        self.linear = linear
        self.lora_A = torch.nn.Parameter(torch.randn(linear.in_features, 128) * 0.1)
        self.lora_B = torch.nn.Parameter(torch.zeros(128, linear.out_features))
        self.linear.weight.requires_grad = False

    def forward(self, x):
        y_linear = self.linear(x)
        y_lora = x.matmul(self.lora_A).matmul(self.lora_B)
        return y_linear + y_lora / 128

1
2
3
4
5

{
  "question": "什么是机器学习？",
  "chosen": "机器学习是一种人工智能方法，通过算法和统计模型使计算机系统能够执行特定任务，而无需明确编程指令。",
  "rejected": "机器学习是计算机科学的一个领域。"
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

class CriticModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.rwtransformer = AutoModel.from_pretrained('facebook/opt-350m', dropout=0.0)
        self.v_head = torch.nn.Linear(512, 1, bias=False)

    def forward(self, input_ids, attention_mask):
        value = self.rwtransformer(input_ids=input_ids,
                                    attention_mask=attention_mask).last_hidden_state
        value = self.v_head(value).squeeze(-1)

        # 对每对 (chosen, rejected) 计算 pairwise loss
        for ... in zip(input_ids[:4], input_ids[4:], value[:4], value[4:]):
            start = (input_ids_chosen == input_ids_rejected).tolist().index(False)
            value_chosen = value_chosen[start:end]
            value_rejected = value_rejected[start:end]
            loss = -torch.nn.functional.logsigmoid(value_chosen - value_rejected).mean()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

from trl import RewardTrainer, RewardConfig

# 数据预处理：将 chosen/rejected 转换为 input_ids_chosen / input_ids_rejected
def process_func(example):
    chosen = example["question"] + example["chosen"]
    rejected = example["question"] + example["rejected"]
    tokenized_chosen = tokenizer(chosen)
    tokenized_rejected = tokenizer(rejected)
    return {
        "input_ids_chosen": tokenized_chosen["input_ids"],
        "attention_mask_chosen": tokenized_chosen["attention_mask"],
        "input_ids_rejected": tokenized_rejected["input_ids"],
        "attention_mask_rejected": tokenized_rejected["attention_mask"],
    }

# 配置与训练
config = RewardConfig(output_dir="./reward_model")
trainer = RewardTrainer(model=model, tokenizer=tokenizer, args=config, train_dataset=dataset)
trainer.train()

1
2
3
4
5
6
7

def get_reward_kl(end, prob_old, prob_ref, reward):
    # 两份预测概率求 KL 散度
    reward_kl = -0.1 * (prob_old - prob_ref)
    # 把 Reward Model 的分数加在最后一个 token 上
    for i, e in enumerate(end):
        reward_kl[i, e] += reward[i].clamp(-5, 5)
    return reward_kl

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

def get_delta(value_old, reward_kl):
    delta = []
    for i in reversed(range(255, value_old.shape[1])):
        value_next = 0.0
        if i != value_old.shape[1] - 1:
            value_next = value_old[:, i + 1]
        d = reward_kl[:, i] + value_next - value_old[:, i]
        if len(delta):
            d += 0.95 * delta[-1]
        delta.append(d)
    delta = torch.stack(delta[::-1], dim=1)
    return delta

1
2
3
4
5
6

def get_loss_actor(prob_new, prob_old, delta, generate_mask):
    ratio = ((prob_new - prob_old) * generate_mask).exp()
    loss1 = delta * ratio
    loss2 = delta * ratio.clamp(0.8, 1.2)  # clip 范围 [0.8, 1.2]
    loss = torch.min(loss1, loss2) * generate_mask
    return -(loss.sum() / generate_mask.sum() / 8)

1
2
3
4
5
6

def get_loss_critic(value_new, value_old, delta, generate_mask):
    loss1 = (value_new - delta - value_old)**2
    value_new_clipped = value_new.clamp(value_old - 0.2, value_old + 0.2)
    loss2 = (value_new_clipped - delta - value_old)**2
    loss = torch.max(loss1, loss2) * generate_mask
    return loss.sum() / 2 / generate_mask.sum() / 8

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

# PPO 配置
config = PPOConfig(
    model_name="Qwen2.5-0.5B-Instruct",
    learning_rate=6e-6,
    batch_size=128,
    mini_batch_size=16,
    target_kl=0.03,
    kl_penalty="kl",
    ppo_epochs=1,
)

# 奖励计算：使用情感分析 pipeline
sentiment_pipe = pipeline("sentiment-analysis", model=reward_model,
                          tokenizer=reward_tokenizer, device=device)

# 训练循环
for step, batch in enumerate(ppo_trainer.dataloader):
    query_tensors = batch["input_ids"]
    response_tensors = ppo_trainer.generate(query_tensors, **generation_kwargs)

    # 拼接 prompt + response 送入奖励模型
    texts_to_score = [q + r for q, r in zip(batch["query"], batch["response"])]
    pipe_outputs = sentiment_pipe(texts_to_score)

    # 提取 POSITIVE 分数作为奖励
    rewards = [torch.tensor(score['score']) for score in pipe_outputs
               if score['label'] == 'POSITIVE']

    # PPO 优化步骤
    stats = ppo_trainer.step(query_tensors, response_tensors, rewards)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer

model = AutoModelForCausalLMWithValueHead.from_pretrained(
    model_path,
    reward_adapter="./reward_model",  # 加载 Reward Model 适配器
    peft_config=peft_config,
    quantization_config=bnb_config
)

ppo_config = PPOConfig(
    kl_penalty="full",
    ppo_epochs=3,
    batch_size=2,
    mini_batch_size=1
)

ppo_trainer = PPOTrainer(
    config=ppo_config, model=model,
    ref_model=None, tokenizer=tokenizer,
    dataset=queries_dataset, data_collator=collator
)

for batch in ppo_trainer.dataloader:
    response_tensors = ppo_trainer.generate(batch, **generation_kwargs)
    scores = [model.compute_reward_score(input_ids=ids)[0, -1, 0]
              for ids in combined_ids]
    stats = ppo_trainer.step(batch, response_tensors, scores)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

def get_prob_diff(actor, input_ids, attention_mask, answer_mask):
    prob = actor(input_ids=input_ids, attention_mask=attention_mask).logits
    input_ids = input_ids[:, 1:]
    answer_mask = answer_mask[:, 1:]
    prob = prob[:, :-1]

    # 取所有字的预测概率（对数域）
    prob = (prob.softmax(2) + 1e-8).log()
    prob = prob.gather(2, index=input_ids.unsqueeze(2)).squeeze(2)

    # 取答案部分的联合概率（对数和）
    prob = (prob * answer_mask).sum(1)

    # chosen 与 rejected 的概率差
    return prob[:b] - prob[b:]

# DPO 训练循环
for i in range(8000):
    data = get_data()
    prob_diff = get_prob_diff(model_actor, **data)
    with torch.no_grad():
        prob_diff_ref = get_prob_diff(model_actor_ref, **data)

    # DPO 损失：beta * (prob_diff - prob_diff_ref)
    loss = 0.1 * (prob_diff - prob_diff_ref)
    loss = -(loss.sigmoid() + 1e-8).log().mean()

    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

from trl import DPOConfig, DPOTrainer

args = DPOConfig(
    output_dir='output_dir',
    loss_type='sigmoid',       # DPO 损失类型
    beta=0.1,                  # KL 系数
    per_device_train_batch_size=8,
    max_steps=80000,
    learning_rate=1e-5,
    optim='rmsprop',
    max_length=100,
    max_prompt_length=100,
)

# 数据集格式需包含 prompt, chosen, rejected 三列
trainer = DPOTrainer(
    model=model_actor,
    ref_model=model_actor_ref,
    args=args,
    train_dataset=dataset,
    tokenizer=tokenizer,
)
trainer.train()